Files
agent_management/api/swarm/router.py
T

609 lines
24 KiB
Python

"""Sub-mode runtime compatibility router for legacy /api/swarms clients."""
import json
import uuid
from datetime import datetime, timedelta
from typing import Dict, Any
from fastapi import APIRouter, Depends, HTTPException, BackgroundTasks, Request
from fastapi.responses import FileResponse, Response
from sqlalchemy.orm import Session
from database import (
get_db, Swarm, SwarmAgent, SwarmMessage,
SwarmStatus, SwarmAgentStatus
)
from k8s_manager import sanitize_k8s_name
from .artifact_store import (
load_azblob_artifact,
load_runtime_artifact,
runtime_artifact_download_path,
runtime_artifact_uri,
store_text_artifact,
)
from .models import (
SwarmCreateRequest, SwarmCreateResponse, SwarmStatusResponse,
SwarmStopRequest, SwarmStopResponse, SwarmAgentInfo, SwarmMetrics,
ApprovalDecisionRequest
)
from .orchestrator import SwarmOrchestrator
from api.status_projection import RuntimeDisplayStatus, project_runtime_run_status
swarms_router = APIRouter(prefix="/api/swarms", tags=["sub-mode-runtime-compatibility"])
def generate_swarm_id() -> str:
"""Generate unique swarm ID"""
return f"swm_{uuid.uuid4().hex[:12]}"
def generate_agent_id(role: str) -> str:
"""Generate unique agent ID"""
return f"agi_{role}_{uuid.uuid4().hex[:8]}"
def _agent_infos_for_swarm(db: Session, swarm_id: str) -> list[SwarmAgentInfo]:
"""Build response agent summaries for a swarm."""
agents = db.query(SwarmAgent).filter(SwarmAgent.swarm_id == swarm_id).all()
return [
SwarmAgentInfo(
agent_id=agent.agent_id,
role=agent.role,
status=agent.status.value,
namespace=agent.namespace,
service_url=agent.service_url,
current_task=agent.current_task,
output=agent.output,
)
for agent in agents
]
def _synthesized_artifacts_for_swarm(db: Session, swarm: Swarm) -> list[Dict[str, Any]]:
"""Return stored artifacts or a compatibility summary for old empty terminal runs."""
if swarm.artifacts:
return swarm.artifacts
if swarm.status not in {SwarmStatus.COMPLETED, SwarmStatus.FAILED, SwarmStatus.STOPPED}:
return []
agents = db.query(SwarmAgent).filter(SwarmAgent.swarm_id == swarm.swarm_id).all()
agent_summaries = []
for agent in agents:
detail = f"{agent.role}:{agent.status.value}"
if agent.output:
detail += f" - {agent.output[:300]}"
agent_summaries.append(detail)
failed = swarm.status == SwarmStatus.FAILED
summary = (
swarm.error_message
if failed and swarm.error_message
else "Runtime reached a terminal state without storing a concrete artifact. "
"This compatibility artifact preserves a visible delivery record for Heicode sub-mode clients."
)
if agent_summaries:
summary += "\nAgents: " + "; ".join(agent_summaries)
if swarm.task_description:
summary = f"Task: {swarm.task_description}\n{summary}"
artifact_id = f"art_{swarm.swarm_id}_{'failure' if failed else 'summary'}"
stored = None
try:
stored = store_text_artifact(swarm.swarm_id, artifact_id, summary[:2000])
except Exception:
stored = None
return [
{
"artifact_id": artifact_id,
"artifact_type": "other" if failed else "document",
"title": "Runtime execution failed" if failed else "Runtime execution summary",
"summary": summary[:2000],
"uri": stored.uri if stored else runtime_artifact_uri(swarm.swarm_id, artifact_id),
"agent_instance_id": None,
"mime_type": stored.mime_type if stored else "text/plain",
"size_bytes": stored.size_bytes if stored else len(summary[:2000].encode("utf-8")),
"stage": swarm.phase or "development",
"checkpoint": "failed" if failed else "artifact_ready",
"metadata": {
"redacted": True,
"source": "agent-manager-sub-mode-runtime",
"runtime_deployment_id": swarm.swarm_id,
"synthesized": True,
"summary_only": True,
"artifact_layout": "single_file_content",
"primary_read_path": "content",
"agent_count": len(agents),
"content_hash": stored.content_hash if stored else None,
"download_path": runtime_artifact_download_path(swarm.swarm_id, artifact_id),
},
}
]
def _build_swarm_status_response(db: Session, swarm: Swarm) -> SwarmStatusResponse:
"""Return a standard status payload for the sub-mode runtime compatibility API."""
elapsed_seconds = int((datetime.utcnow() - swarm.created_at).total_seconds())
display_status = project_runtime_run_status(swarm.status)
return SwarmStatusResponse(
deployment_id=swarm.swarm_id,
swarm_id=swarm.swarm_id,
status=display_status,
phase=swarm.phase,
progress=swarm.progress,
agents=_agent_infos_for_swarm(db, swarm.swarm_id),
metrics=SwarmMetrics(
total_messages=swarm.total_messages,
tokens_used=swarm.tokens_used,
elapsed_seconds=elapsed_seconds,
),
artifacts=_synthesized_artifacts_for_swarm(db, swarm),
error_message=swarm.error_message,
created_at=swarm.created_at,
updated_at=swarm.updated_at,
)
def _stop_swarm_record(db: Session, swarm_id: str, request: SwarmStopRequest) -> SwarmStopResponse:
"""Idempotently stop a compatibility runtime-run database record."""
swarm = db.query(Swarm).filter(Swarm.swarm_id == swarm_id).first()
if not swarm:
raise HTTPException(status_code=404, detail="Runtime deployment not found")
stopped_at = datetime.utcnow()
if swarm.status != SwarmStatus.STOPPED:
swarm.status = SwarmStatus.STOPPED
swarm.error_message = request.reason
swarm.updated_at = stopped_at
db.query(SwarmAgent).filter(SwarmAgent.swarm_id == swarm_id).update(
{"status": SwarmAgentStatus.FAILED if request.reason else SwarmAgentStatus.COMPLETED}
)
db.commit()
return SwarmStopResponse(
deployment_id=swarm_id,
swarm_id=swarm_id,
status=RuntimeDisplayStatus.STOPPED.value,
stopped_at=stopped_at,
)
async def initialize_and_execute_swarm(swarm_id: str, db_url: str):
"""Background task to initialize and execute a sub-mode runtime run."""
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker
# Create new database session for background task
engine = create_engine(db_url)
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
db = SessionLocal()
try:
orchestrator = SwarmOrchestrator(swarm_id, db)
# Initialize swarm
success = await orchestrator.initialize()
if not success:
return
# Execute swarm task
await orchestrator.execute()
except Exception as e:
print(f"Error in background sub-mode runtime execution: {e}")
finally:
await orchestrator.cleanup()
db.close()
async def create_swarm(
request: SwarmCreateRequest,
background_tasks: BackgroundTasks,
db: Session = Depends(get_db)
):
"""Create a new sub-mode runtime run."""
swarm_id = generate_swarm_id()
project_context = request.project_context.dict() if request.project_context else {}
if request.callback:
project_context["_callback"] = request.callback.dict(exclude_none=True)
# Create swarm record
swarm = Swarm(
swarm_id=swarm_id,
task_description=request.task_description,
project_context=project_context,
orchestration_strategy=request.orchestration.strategy,
max_iterations=request.orchestration.max_iterations,
timeout_minutes=request.orchestration.timeout_minutes,
callback_url=request.callback.url if request.callback else None,
callback_method=request.callback.method if request.callback else "POST",
owner_id=request.owner_id,
status=SwarmStatus.INITIALIZING
)
db.add(swarm)
db.commit()
db.refresh(swarm)
# Create agent records (pods will be created by background task)
agent_infos = []
for agent_config in request.agents:
for i in range(agent_config.replicas):
agent_id = generate_agent_id(agent_config.role)
replica_suffix = f"-{i+1}" if agent_config.replicas > 1 else ""
namespace = sanitize_k8s_name(f"swarm-{swarm_id[:8]}-{agent_config.role}{replica_suffix}")
pod_name = sanitize_k8s_name(f"agent-{agent_id}")
swarm_agent = SwarmAgent(
agent_id=agent_id,
swarm_id=swarm_id,
role=agent_config.role,
template=agent_config.template,
model=agent_config.model,
capabilities=agent_config.capabilities,
system_prompt=agent_config.system_prompt,
namespace=namespace,
pod_name=pod_name,
service_url=f"http://{pod_name}.{namespace}.svc.cluster.local:8000",
status=SwarmAgentStatus.PENDING
)
db.add(swarm_agent)
db.commit()
db.refresh(swarm_agent)
agent_infos.append(SwarmAgentInfo(
agent_id=swarm_agent.agent_id,
role=swarm_agent.role,
status=swarm_agent.status.value,
namespace=swarm_agent.namespace,
service_url=swarm_agent.service_url
))
# Schedule background task to initialize and execute swarm
# Note: In production, this should use a proper task queue like Celery
from database import DATABASE_URL
background_tasks.add_task(initialize_and_execute_swarm, swarm_id, DATABASE_URL)
return SwarmCreateResponse(
deployment_id=swarm_id,
swarm_id=swarm_id,
status=project_runtime_run_status(swarm.status),
agents=agent_infos,
created_at=swarm.created_at,
estimated_ready_at=swarm.created_at + timedelta(minutes=2)
)
@swarms_router.post("", response_model=SwarmCreateResponse)
async def create_swarm_compat(
payload: Dict[str, Any],
request: Request,
background_tasks: BackgroundTasks,
db: Session = Depends(get_db)
):
"""Compatibility entrypoint for Heicode sub-mode Runtime adapters."""
if payload.get("dry_run") is True:
raise HTTPException(status_code=422, detail="dry_run is not supported by Runtime create; no swarm was created")
plan = payload.get("orchestration_plan")
if not isinstance(plan, dict):
raise HTTPException(status_code=422, detail="orchestration_plan must be an object")
if not plan.get("sub_mode"):
raise HTTPException(status_code=422, detail="orchestration_plan.sub_mode is required")
if plan.get("sub_mode") not in {"agile", "waterfall"}:
raise HTTPException(status_code=422, detail="orchestration_plan.sub_mode must be agile or waterfall")
if not isinstance(plan.get("user_context"), dict) or not plan["user_context"].get("user_id"):
raise HTTPException(status_code=422, detail="orchestration_plan.user_context.user_id is required")
callback = payload.get("callback") or plan.get("callback")
if not isinstance(callback, dict) or not callback.get("url"):
raise HTTPException(status_code=422, detail="callback.url is required")
idempotency_key = request.headers.get("X-Idempotency-Key") or request.headers.get("Idempotency-Key")
if idempotency_key:
existing = None
for candidate in db.query(Swarm).order_by(Swarm.created_at.asc()).all():
context = candidate.project_context or {}
if context.get("idempotency_key") == idempotency_key:
existing = candidate
break
if existing:
return _build_swarm_status_response(db, existing)
agents = plan.get("agents") or payload.get("agents") or [
{"role": "backend", "capabilities": ["code", "test"]},
{"role": "frontend", "capabilities": ["ui", "test"]},
{"role": "reviewer", "capabilities": ["review"]},
]
budget = plan.get("budget") or {}
agile_context = plan.get("agile_context") or payload.get("agile_context") or {}
user_context = plan.get("user_context") or {}
metadata = plan.get("metadata") or payload.get("metadata") or {}
project_context = plan.get("project_context") or {}
if not isinstance(project_context, dict):
project_context = {}
project_context = {
**project_context,
"intent_id": plan.get("intent_id"),
"template_hint": plan.get("template_hint"),
"binding_scope": user_context.get("binding_scope"),
"sub_mode": plan.get("sub_mode", "agile"),
"agile_context": agile_context,
"budget": budget,
"billing_context": plan.get("billing_context") or payload.get("billing_context") or {},
"resource_grants": plan.get("resource_grants") or payload.get("resource_grants") or [],
"idempotency_key": idempotency_key,
"correlation_id": metadata.get("correlation_id") or payload.get("correlation_id"),
"manager_deployment_id": payload.get("deployment_id")
or metadata.get("manager_deployment_id")
or metadata.get("heicode_deployment_id"),
"heicode_deployment_id": payload.get("deployment_id")
or metadata.get("heicode_deployment_id")
or metadata.get("manager_deployment_id"),
}
swarm_request = SwarmCreateRequest(
task_description=plan.get("objective") or plan.get("intent_id") or "Heicode sub-mode task",
project_context=project_context,
agents=agents,
orchestration={
"strategy": "sequential" if plan.get("sub_mode", "agile") == "waterfall" else "hybrid",
"max_iterations": agile_context.get("max_iterations", 3),
"timeout_minutes": max(1, int((budget.get("max_duration_sec") or 1800) / 60)),
},
callback=callback,
owner_id=str(user_context.get("user_id") or payload.get("owner_id") or "default"),
)
return await create_swarm(swarm_request, background_tasks, db)
@swarms_router.get("/{swarm_id}", response_model=SwarmStatusResponse)
async def get_swarm_detail_compat(swarm_id: str, db: Session = Depends(get_db)):
"""Compatibility detail endpoint for legacy Manager runtime clients."""
swarm = db.query(Swarm).filter(Swarm.swarm_id == swarm_id).first()
if not swarm:
raise HTTPException(status_code=404, detail="Runtime deployment not found")
return _build_swarm_status_response(db, swarm)
@swarms_router.get("/{swarm_id}/status", response_model=SwarmStatusResponse)
async def get_swarm_status_compat(swarm_id: str, db: Session = Depends(get_db)):
"""Compatibility status endpoint for legacy Manager runtime clients."""
swarm = db.query(Swarm).filter(Swarm.swarm_id == swarm_id).first()
if not swarm:
raise HTTPException(status_code=404, detail="Runtime deployment not found")
return _build_swarm_status_response(db, swarm)
@swarms_router.post("/{swarm_id}/stop", response_model=SwarmStopResponse)
async def stop_swarm_compat(
swarm_id: str,
request: SwarmStopRequest,
db: Session = Depends(get_db)
):
"""Compatibility stop endpoint used by Heicode Manager."""
return _stop_swarm_record(db, swarm_id, request)
async def get_swarm_logs(swarm_id: str, db: Session = Depends(get_db)):
"""Get aggregated logs from all agents in a sub-mode runtime run."""
swarm = db.query(Swarm).filter(Swarm.swarm_id == swarm_id).first()
if not swarm:
raise HTTPException(status_code=404, detail="Runtime deployment not found")
agents = db.query(SwarmAgent).filter(SwarmAgent.swarm_id == swarm_id).all()
logs = {
"swarm_id": swarm_id,
"agents": []
}
messages_by_agent = {}
swarm_messages = (
db.query(SwarmMessage)
.filter(SwarmMessage.swarm_id == swarm_id)
.order_by(SwarmMessage.created_at.asc())
.all()
)
for message in swarm_messages:
related_agent_id = message.from_agent_id or message.to_agent_id
if not related_agent_id:
continue
messages_by_agent.setdefault(related_agent_id, []).append(message)
for agent in agents:
agent_messages = messages_by_agent.get(agent.agent_id, [])
log_lines = [
f"status={agent.status.value}",
f"role={agent.role}",
]
if agent.current_task:
log_lines.append(f"current_task={agent.current_task}")
if agent.output:
log_lines.append(f"last_output={agent.output[:500]}")
if agent_messages:
for message in agent_messages[-5:]:
metadata = message.message_metadata or {}
request_id = metadata.get("newapi_request_id")
usage = metadata.get("model_usage") or {}
suffix_parts = []
if request_id:
suffix_parts.append(f"newapi_request_id={request_id}")
if usage.get("total_tokens"):
suffix_parts.append(f"tokens={usage['total_tokens']}")
suffix = f" ({', '.join(suffix_parts)})" if suffix_parts else ""
log_lines.append(
f"{message.created_at.isoformat()} {message.message_type}: {message.content[:300]}{suffix}"
)
else:
log_lines.append("no_runtime_messages_recorded")
logs["agents"].append({
"agent_id": agent.agent_id,
"role": agent.role,
"namespace": agent.namespace,
"pod_name": agent.pod_name,
"logs": "\n".join(log_lines),
"messages": [
{
"message_id": message.message_id,
"message_type": message.message_type,
"created_at": message.created_at,
"newapi_request_id": (message.message_metadata or {}).get("newapi_request_id"),
"model_usage": (message.message_metadata or {}).get("model_usage"),
"metadata": message.message_metadata or {},
}
for message in agent_messages[-20:]
],
})
return logs
@swarms_router.get("/{swarm_id}/logs")
async def get_swarm_logs_compat(swarm_id: str, db: Session = Depends(get_db)):
"""Compatibility logs endpoint under /api/swarms."""
return await get_swarm_logs(swarm_id, db)
@swarms_router.get("/{swarm_id}/events")
async def get_swarm_events_compat(swarm_id: str, db: Session = Depends(get_db)):
"""Return runtime messages as compatibility events for polling fallback."""
swarm = db.query(Swarm).filter(Swarm.swarm_id == swarm_id).first()
if not swarm:
raise HTTPException(status_code=404, detail="Runtime deployment not found")
messages = (
db.query(SwarmMessage)
.filter(SwarmMessage.swarm_id == swarm_id)
.order_by(SwarmMessage.created_at.asc())
.all()
)
events = [
{
"event_id": message.message_id,
"event_type": f"runtime.message.{message.message_type}",
"swarm_id": swarm_id,
"agent_instance_id": message.from_agent_id or message.to_agent_id,
"occurred_at": message.created_at,
"payload": {
"from_agent_id": message.from_agent_id,
"to_agent_id": message.to_agent_id,
"message_type": message.message_type,
"summary": message.content[:300] if message.content else None,
"metadata": message.message_metadata or {},
},
}
for message in messages
]
if not any(event["event_type"] == "artifact.created" for event in events):
for artifact in _synthesized_artifacts_for_swarm(db, swarm):
events.append(
{
"event_id": artifact.get("artifact_id"),
"event_type": "artifact.created",
"swarm_id": swarm_id,
"agent_instance_id": artifact.get("agent_instance_id"),
"occurred_at": swarm.completed_at or swarm.updated_at,
"payload": artifact,
}
)
return {"success": True, "swarm_id": swarm_id, "events": events}
@swarms_router.get("/{swarm_id}/metrics")
async def get_swarm_metrics_compat(swarm_id: str, db: Session = Depends(get_db)):
"""Return basic runtime metrics for compatibility polling fallback."""
swarm = db.query(Swarm).filter(Swarm.swarm_id == swarm_id).first()
if not swarm:
raise HTTPException(status_code=404, detail="Runtime deployment not found")
elapsed_seconds = int((datetime.utcnow() - swarm.created_at).total_seconds())
return {
"success": True,
"swarm_id": swarm_id,
"status": project_runtime_run_status(swarm.status),
"stage": swarm.phase,
"checkpoint": "completed" if swarm.status == SwarmStatus.COMPLETED else "agent_running",
"metrics": {
"tokens_used": swarm.tokens_used,
"duration_ms": elapsed_seconds * 1000,
"artifact_count": len(_synthesized_artifacts_for_swarm(db, swarm)),
"total_messages": swarm.total_messages,
},
}
@swarms_router.get("/{swarm_id}/artifacts/{artifact_id}/content")
async def get_swarm_artifact_content(
swarm_id: str,
artifact_id: str,
db: Session = Depends(get_db),
):
"""Return full persisted content for a sub-mode runtime artifact URI."""
swarm = db.query(Swarm).filter(Swarm.swarm_id == swarm_id).first()
if not swarm:
raise HTTPException(status_code=404, detail="Runtime deployment not found")
artifacts_by_id = {
artifact.get("artifact_id"): artifact
for artifact in _synthesized_artifacts_for_swarm(db, swarm)
}
artifact = artifacts_by_id.get(artifact_id)
if not artifact:
raise HTTPException(status_code=404, detail="Artifact not found")
stored = load_runtime_artifact(swarm_id, artifact_id)
if stored:
return FileResponse(
path=stored.path,
media_type=stored.mime_type,
filename=stored.path.name,
)
blob_content = load_azblob_artifact(artifact.get("uri"))
if blob_content:
content, mime_type, filename = blob_content
return Response(
content=content,
media_type=mime_type,
headers={"Content-Disposition": f'attachment; filename="{filename}"'},
)
raise HTTPException(status_code=404, detail="Artifact content not found")
@swarms_router.post("/{swarm_id}/approvals/{approval_id}")
async def receive_swarm_approval_decision(
swarm_id: str,
approval_id: str,
request: ApprovalDecisionRequest,
db: Session = Depends(get_db)
):
"""Accept Manager approval decisions for paused high-risk runtime actions."""
swarm = db.query(Swarm).filter(Swarm.swarm_id == swarm_id).first()
if not swarm:
raise HTTPException(status_code=404, detail="Runtime deployment not found")
if request.approval_id != approval_id:
raise HTTPException(status_code=422, detail="approval_id path/body mismatch")
if request.decision not in {"approved", "rejected"}:
raise HTTPException(status_code=422, detail="decision must be approved or rejected")
message = SwarmMessage(
message_id=f"appr_{uuid.uuid4().hex[:12]}",
swarm_id=swarm_id,
from_agent_id=None,
to_agent_id=None,
message_type="approval_decision",
content=request.decision,
message_metadata=request.model_dump(exclude_none=True),
)
db.add(message)
db.commit()
return {
"success": True,
"swarm_id": swarm_id,
"approval_id": approval_id,
"decision": request.decision,
"status": "accepted",
}