Songhaoz666andClaude Opus 4.8
a289495823
回应 Fasthei 的 Request changes 两个阻塞项：
1) 安全 / fail-closed 沙箱隔离（原仅靠 ENABLE_QUALITY_EVAL + 运维约定）：
- 新增第二道显式确认 HEICODE_SANDBOX_ISOLATED（断言运行在隔离 Pod 内）。
- sandbox.run_tests() 与 quality.evaluate_run_quality() 执行任何代码前调用
assert_isolated()，未确认即抛 SandboxIsolationError——不写文件、不起子进程。
- 启动期 assert_quality_eval_safe()：ENABLE_QUALITY_EVAL 开但隔离未确认 → 拒绝启动
（平台级硬失败，非运维口头约定）。
- 文档（security-boundary §8.1/§9、CLAUDE.md）与测试同步：test-sandbox/test-quality
先断言未确认时硬失败，再显式确认后继续。
2) #10 DecisionTrace 可回放（原仅存被选中任务的标量）：
- Decision 现记录完整重放上下文：整个候选集（每候选 tau/eta/weight/p_norm/dependents）、
alpha/beta/epsilon、seed、free_slots、total_weight、explore_draw、select_pick、
select_index、explored 分支。
- 新增 DecisionEngine.replay_decision(trace)：仅凭一条 trace（无 RNG/活体状态）复现被选任务；
test-decision-engine 断言「重放==实选」跨 50 次决策（探索+利用）成立。
- decision-engine.md §3.3 更新为可回放 DecisionTrace。
附：新增 docs/TESTING.md（reviewer 速查：依赖安装 + 每套测试命令，复审者此前因缺 fakeredis
未能跑到断言）。本地 11 项 gate 全绿。
影响范围：Swarm（orchestrator + 测试 + 文档）。不改 Manager↔Swarm 契约；新增开关
HEICODE_SANDBOX_ISOLATED（默认未设=拒绝执行）。仍非验收：gain/Benchmark_Agent 仍 NaN。
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2293 lines
94 KiB
Python
2293 lines
94 KiB
Python
"""FastAPI orchestrator with WebSocket support for agent coordination."""
|
|
import asyncio
|
|
import json
|
|
import logging
|
|
import os
|
|
import time
|
|
from collections import defaultdict
|
|
from contextlib import asynccontextmanager
|
|
from typing import Any, Dict, List, Optional
|
|
|
|
from fastapi import FastAPI, Header, Request, WebSocket, WebSocketDisconnect, HTTPException
|
|
from fastapi.responses import JSONResponse, PlainTextResponse
|
|
from pydantic import BaseModel, Field
|
|
from prometheus_client import Counter, Gauge, Histogram, generate_latest, CONTENT_TYPE_LATEST
|
|
from dotenv import load_dotenv
|
|
|
|
# Load a local .env (gitignored) before the planner is imported, so its OpenAI client and the
|
|
# review/synthesis steps pick up credentials without them being passed on the command line.
|
|
load_dotenv()
|
|
|
|
from .redis_client import redis_client
|
|
from .agent_registry import agent_registry, AgentStatus, AgentMetadata
|
|
from .handoff_manager import handoff_manager, HandoffRequest
|
|
from .task_queue import task_queue, TaskStatus
|
|
from .swarm_runtime import RuntimeValidationError, swarm_runtime
|
|
from .planner import planner
|
|
from .master_agent import master_agent
|
|
from .quality import evaluate_run_quality
|
|
from .decision_engine import decision_engine, aco_dispatch_enabled
|
|
|
|
# Configure logging
|
|
logging.basicConfig(
|
|
level=logging.INFO,
|
|
format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
|
|
)
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
# Prometheus metrics
|
|
AGENTS_CREATED = Counter('swarm_agents_created_total', 'Total agents created')
|
|
AGENTS_FAILED = Counter('swarm_agents_failed_total', 'Total agents failed')
|
|
AGENTS_ACTIVE = Gauge('swarm_agents_active', 'Currently active agents')
|
|
TASKS_CREATED = Counter('swarm_tasks_created_total', 'Total tasks created')
|
|
TASKS_COMPLETED = Counter('swarm_tasks_completed_total', 'Total tasks completed')
|
|
TASKS_FAILED = Counter('swarm_tasks_failed_total', 'Total tasks failed')
|
|
HANDOFFS_TOTAL = Counter('swarm_handoffs_total', 'Total handoffs performed')
|
|
HANDOFF_DURATION = Histogram('swarm_handoff_duration_seconds', 'Handoff latency in seconds')
|
|
TASK_DURATION = Histogram('swarm_task_duration_seconds', 'Task duration in seconds')
|
|
AGENT_STATUS = Gauge('swarm_agent_status', 'Agent status by ID', ['agent_id', 'status'])
|
|
WEBSOCKET_CONNECTIONS = Gauge('swarm_websocket_connections', 'Active WebSocket connections')
|
|
ERRORS_TOTAL = Counter('swarm_errors_total', 'Total errors', ['error_type'])
|
|
_AUTH_WARNING_EMITTED = False
|
|
|
|
|
|
# WebSocket connection manager
|
|
class ConnectionManager:
|
|
"""Manages WebSocket connections for agents."""
|
|
|
|
def __init__(self):
|
|
self.active_connections: Dict[str, WebSocket] = {}
|
|
|
|
async def connect(self, agent_id: str, websocket: WebSocket):
|
|
"""Accept and store WebSocket connection."""
|
|
await websocket.accept()
|
|
self.active_connections[agent_id] = websocket
|
|
logger.info(f"Agent {agent_id} connected via WebSocket")
|
|
|
|
def disconnect(self, agent_id: str):
|
|
"""Remove WebSocket connection."""
|
|
if agent_id in self.active_connections:
|
|
del self.active_connections[agent_id]
|
|
logger.info(f"Agent {agent_id} disconnected")
|
|
|
|
async def send_message(self, agent_id: str, message: dict) -> bool:
|
|
"""Send message to specific agent. Returns True if it was delivered."""
|
|
if agent_id in self.active_connections:
|
|
try:
|
|
await self.active_connections[agent_id].send_json(message)
|
|
return True
|
|
except Exception as e:
|
|
logger.error(f"Failed to send message to agent {agent_id}: {e}")
|
|
return False
|
|
return False
|
|
|
|
async def broadcast(self, message: dict):
|
|
"""Broadcast message to all connected agents."""
|
|
for agent_id, connection in self.active_connections.items():
|
|
try:
|
|
await connection.send_json(message)
|
|
except Exception as e:
|
|
logger.error(f"Failed to broadcast to agent {agent_id}: {e}")
|
|
|
|
|
|
manager = ConnectionManager()
|
|
|
|
# Last self-reported free capacity per agent, learned from register/heartbeat/accept/reject
|
|
# messages sent by the merged agent runtime. Used to avoid dispatching to an agent that
|
|
# reports it is momentarily full. Older agents omit these fields and are treated as having
|
|
# capacity (default 1), preserving the previous one-task-per-agent behavior.
|
|
AGENT_SLOTS: Dict[str, int] = {}
|
|
|
|
|
|
def record_agent_slots(agent_id: str, message: dict):
|
|
"""Update the cached free-slot count for an agent from a protocol message."""
|
|
slots = message.get("available_slots")
|
|
if isinstance(slots, int):
|
|
AGENT_SLOTS[agent_id] = max(0, slots)
|
|
|
|
|
|
def agent_has_capacity(agent_id: str) -> bool:
|
|
"""Return whether an agent currently reports free capacity (default: yes)."""
|
|
return AGENT_SLOTS.get(agent_id, 1) > 0
|
|
|
|
|
|
async def deposit_pheromone(task, agent_id: str, *, success: bool, result=None, run=None):
|
|
"""Group A learning (always on): update the agent's τ trail for this task's role.
|
|
|
|
Passive observation — never changes dispatch behavior by itself (selection is gated
|
|
separately by ENABLE_ACO_DISPATCH) and never fails the caller. cost_ratio is the task's
|
|
model cost as a fraction of the run budget when both are known, else 0 (no fabricated
|
|
penalty — see decision_engine module docstring).
|
|
"""
|
|
if not task or not agent_id:
|
|
return
|
|
try:
|
|
cost_ratio = 0.0
|
|
if run is not None and isinstance(result, dict):
|
|
budget = ((run.request_body or {}).get("orchestration_plan") or {}).get("budget") or {}
|
|
max_cost = budget.get("max_cost_usd")
|
|
if isinstance(max_cost, (int, float)) and max_cost > 0:
|
|
cost = float((result.get("usage") or {}).get("model_cost_usd") or 0.0)
|
|
cost_ratio = cost / float(max_cost)
|
|
await decision_engine.deposit(
|
|
agent_role=task.agent_role,
|
|
agent_id=agent_id,
|
|
success=success,
|
|
cost_ratio=cost_ratio,
|
|
)
|
|
except Exception as exc:
|
|
logger.debug(f"pheromone deposit skipped for {getattr(task, 'task_id', '?')}: {exc}")
|
|
|
|
|
|
# Background task for failure detection
|
|
async def failure_detection_loop():
|
|
"""Periodically check for failed agents and reassign their tasks."""
|
|
while True:
|
|
try:
|
|
await asyncio.sleep(10) # Check every 10 seconds
|
|
|
|
failed_agents = await agent_registry.check_failed_agents()
|
|
|
|
for agent_id in failed_agents:
|
|
# Reassign tasks from failed agent
|
|
reassigned_tasks = await task_queue.reassign_agent_tasks(agent_id)
|
|
|
|
if reassigned_tasks:
|
|
logger.info(
|
|
f"Reassigned {len(reassigned_tasks)} tasks from failed agent {agent_id}"
|
|
)
|
|
|
|
# Disconnect WebSocket
|
|
manager.disconnect(agent_id)
|
|
|
|
recovered_tasks = await task_queue.recover_orphaned_tasks(
|
|
set(manager.active_connections.keys())
|
|
)
|
|
if recovered_tasks:
|
|
logger.info(f"Recovered {len(recovered_tasks)} orphaned tasks")
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error in failure detection loop: {e}")
|
|
|
|
|
|
def review_loop_enabled() -> bool:
|
|
"""Whether the master critic runs before a run is declared complete (default off)."""
|
|
return os.getenv("ENABLE_REVIEW_LOOP", "false").lower() in {"1", "true", "yes"}
|
|
|
|
|
|
def review_max_cycles() -> int:
|
|
try:
|
|
return int(os.getenv("MAX_REVIEW_CYCLES", "2") or 2)
|
|
except ValueError:
|
|
return 2
|
|
|
|
|
|
async def build_dispatch_context(run, task) -> Dict[str, Any]:
|
|
"""Enrich a task's context at dispatch with completed dependency artifacts and live peers.
|
|
|
|
This is what lets specialists see each other's work: dependency_artifacts carries the
|
|
results of completed upstream tasks, and peer_agents lists other agents currently connected
|
|
on the run so the executor can consult them for overlapping fields. Purely additive — it
|
|
does not change the Manager-facing API.
|
|
"""
|
|
context = dict(task.context or {})
|
|
|
|
dependency_artifacts = []
|
|
for dep_id in task.depends_on or []:
|
|
dep = await task_queue.get_task(dep_id)
|
|
if dep and dep.status == TaskStatus.COMPLETED:
|
|
parsed = parse_task_result(dep) or {}
|
|
dependency_artifacts.append({
|
|
"task_id": dep_id,
|
|
"agent_role": dep.agent_role,
|
|
"summary": summarize_task_result(parsed),
|
|
"files_modified": parsed.get("files_modified") or [],
|
|
"changes": parsed.get("changes"),
|
|
})
|
|
if dependency_artifacts:
|
|
context["dependency_artifacts"] = dependency_artifacts
|
|
|
|
peer_agents = []
|
|
for other_id in run.task_ids:
|
|
if other_id == task.task_id:
|
|
continue
|
|
other = await task_queue.get_task(other_id)
|
|
if not other or not other.assigned_agent_id:
|
|
continue
|
|
if other.assigned_agent_id in manager.active_connections:
|
|
peer_agents.append({
|
|
"agent_id": other.assigned_agent_id,
|
|
"role": other.agent_role,
|
|
"capabilities": other.required_capabilities,
|
|
})
|
|
if peer_agents:
|
|
context["peer_agents"] = peer_agents
|
|
|
|
# Surface the task's role as the executor's specialist role, so the LLM prompt is
|
|
# role-aware and the testing/documentation peer-consult path can activate. The Task's
|
|
# agent_role is authoritative; fall back to any role already in context, then "general".
|
|
context.setdefault(
|
|
"specialist_role",
|
|
getattr(task, "agent_role", None) or context.get("agent_role") or "general",
|
|
)
|
|
context.setdefault("run_goal", run.objective)
|
|
context.setdefault("user_prompt", run.objective)
|
|
return context
|
|
|
|
|
|
async def maybe_run_review_cycle(run, tasks) -> bool:
|
|
"""Run the master critic on completed work; reopen rejected tasks if budget remains.
|
|
|
|
Returns True when tasks were re-opened (the run stays 'running' and will re-finalize after
|
|
the redo completes), False when the work is accepted or the cycle budget is exhausted.
|
|
"""
|
|
cycles = int(run.metadata.get("review_cycles", 0) or 0)
|
|
if cycles >= review_max_cycles():
|
|
return False
|
|
|
|
completed = [t for t in tasks if t.status == TaskStatus.COMPLETED]
|
|
if not completed:
|
|
return False
|
|
|
|
results = {t.task_id: {"result": parse_task_result(t) or {}} for t in completed}
|
|
verdict = await master_agent.review_and_decide(
|
|
run.objective,
|
|
[task_payload(t) for t in completed],
|
|
results,
|
|
)
|
|
if verdict.get("accepted", True):
|
|
run.metadata["review_summary"] = verdict.get("summary")
|
|
await swarm_runtime.save_run(run)
|
|
return False
|
|
|
|
retry_tasks = [tid for tid in (verdict.get("retry_tasks") or []) if tid in run.task_ids]
|
|
reopened = [tid for tid in retry_tasks if await task_queue.reopen_task(tid)]
|
|
if not reopened:
|
|
# Rejected but nothing actionable to redo: accept rather than loop forever.
|
|
run.metadata["review_summary"] = verdict.get("summary")
|
|
await swarm_runtime.save_run(run)
|
|
return False
|
|
|
|
run.metadata["review_cycles"] = cycles + 1
|
|
run.metadata["review_summary"] = verdict.get("summary")
|
|
run.status = "running"
|
|
await swarm_runtime.save_run(run)
|
|
await swarm_runtime.emit_event(
|
|
run,
|
|
"deployment.status_changed",
|
|
payload=swarm_runtime.status_payload(run, phase="Review", reason=verdict.get("summary")),
|
|
)
|
|
await swarm_runtime.emit_event(run, "timeline.updated", payload={
|
|
"summary": f"Review cycle {cycles + 1}: {verdict.get('summary')}",
|
|
"status": "running",
|
|
"retry_tasks": reopened,
|
|
})
|
|
logger.info(f"Review rejected run {run.swarm_id}; reopened {reopened} (cycle {cycles + 1})")
|
|
return True
|
|
|
|
|
|
async def task_dispatch_loop():
|
|
"""Assign pending tasks to idle, connected agents."""
|
|
while True:
|
|
try:
|
|
await asyncio.sleep(2)
|
|
|
|
pending_count = await task_queue.get_pending_count()
|
|
if pending_count == 0:
|
|
continue
|
|
|
|
idle_agents = await agent_registry.get_idle_agents()
|
|
connected_idle_agents = [
|
|
agent for agent in idle_agents
|
|
if agent.agent_id in manager.active_connections
|
|
and agent_has_capacity(agent.agent_id)
|
|
]
|
|
|
|
for agent in connected_idle_agents:
|
|
decision = None
|
|
if aco_dispatch_enabled():
|
|
# Group A (Option A, score-at-pull): enumerate ALL eligible ready tasks for
|
|
# this agent and SAMPLE one with P=τ^α·η^β/Σ instead of taking the first
|
|
# match. One-sided by design: the agent is fixed by arrival order.
|
|
candidates = await task_queue.get_ready_pending_tasks(agent.capabilities)
|
|
if not candidates:
|
|
break
|
|
dependents_counts: Dict[str, int] = {}
|
|
for t in await task_queue.get_all_tasks():
|
|
for dep in t.depends_on:
|
|
dependents_counts[dep] = dependents_counts.get(dep, 0) + 1
|
|
decision = await decision_engine.select(
|
|
agent.agent_id,
|
|
agent.capabilities,
|
|
candidates,
|
|
free_slots=AGENT_SLOTS.get(agent.agent_id, 1),
|
|
dependents_counts=dependents_counts,
|
|
)
|
|
task = next(t for t in candidates if t.task_id == decision.task_id)
|
|
await task_queue.remove_pending_task(task.task_id)
|
|
else:
|
|
task = await task_queue.get_ready_pending_task(agent.capabilities)
|
|
if not task:
|
|
break
|
|
|
|
success = await task_queue.assign_task(task.task_id, agent.agent_id)
|
|
if not success:
|
|
latest_task = await task_queue.get_task(task.task_id)
|
|
if latest_task and latest_task.status == TaskStatus.PENDING:
|
|
await task_queue.requeue_task(task.task_id)
|
|
continue
|
|
|
|
run = await swarm_runtime.get_run_for_task(task.task_id)
|
|
if run:
|
|
await swarm_runtime.emit_event(
|
|
run,
|
|
"task.claimed",
|
|
task_id=task.task_id,
|
|
agent_instance_id=agent.agent_id,
|
|
payload={
|
|
"task_id": task.task_id,
|
|
"agent_role": task.agent_role,
|
|
"agent_id": agent.agent_id,
|
|
},
|
|
)
|
|
if decision is not None:
|
|
# Group A telemetry: feeds tau/eta/p_decision in the collector.
|
|
await swarm_runtime.record_decision(run, decision.telemetry())
|
|
|
|
dispatch_context = (
|
|
await build_dispatch_context(run, task) if run else task.context
|
|
)
|
|
await manager.send_message(agent.agent_id, {
|
|
"type": "task_assignment",
|
|
"task_id": task.task_id,
|
|
"description": task.description,
|
|
"context": dispatch_context,
|
|
})
|
|
|
|
logger.info(
|
|
f"Dispatched task {task.task_id} to connected idle agent {agent.agent_id}"
|
|
)
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error in task dispatch loop: {e}")
|
|
|
|
|
|
async def refresh_swarm_run_status(run):
|
|
"""Update a swarm run once all known tasks have reached terminal states."""
|
|
if run.status == "stopped":
|
|
return
|
|
if not run.task_ids:
|
|
return
|
|
|
|
tasks = []
|
|
for task_id in run.task_ids:
|
|
task = await task_queue.get_task(task_id)
|
|
if task:
|
|
tasks.append(task)
|
|
|
|
if not tasks:
|
|
return
|
|
|
|
if any(task.status in {TaskStatus.PENDING, TaskStatus.ASSIGNED, TaskStatus.IN_PROGRESS} for task in tasks):
|
|
if run.status != "running":
|
|
run.status = "running"
|
|
await swarm_runtime.save_run(run)
|
|
return
|
|
|
|
if any(task.status == TaskStatus.BLOCKED for task in tasks):
|
|
blocked_with_active_child = False
|
|
for task in tasks:
|
|
if task.status != TaskStatus.BLOCKED:
|
|
continue
|
|
for child_task_id in task.child_task_ids:
|
|
child_task = await task_queue.get_task(child_task_id)
|
|
if child_task and child_task.status in {
|
|
TaskStatus.PENDING,
|
|
TaskStatus.ASSIGNED,
|
|
TaskStatus.IN_PROGRESS,
|
|
TaskStatus.BLOCKED,
|
|
}:
|
|
blocked_with_active_child = True
|
|
break
|
|
if blocked_with_active_child:
|
|
break
|
|
if blocked_with_active_child:
|
|
if run.status != "running":
|
|
run.status = "running"
|
|
await swarm_runtime.save_run(run)
|
|
return
|
|
|
|
next_status = "failed" if any(task.status == TaskStatus.FAILED for task in tasks) else "completed"
|
|
|
|
# Master review-and-iterate gate: before declaring success, optionally run the critic and
|
|
# send rejected work back to the specialists. Off unless ENABLE_REVIEW_LOOP is set, so the
|
|
# default completion semantics (and Manager contract) are unchanged.
|
|
if next_status == "completed" and review_loop_enabled():
|
|
if await maybe_run_review_cycle(run, tasks):
|
|
return
|
|
|
|
if run.status == next_status:
|
|
return
|
|
|
|
run.status = next_status
|
|
await swarm_runtime.save_run(run)
|
|
|
|
deliverable = None
|
|
final_summary = None
|
|
if next_status == "completed":
|
|
deliverable = build_run_deliverable(run, tasks)
|
|
if review_loop_enabled():
|
|
# Synthesize the specialist results into one coherent, user-facing answer.
|
|
completed = [t for t in tasks if t.status == TaskStatus.COMPLETED]
|
|
results = {t.task_id: {"result": parse_task_result(t) or {}} for t in completed}
|
|
final_summary = await master_agent.synthesize(run.objective, results)
|
|
run.metadata["final_summary"] = final_summary
|
|
await swarm_runtime.save_run(run)
|
|
|
|
# Benchmark Group B: grade the run's generated code against its held-out fixture tests in
|
|
# the sandbox. Gated (ENABLE_QUALITY_EVAL) + fixture-bound; a no-op otherwise. Never fails
|
|
# the run — a grading error just leaves quality unrecorded (reward stays NaN).
|
|
try:
|
|
quality = await evaluate_run_quality(run, tasks)
|
|
if quality:
|
|
await swarm_runtime.record_quality(run, quality)
|
|
except Exception as exc:
|
|
logger.warning("quality eval failed for run %s: %s", run.swarm_id, exc)
|
|
|
|
await swarm_runtime.emit_event(
|
|
run,
|
|
"deployment.status_changed",
|
|
payload=swarm_runtime.status_payload(
|
|
run,
|
|
phase="Deliver" if next_status == "completed" else "Execute",
|
|
deliverable=deliverable,
|
|
),
|
|
)
|
|
timeline_payload = {
|
|
"summary": final_summary or f"Swarm run {next_status}",
|
|
"status": next_status,
|
|
"task_count": len(tasks),
|
|
}
|
|
if run.metadata.get("review_summary"):
|
|
timeline_payload["review_summary"] = run.metadata["review_summary"]
|
|
await swarm_runtime.emit_event(run, "timeline.updated", payload=timeline_payload)
|
|
await maybe_emit_budget_alert(run)
|
|
|
|
|
|
# Lifespan context manager
|
|
@asynccontextmanager
|
|
async def lifespan(app: FastAPI):
|
|
"""Startup and shutdown events."""
|
|
# Startup
|
|
logger.info("Starting orchestrator...")
|
|
# Fail-closed at boot: if code-executing quality eval is enabled, isolation MUST be confirmed.
|
|
# Refuse to start otherwise — a platform safety guarantee, not an ops-only convention.
|
|
from .quality import assert_quality_eval_safe
|
|
assert_quality_eval_safe()
|
|
await redis_client.connect()
|
|
|
|
# Start background tasks
|
|
failure_task = asyncio.create_task(failure_detection_loop())
|
|
dispatch_task = asyncio.create_task(task_dispatch_loop())
|
|
|
|
yield
|
|
|
|
# Shutdown
|
|
logger.info("Shutting down orchestrator...")
|
|
failure_task.cancel()
|
|
dispatch_task.cancel()
|
|
await redis_client.disconnect()
|
|
|
|
|
|
# Create FastAPI app
|
|
app = FastAPI(title="Swarm Orchestrator", version="1.0.0", lifespan=lifespan)
|
|
|
|
|
|
# Pydantic models for API
|
|
class RegisterRequest(BaseModel):
|
|
"""Agent registration request."""
|
|
agent_id: str
|
|
capabilities: List[str]
|
|
|
|
|
|
class TaskCreateRequest(BaseModel):
|
|
"""Task creation request."""
|
|
description: str
|
|
context: Dict = Field(default_factory=dict)
|
|
max_retries: int = 3
|
|
|
|
|
|
class TaskAssignRequest(BaseModel):
|
|
"""Task assignment request."""
|
|
task_id: str
|
|
agent_id: str
|
|
|
|
|
|
class StopDeploymentRequest(BaseModel):
|
|
"""Agent Manager stop request."""
|
|
reason: str = "Heicode Manager requested stop"
|
|
manager_deployment_id: Optional[str] = None
|
|
|
|
|
|
def error_response(
|
|
status_code: int,
|
|
code: str,
|
|
message: str,
|
|
request_id: Optional[str] = None,
|
|
) -> JSONResponse:
|
|
"""Return the Agent Manager compatible error envelope."""
|
|
return JSONResponse(
|
|
status_code=status_code,
|
|
content={
|
|
"success": False,
|
|
"error": {
|
|
"code": code,
|
|
"message": message,
|
|
"request_id": request_id,
|
|
},
|
|
},
|
|
)
|
|
|
|
|
|
async def require_runtime_auth(request: Request) -> Optional[JSONResponse]:
|
|
"""Validate service-to-service runtime auth when configured."""
|
|
global _AUTH_WARNING_EMITTED
|
|
token = (
|
|
os.getenv("AGENT_RUNTIME_SERVICE_TOKEN")
|
|
or os.getenv("AGNET_RUNTIME_SERVICE_TOKEN")
|
|
)
|
|
if not token:
|
|
if not _AUTH_WARNING_EMITTED:
|
|
logger.warning(
|
|
"AGENT_RUNTIME_SERVICE_TOKEN/AGNET_RUNTIME_SERVICE_TOKEN is not configured; runtime API is in insecure dev mode"
|
|
)
|
|
_AUTH_WARNING_EMITTED = True
|
|
return None
|
|
|
|
authorization = request.headers.get("authorization", "")
|
|
expected = f"Bearer {token}"
|
|
if authorization != expected:
|
|
return error_response(
|
|
401,
|
|
"UNAUTHORIZED",
|
|
"Missing or invalid service token",
|
|
request.headers.get("x-correlation-id"),
|
|
)
|
|
return None
|
|
|
|
|
|
def request_context_headers(request: Request) -> Dict[str, Optional[str]]:
|
|
"""Capture common Agent Manager request headers for audit/debug state."""
|
|
return {
|
|
"authorization_present": "authorization" in request.headers,
|
|
"x_user_id": request.headers.get("x-user-id"),
|
|
"x_binding_scope": request.headers.get("x-binding-scope"),
|
|
"x_agent_request_id": request.headers.get("x-agent-request-id"),
|
|
"x_correlation_id": request.headers.get("x-correlation-id"),
|
|
"x_idempotency_key": request.headers.get("x-idempotency-key"),
|
|
}
|
|
|
|
|
|
def planner_fallback_enabled(body: Dict[str, Any]) -> bool:
|
|
"""Return whether the LLM planner should synthesize the task graph.
|
|
|
|
Manager-first: only when the operator opts in via ENABLE_PLANNER_FALLBACK AND the
|
|
Manager supplied no explicit agent breakdown. A Manager-provided plan is never overridden.
|
|
"""
|
|
if os.getenv("ENABLE_PLANNER_FALLBACK", "false").lower() not in {"1", "true", "yes"}:
|
|
return False
|
|
normalized = swarm_runtime.normalize_create_request(body)
|
|
plan = normalized.get("orchestration_plan") or {}
|
|
agents = plan.get("agents") or normalized.get("agents")
|
|
return not agents
|
|
|
|
|
|
async def build_planner_task_specs(run, body: Dict[str, Any], base_specs: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
|
|
"""Map planner subtasks into the build_task_descriptions spec shape.
|
|
|
|
Reuses the base spec's context so orchestration_plan/resource_grants/billing plumbing
|
|
is preserved identically to the Manager-driven path.
|
|
"""
|
|
objective = run.objective or "Complete swarm objective"
|
|
base_context = (base_specs[0].get("context") if base_specs else {}) or {}
|
|
subtasks = await master_agent.plan(run.swarm_id, objective)
|
|
|
|
prefix = f"{run.swarm_id}-"
|
|
|
|
def _local_id(value: str) -> str:
|
|
# Planner ids are prefixed with the run id; strip it so create_tasks_for_run prepends
|
|
# the swarm id exactly once (avoids doubled swarm-X-swarm-X-... task ids).
|
|
return value[len(prefix):] if value.startswith(prefix) else value
|
|
|
|
specs: List[Dict[str, Any]] = []
|
|
for index, sub in enumerate(subtasks, start=1):
|
|
role = sub.get("role") or "general"
|
|
task_id = _local_id(sub.get("subtask_id") or f"task-{index}")
|
|
specs.append({
|
|
"task_id": task_id,
|
|
"title": f"{role} task",
|
|
"description": sub.get("description") or objective,
|
|
"agent_role": role,
|
|
"required_capabilities": sub.get("required_capabilities") or [role],
|
|
"depends_on": [_local_id(dep) for dep in (sub.get("depends_on") or [])],
|
|
"parent_task_id": None,
|
|
"root_task_id": task_id,
|
|
"source": "planner",
|
|
"workflow_mode": "multi_agent",
|
|
"allow_handoff": True,
|
|
"context": {
|
|
**base_context,
|
|
"agent_role": role,
|
|
"workflow_mode": "multi_agent",
|
|
},
|
|
})
|
|
|
|
if not specs:
|
|
return base_specs
|
|
|
|
# Drop any dependency that does not reference a known subtask so a malformed plan
|
|
# cannot leave tasks permanently blocked on a phantom dependency.
|
|
known_ids = {spec["task_id"] for spec in specs}
|
|
for spec in specs:
|
|
spec["depends_on"] = [dep for dep in spec["depends_on"] if dep in known_ids]
|
|
return specs
|
|
|
|
|
|
async def create_tasks_for_run(run, body: Dict[str, Any]) -> int:
|
|
"""Create the runtime task graph and emit task.created events."""
|
|
created_count = 0
|
|
task_specs = swarm_runtime.build_task_descriptions(body)
|
|
if planner_fallback_enabled(body):
|
|
task_specs = await build_planner_task_specs(run, body, task_specs)
|
|
task_id_map = {
|
|
task_spec["task_id"]: f"{run.swarm_id}-{task_spec['task_id']}"
|
|
for task_spec in task_specs
|
|
if task_spec.get("task_id")
|
|
}
|
|
for task_spec in task_specs:
|
|
graph_task_id = task_spec.get("task_id")
|
|
runtime_task_id = task_id_map.get(graph_task_id, graph_task_id)
|
|
runtime_depends_on = [
|
|
task_id_map.get(dependency_id, dependency_id)
|
|
for dependency_id in (task_spec.get("depends_on") or [])
|
|
]
|
|
runtime_parent_task_id = task_id_map.get(task_spec.get("parent_task_id"))
|
|
runtime_root_task_id = task_id_map.get(
|
|
task_spec.get("root_task_id") or graph_task_id,
|
|
runtime_task_id,
|
|
)
|
|
task_context = {
|
|
**task_spec.get("context", {}),
|
|
"runtime_deployment_id": run.deployment_id,
|
|
"swarm_id": run.swarm_id,
|
|
"manager_deployment_id": run.manager_deployment_id,
|
|
"correlation_id": run.correlation_id,
|
|
"task_graph_id": graph_task_id,
|
|
"task_title": task_spec.get("title"),
|
|
"depends_on": runtime_depends_on,
|
|
"agent_role": task_spec.get("agent_role", "general"),
|
|
"required_capabilities": task_spec.get("required_capabilities") or [],
|
|
"parent_task_id": runtime_parent_task_id,
|
|
"root_task_id": runtime_root_task_id,
|
|
"source": task_spec.get("source", "runtime_bridge"),
|
|
"workflow_mode": task_spec.get("workflow_mode", "single_agent"),
|
|
"allow_handoff": task_spec.get("allow_handoff", False),
|
|
"model_id": (body.get("orchestration_plan") or {}).get("model_id")
|
|
or body.get("model_id"),
|
|
}
|
|
task = await task_queue.create_task(
|
|
task_id=runtime_task_id,
|
|
title=task_spec.get("title"),
|
|
description=task_spec["description"],
|
|
agent_role=task_spec.get("agent_role", "general"),
|
|
required_capabilities=task_spec.get("required_capabilities") or [],
|
|
depends_on=runtime_depends_on,
|
|
parent_task_id=runtime_parent_task_id,
|
|
root_task_id=runtime_root_task_id,
|
|
source=task_spec.get("source", "runtime_bridge"),
|
|
context=task_context,
|
|
max_retries=body.get("max_retries", 3),
|
|
)
|
|
await swarm_runtime.attach_task(run, task.task_id)
|
|
await swarm_runtime.emit_event(
|
|
run,
|
|
"task.created",
|
|
task_id=task.task_id,
|
|
payload=swarm_runtime.task_event_payload(task, task_spec),
|
|
)
|
|
TASKS_CREATED.inc()
|
|
created_count += 1
|
|
if created_count:
|
|
await swarm_runtime.emit_event(run, "timeline.updated", payload={
|
|
"summary": f"Created {created_count} task(s)",
|
|
"task_count": created_count,
|
|
})
|
|
return created_count
|
|
|
|
|
|
def task_payload(task) -> Dict[str, Any]:
|
|
"""Return the common task payload shape used by task APIs."""
|
|
return {
|
|
"task_id": task.task_id,
|
|
"title": task.title or task.context.get("task_title") or task.description[:80],
|
|
"description": task.description,
|
|
"status": task.status.value,
|
|
"agent_role": task.agent_role,
|
|
"required_capabilities": task.required_capabilities,
|
|
"depends_on": task.depends_on,
|
|
"parent_task_id": task.parent_task_id,
|
|
"root_task_id": task.root_task_id,
|
|
"source": task.source,
|
|
"task_graph_id": task.context.get("task_graph_id"),
|
|
"assigned_agent_id": task.assigned_agent_id,
|
|
"created_at": task.created_at,
|
|
"started_at": task.started_at,
|
|
"completed_at": task.completed_at,
|
|
"attempt": task.retry_count,
|
|
"max_retries": task.max_retries,
|
|
"retry_count": task.retry_count,
|
|
"blocked_reason": task.blocked_reason,
|
|
"context": task.context,
|
|
}
|
|
|
|
|
|
def handoff_workflow_enabled(task) -> bool:
|
|
"""Return whether a task should use the DAG-style handoff workflow."""
|
|
return (
|
|
os.getenv("ENABLE_SUBTASK_HANDOFF", "false").lower() in {"1", "true", "yes"}
|
|
and task.context.get("workflow_mode") == "multi_agent"
|
|
)
|
|
|
|
|
|
def summarize_task_result(result: Any) -> str:
|
|
"""Return a user-facing summary for task completion events."""
|
|
if isinstance(result, dict):
|
|
summary = result.get("summary")
|
|
if isinstance(summary, str) and summary.strip():
|
|
return summary.strip()
|
|
subtasks = result.get("subtasks") or []
|
|
for subtask_result in reversed(subtasks):
|
|
summary = subtask_result.get("summary")
|
|
if isinstance(summary, str) and summary.strip():
|
|
return summary.strip()
|
|
if isinstance(result, str) and result.strip():
|
|
return result.strip()[:500]
|
|
return "Task completed"
|
|
|
|
|
|
def parse_task_result(task) -> Optional[Dict[str, Any]]:
|
|
"""Parse a task result into a structured dictionary when possible."""
|
|
if not getattr(task, "result", None):
|
|
return None
|
|
if isinstance(task.result, dict):
|
|
return task.result
|
|
if isinstance(task.result, str):
|
|
try:
|
|
parsed = json.loads(task.result)
|
|
return parsed if isinstance(parsed, dict) else {"summary": str(parsed)}
|
|
except json.JSONDecodeError:
|
|
return {"summary": task.result}
|
|
return None
|
|
|
|
|
|
def build_deliverable_fact(
|
|
result: Optional[Dict[str, Any]],
|
|
artifact: Optional[Dict[str, Any]] = None,
|
|
) -> Dict[str, Any]:
|
|
"""Build the structured deliverable fact consumed by Manager."""
|
|
result = result or {}
|
|
files_modified = list(result.get("files_modified") or [])
|
|
files_deleted = list(result.get("files_deleted") or [])
|
|
has_branch = bool(result.get("git_branch"))
|
|
commit_sha = result.get("commit_sha") or ""
|
|
artifact_ids = [artifact["artifact_id"]] if artifact else []
|
|
has_deliverable = bool(artifact_ids or has_branch or files_modified or files_deleted)
|
|
summary_only = not bool(has_branch or files_modified or files_deleted)
|
|
return {
|
|
"has_deliverable": has_deliverable,
|
|
"summary_only": summary_only,
|
|
"artifact_ids": artifact_ids,
|
|
"files_modified": files_modified,
|
|
"files_deleted": files_deleted,
|
|
"has_diff": bool(has_branch or commit_sha or files_modified or files_deleted),
|
|
"commit_sha": commit_sha,
|
|
}
|
|
|
|
|
|
def build_run_deliverable(run, tasks: List[Any]) -> Dict[str, Any]:
|
|
"""Aggregate deliverable facts across all runtime tasks."""
|
|
artifact_ids: List[str] = []
|
|
files_modified: List[str] = []
|
|
files_deleted: List[str] = []
|
|
commit_sha = ""
|
|
has_deliverable = False
|
|
summary_only = True
|
|
|
|
for task in tasks:
|
|
parsed_result = parse_task_result(task)
|
|
artifact = build_result_artifact(run, task, parsed_result) if parsed_result else None
|
|
fact = build_deliverable_fact(parsed_result, artifact)
|
|
artifact_ids.extend(fact["artifact_ids"])
|
|
files_modified.extend(fact["files_modified"])
|
|
files_deleted.extend(fact["files_deleted"])
|
|
has_deliverable = has_deliverable or fact["has_deliverable"]
|
|
summary_only = summary_only and fact["summary_only"]
|
|
if fact["commit_sha"] and not commit_sha:
|
|
commit_sha = fact["commit_sha"]
|
|
|
|
return {
|
|
"has_deliverable": has_deliverable,
|
|
"summary_only": summary_only,
|
|
"artifact_ids": artifact_ids,
|
|
"files_modified": files_modified,
|
|
"files_deleted": files_deleted,
|
|
"has_diff": bool(commit_sha or files_modified or files_deleted),
|
|
"commit_sha": commit_sha,
|
|
}
|
|
|
|
|
|
def infer_tool_count(result: Optional[Dict[str, Any]]) -> int:
|
|
"""Infer tool usage count from a task result payload."""
|
|
if not result:
|
|
return 0
|
|
if isinstance(result.get("tool_count"), int):
|
|
return result["tool_count"]
|
|
tool_calls = result.get("tool_calls")
|
|
if isinstance(tool_calls, list):
|
|
return len(tool_calls)
|
|
return 0
|
|
|
|
|
|
def collect_run_artifacts(run, tasks: List[Any]) -> List[Dict[str, Any]]:
|
|
"""Collect stable artifact descriptors from runtime tasks."""
|
|
artifacts = []
|
|
for task in tasks:
|
|
parsed_result = parse_task_result(task)
|
|
artifact = build_result_artifact(run, task, parsed_result) if parsed_result else None
|
|
if artifact:
|
|
artifacts.append(artifact)
|
|
return artifacts
|
|
|
|
|
|
def build_workflow_phases(run, tasks: List[Any], events: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
|
|
"""Build the fixed workflow phase view for the swarm runtime."""
|
|
agents_by_role: Dict[str, List[Dict[str, Any]]] = defaultdict(list)
|
|
for task in tasks:
|
|
parsed_result = parse_task_result(task)
|
|
artifact = build_result_artifact(run, task, parsed_result) if parsed_result else None
|
|
usage = (parsed_result or {}).get("usage") or {}
|
|
elapsed_seconds = 0
|
|
if task.started_at and task.completed_at and task.completed_at >= task.started_at:
|
|
elapsed_seconds = int(task.completed_at - task.started_at)
|
|
agents_by_role[task.agent_role].append({
|
|
"agent_id": task.task_id,
|
|
"name": task.title or task.task_id,
|
|
"role": task.agent_role,
|
|
"status": task.status.value,
|
|
"tokens": usage.get("model_tokens", 0),
|
|
"tools": infer_tool_count(parsed_result),
|
|
"elapsed_seconds": elapsed_seconds,
|
|
"artifact_ids": [artifact["artifact_id"]] if artifact else [],
|
|
})
|
|
|
|
handoff_seen = any(event.get("event_type", "").startswith("handoff.") for event in events)
|
|
approvals = list(run.approvals.values())
|
|
task_statuses = {task.status.value for task in tasks}
|
|
terminal_statuses = {"completed", "failed", "cancelled"}
|
|
|
|
def phase_status(name: str) -> str:
|
|
if name == "Plan":
|
|
return "completed"
|
|
if name == "Dispatch":
|
|
if not tasks:
|
|
return "pending"
|
|
if any(status in {"assigned", "in_progress", "blocked", "completed", "failed", "cancelled"} for status in task_statuses):
|
|
return "completed"
|
|
return "running"
|
|
if name == "Execute":
|
|
if "failed" in task_statuses or run.status == "failed":
|
|
return "failed"
|
|
if task_statuses and task_statuses.issubset(terminal_statuses):
|
|
return "completed"
|
|
if any(status in {"assigned", "in_progress", "blocked"} for status in task_statuses):
|
|
return "running"
|
|
return "pending"
|
|
if name == "Handoff":
|
|
if not handoff_seen:
|
|
return "pending"
|
|
if any(event.get("event_type") == "handoff.completed" for event in events):
|
|
return "completed"
|
|
return "running"
|
|
if name == "Review":
|
|
if run.status == "waiting_approval":
|
|
return "running"
|
|
if run.status == "blocked":
|
|
return "failed"
|
|
if approvals:
|
|
return "completed"
|
|
return "pending"
|
|
if name == "Deliver":
|
|
if run.status == "completed":
|
|
return "completed"
|
|
if run.status in {"failed", "stopped", "blocked"}:
|
|
return "failed" if run.status == "failed" else run.status
|
|
if run.status == "running":
|
|
return "running"
|
|
return "pending"
|
|
return "pending"
|
|
|
|
phase_agents = {
|
|
"Plan": [],
|
|
"Dispatch": [],
|
|
"Execute": [agent for group in agents_by_role.values() for agent in group],
|
|
"Handoff": [agent for group in agents_by_role.values() for agent in group if agent["status"] == "blocked"] if handoff_seen else [],
|
|
"Review": [],
|
|
"Deliver": [],
|
|
}
|
|
return [
|
|
{
|
|
"phase_id": phase_name.lower(),
|
|
"name": phase_name,
|
|
"status": phase_status(phase_name),
|
|
"agents": phase_agents.get(phase_name, []),
|
|
}
|
|
for phase_name in ["Plan", "Dispatch", "Execute", "Handoff", "Review", "Deliver"]
|
|
]
|
|
|
|
|
|
def build_result_artifact(run, task, result: Any) -> Optional[Dict[str, Any]]:
|
|
"""Build a callback-safe artifact for completed tasks."""
|
|
if not isinstance(result, dict):
|
|
return None
|
|
|
|
summary = summarize_task_result(result)
|
|
base_metadata = {
|
|
"redacted": True,
|
|
"agent_role": task.agent_role,
|
|
"runtime_deployment_id": run.deployment_id,
|
|
"manager_deployment_id": run.manager_deployment_id,
|
|
"task_id": task.task_id,
|
|
"parent_task_id": task.parent_task_id,
|
|
"root_task_id": task.root_task_id,
|
|
}
|
|
|
|
if result.get("git_branch"):
|
|
return {
|
|
"artifact_id": f"art_{task.task_id}",
|
|
"artifact_type": "code_patch",
|
|
"title": task.title or "Agent task result branch",
|
|
"summary": summary,
|
|
"uri": f"git://repo#{result.get('git_branch')}",
|
|
"checksum": result.get("commit_sha"),
|
|
"metadata": {
|
|
**base_metadata,
|
|
"commit_sha": result.get("commit_sha"),
|
|
"git_branch": result.get("git_branch"),
|
|
},
|
|
}
|
|
|
|
files_modified = result.get("files_modified") or []
|
|
files_deleted = result.get("files_deleted") or []
|
|
artifact_type = "deployment_manifest" if files_modified or files_deleted else "document"
|
|
return {
|
|
"artifact_id": f"art_{task.task_id}_summary",
|
|
"artifact_type": artifact_type,
|
|
"title": task.title or task.context.get("task_title") or "Task execution result",
|
|
"summary": summary,
|
|
"uri": f"runtime://{run.swarm_id}/artifacts/{task.task_id}",
|
|
"metadata": {
|
|
**base_metadata,
|
|
"files_modified": files_modified,
|
|
"files_deleted": files_deleted,
|
|
"changes": result.get("changes"),
|
|
"git_skipped": result.get("git_skipped"),
|
|
"git_error": result.get("git_error"),
|
|
},
|
|
}
|
|
|
|
|
|
async def emit_task_completion_events(run, task, agent_id: str, result: Any):
|
|
"""Emit the completion callback set expected by Agent Manager."""
|
|
summary = summarize_task_result(result)
|
|
artifact = build_result_artifact(run, task, result)
|
|
deliverable = build_deliverable_fact(result, artifact)
|
|
payload = {
|
|
"task_id": task.task_id,
|
|
"agent_role": task.agent_role,
|
|
"agent_id": agent_id,
|
|
"status": "completed",
|
|
"summary": summary,
|
|
"parent_task_id": task.parent_task_id,
|
|
"root_task_id": task.root_task_id,
|
|
"deliverable": deliverable,
|
|
}
|
|
await swarm_runtime.emit_event(
|
|
run,
|
|
"task.completed",
|
|
task_id=task.task_id,
|
|
agent_instance_id=agent_id,
|
|
payload=payload,
|
|
)
|
|
|
|
if artifact:
|
|
await swarm_runtime.emit_event(
|
|
run,
|
|
"artifact.created",
|
|
task_id=task.task_id,
|
|
agent_instance_id=agent_id,
|
|
artifact=artifact,
|
|
)
|
|
|
|
await swarm_runtime.emit_event(
|
|
run,
|
|
"timeline.updated",
|
|
task_id=task.task_id,
|
|
agent_instance_id=agent_id,
|
|
payload={
|
|
"summary": summary,
|
|
"task_id": task.task_id,
|
|
"status": "completed",
|
|
},
|
|
)
|
|
|
|
|
|
async def finalize_parent_after_child(run, child_task, agent_id: str, success: bool, summary: str):
|
|
"""Resolve a blocked parent task after a delegated child reaches a terminal state."""
|
|
parent_task_id = child_task.parent_task_id
|
|
if not parent_task_id:
|
|
return
|
|
|
|
parent_task = await task_queue.get_task(parent_task_id)
|
|
if not parent_task or parent_task.status != TaskStatus.BLOCKED:
|
|
return
|
|
|
|
payload = {
|
|
"task_id": parent_task.task_id,
|
|
"agent_role": parent_task.agent_role,
|
|
"delegated_child_task_id": child_task.task_id,
|
|
"status": "completed" if success else "failed",
|
|
"summary": summary if success else None,
|
|
"reason": None if success else summary,
|
|
}
|
|
if success:
|
|
await task_queue.complete_task(
|
|
parent_task.task_id,
|
|
json.dumps(
|
|
{
|
|
"delegated_to": child_task.task_id,
|
|
"result": child_task.result,
|
|
"summary": summary,
|
|
}
|
|
),
|
|
)
|
|
await swarm_runtime.emit_event(
|
|
run,
|
|
"handoff.completed",
|
|
task_id=parent_task.task_id,
|
|
agent_instance_id=agent_id,
|
|
payload={
|
|
**payload,
|
|
"child_task_id": child_task.task_id,
|
|
"parent_task_id": parent_task.task_id,
|
|
"from_role": parent_task.agent_role,
|
|
"to_role": child_task.agent_role,
|
|
},
|
|
)
|
|
await swarm_runtime.emit_event(
|
|
run,
|
|
"task.completed",
|
|
task_id=parent_task.task_id,
|
|
agent_instance_id=agent_id,
|
|
payload=payload,
|
|
)
|
|
else:
|
|
parent_task.retry_count = max(parent_task.max_retries - 1, 0)
|
|
await task_queue._save_task(parent_task)
|
|
await task_queue.fail_task(parent_task.task_id, summary)
|
|
await swarm_runtime.emit_event(
|
|
run,
|
|
"task.failed",
|
|
task_id=parent_task.task_id,
|
|
agent_instance_id=agent_id,
|
|
payload=payload,
|
|
)
|
|
await swarm_runtime.emit_event(
|
|
run,
|
|
"timeline.updated",
|
|
task_id=parent_task.task_id,
|
|
agent_instance_id=agent_id,
|
|
payload={
|
|
"summary": (
|
|
f"Delegated parent task {parent_task.task_id} completed"
|
|
if success else f"Delegated parent task {parent_task.task_id} failed"
|
|
),
|
|
"task_id": parent_task.task_id,
|
|
"child_task_id": child_task.task_id,
|
|
"status": "completed" if success else "failed",
|
|
},
|
|
)
|
|
|
|
|
|
async def create_swarm_run_from_request(
|
|
request: Request,
|
|
body: Dict[str, Any],
|
|
x_correlation_id: Optional[str],
|
|
x_idempotency_key: Optional[str],
|
|
):
|
|
"""Shared implementation for /api/swarms and /api/agnet/deployments."""
|
|
auth_error = await require_runtime_auth(request)
|
|
if auth_error:
|
|
return auth_error
|
|
|
|
body = swarm_runtime.normalize_create_request(body)
|
|
correlation_id = x_correlation_id or (body.get("metadata") or {}).get("correlation_id")
|
|
try:
|
|
swarm_runtime.validate_create_request(body)
|
|
except RuntimeValidationError as exc:
|
|
return error_response(422, exc.code, exc.message, correlation_id)
|
|
|
|
body.setdefault("metadata", {})
|
|
body["metadata"]["runtime_headers"] = request_context_headers(request)
|
|
|
|
run, created = await swarm_runtime.get_or_create_run(
|
|
body=body,
|
|
idempotency_key=x_idempotency_key,
|
|
correlation_id=correlation_id,
|
|
)
|
|
|
|
if created and run.status == "running":
|
|
await create_tasks_for_run(run, body)
|
|
|
|
return {
|
|
"success": True,
|
|
"data": {
|
|
"deployment_id": run.deployment_id,
|
|
"runtime_deployment_id": run.deployment_id,
|
|
"manager_deployment_id": run.manager_deployment_id,
|
|
"swarm_id": run.swarm_id,
|
|
"mode": run.mode,
|
|
"status": run.status,
|
|
"runtime_execution_status": run.status,
|
|
"created": created,
|
|
},
|
|
}
|
|
|
|
|
|
async def stop_swarm_run(deployment_or_swarm_id: str, reason: str):
|
|
"""Stop a run and cancel non-terminal tasks."""
|
|
run = await swarm_runtime.stop_run(deployment_or_swarm_id, reason)
|
|
if not run:
|
|
raise HTTPException(status_code=404, detail="Runtime deployment not found")
|
|
for task_id in run.task_ids:
|
|
# Ask the owning agent to abort the in-flight task before we mark it cancelled,
|
|
# so the merged agent runtime tears down its execution coroutine promptly.
|
|
task = await task_queue.get_task(task_id)
|
|
if task and task.assigned_agent_id and task.status in {TaskStatus.ASSIGNED, TaskStatus.IN_PROGRESS}:
|
|
await manager.send_message(task.assigned_agent_id, {
|
|
"type": "cancel_task",
|
|
"task_id": task_id,
|
|
})
|
|
await task_queue.cancel_task(task_id, reason)
|
|
await swarm_runtime.emit_event(run, "timeline.updated", payload={
|
|
"summary": "Swarm stopped by Manager",
|
|
"reason": reason,
|
|
})
|
|
return run
|
|
|
|
|
|
async def maybe_emit_budget_alert(run):
|
|
"""Emit a simple duration-based budget alert once per run."""
|
|
budget = ((run.request_body.get("orchestration_plan") or {}).get("budget") or {})
|
|
duration_budget = budget.get("duration_seconds") or budget.get("max_duration_seconds")
|
|
if not duration_budget or run.metadata.get("budget_duration_alerted"):
|
|
return
|
|
elapsed = time.time() - run.created_at
|
|
try:
|
|
ratio = elapsed / float(duration_budget)
|
|
except (TypeError, ValueError, ZeroDivisionError):
|
|
return
|
|
if ratio < 0.8:
|
|
return
|
|
|
|
run.metadata["budget_duration_alerted"] = True
|
|
await swarm_runtime.save_run(run)
|
|
await swarm_runtime.emit_event(run, "budget.alert", payload={
|
|
"budget_type": "duration",
|
|
"threshold": 0.8,
|
|
"ratio": ratio,
|
|
"elapsed_seconds": elapsed,
|
|
"limit_seconds": duration_budget,
|
|
})
|
|
|
|
|
|
async def emit_usage_event(run, task, agent_id: str, result: Any):
|
|
"""Emit model/runtime usage when the agent returned provider usage."""
|
|
if not isinstance(result, dict):
|
|
return
|
|
usage = result.get("usage")
|
|
if not isinstance(usage, dict):
|
|
return
|
|
|
|
plan = run.request_body.get("orchestration_plan") or {}
|
|
budget = plan.get("budget") or {}
|
|
payload = {
|
|
"model_id": usage.get("model_id") or task.context.get("model_id"),
|
|
"model_tokens": usage.get("model_tokens", 0),
|
|
"prompt_tokens": usage.get("prompt_tokens", 0),
|
|
"completion_tokens": usage.get("completion_tokens", 0),
|
|
"model_cost_usd": usage.get("model_cost_usd", 0),
|
|
"runtime_seconds": usage.get("runtime_seconds", 0),
|
|
"billing_source": usage.get("billing_source", "unknown"),
|
|
"manager_deployment_id": run.manager_deployment_id,
|
|
"swarm_id": run.swarm_id,
|
|
"task_id": task.task_id,
|
|
"agent_role": task.agent_role,
|
|
"correlation_id": run.correlation_id,
|
|
"budget": {
|
|
"max_tokens": budget.get("max_tokens") or budget.get("token_limit"),
|
|
"max_cost_usd": budget.get("max_cost_usd"),
|
|
"consumed_usd": usage.get("model_cost_usd", 0),
|
|
"remaining_usd": (
|
|
budget.get("max_cost_usd") - usage.get("model_cost_usd", 0)
|
|
if isinstance(budget.get("max_cost_usd"), (int, float))
|
|
else None
|
|
),
|
|
},
|
|
}
|
|
await swarm_runtime.emit_event(
|
|
run,
|
|
"budget.alert",
|
|
task_id=task.task_id,
|
|
agent_instance_id=agent_id,
|
|
payload=payload,
|
|
)
|
|
await swarm_runtime.emit_event(
|
|
run,
|
|
"timeline.updated",
|
|
task_id=task.task_id,
|
|
agent_instance_id=agent_id,
|
|
payload={
|
|
"summary": "Usage updated",
|
|
"task_id": task.task_id,
|
|
"model_id": payload["model_id"],
|
|
"model_tokens": payload["model_tokens"],
|
|
"model_cost_usd": payload["model_cost_usd"],
|
|
"runtime_seconds": payload["runtime_seconds"],
|
|
},
|
|
)
|
|
|
|
|
|
@app.exception_handler(HTTPException)
|
|
async def runtime_http_exception_handler(request: Request, exc: HTTPException):
|
|
"""Use Manager error envelopes for runtime API errors."""
|
|
if request.url.path.startswith("/api/"):
|
|
code = "NOT_FOUND" if exc.status_code == 404 else "HTTP_ERROR"
|
|
return error_response(
|
|
exc.status_code,
|
|
code,
|
|
str(exc.detail),
|
|
request.headers.get("x-correlation-id"),
|
|
)
|
|
return JSONResponse(status_code=exc.status_code, content={"detail": exc.detail})
|
|
|
|
|
|
# REST API endpoints
|
|
@app.get("/")
|
|
async def root():
|
|
"""Health check endpoint."""
|
|
return {"status": "ok", "service": "swarm-orchestrator"}
|
|
|
|
|
|
@app.get("/health")
|
|
async def health():
|
|
"""Detailed health check."""
|
|
try:
|
|
await redis_client.client.ping()
|
|
redis_status = "connected"
|
|
except Exception:
|
|
redis_status = "disconnected"
|
|
|
|
return {
|
|
"status": "ok",
|
|
"redis": redis_status,
|
|
"active_connections": len(manager.active_connections)
|
|
}
|
|
|
|
async def get_runtime_run_or_404(request: Request, deployment_id: str):
|
|
"""Load a runtime run after validating service auth."""
|
|
auth_error = await require_runtime_auth(request)
|
|
if auth_error:
|
|
return None, auth_error
|
|
run = await swarm_runtime.get_run_by_identifier(deployment_id)
|
|
if not run:
|
|
raise HTTPException(status_code=404, detail="Swarm run not found")
|
|
return run, None
|
|
|
|
|
|
async def load_runtime_tasks(run) -> List[Any]:
|
|
"""Load all known task records for a swarm run."""
|
|
tasks = [await task_queue.get_task(task_id) for task_id in run.task_ids]
|
|
return [task for task in tasks if task]
|
|
|
|
|
|
def runtime_deployment_response(run) -> Dict[str, Any]:
|
|
"""Return the standard runtime deployment envelope."""
|
|
return {
|
|
"deployment_id": run.deployment_id,
|
|
"runtime_deployment_id": run.deployment_id,
|
|
"manager_deployment_id": run.manager_deployment_id,
|
|
"swarm_id": run.swarm_id,
|
|
"mode": run.mode,
|
|
"objective": run.objective,
|
|
"status": run.status,
|
|
"runtime_execution_status": run.status,
|
|
"created_at": run.created_at,
|
|
"updated_at": run.updated_at,
|
|
"task_count": len(run.task_ids),
|
|
"approvals": list(run.approvals.values()),
|
|
}
|
|
|
|
|
|
async def build_runtime_workflow(run) -> Dict[str, Any]:
|
|
"""Build the workflow-oriented runtime response consumed by Manager."""
|
|
tasks = await load_runtime_tasks(run)
|
|
events_data = await swarm_runtime.list_events(run.swarm_id, limit=500)
|
|
events = events_data["events"]
|
|
artifacts = collect_run_artifacts(run, tasks)
|
|
duration_seconds = int(time.time() - run.created_at)
|
|
tokens = 0
|
|
tools = 0
|
|
for task in tasks:
|
|
parsed = parse_task_result(task) or {}
|
|
usage = parsed.get("usage") or {}
|
|
tokens += int(usage.get("model_tokens", 0) or 0)
|
|
tools += infer_tool_count(parsed)
|
|
summary = run.objective
|
|
for event in reversed(events):
|
|
if event.get("event_type") == "timeline.updated":
|
|
summary = (event.get("payload") or {}).get("summary") or summary
|
|
if summary:
|
|
break
|
|
return {
|
|
"workflow_id": run.swarm_id,
|
|
"deployment_id": run.deployment_id,
|
|
"runtime_deployment_id": run.deployment_id,
|
|
"manager_deployment_id": run.manager_deployment_id,
|
|
"task_id": run.manager_deployment_id or run.deployment_id,
|
|
"title": run.objective,
|
|
"mode": run.mode,
|
|
"status": run.status,
|
|
"runtime_execution_status": run.status,
|
|
"summary": summary,
|
|
"agent_count": len(tasks),
|
|
"tokens": tokens,
|
|
"tools": tools,
|
|
"elapsed_seconds": duration_seconds,
|
|
"phases": build_workflow_phases(run, tasks, events),
|
|
"artifacts": artifacts,
|
|
}
|
|
|
|
|
|
async def build_runtime_metrics(run, window: str, step: str) -> Dict[str, Any]:
|
|
"""Build metrics for the current swarm run."""
|
|
tasks = await load_runtime_tasks(run)
|
|
agents = await agent_registry.get_all_agents()
|
|
status_counts: Dict[str, int] = {}
|
|
for task in tasks:
|
|
status_counts[task.status.value] = status_counts.get(task.status.value, 0) + 1
|
|
completed_durations = [
|
|
task.completed_at - task.started_at
|
|
for task in tasks
|
|
if task.completed_at and task.started_at and task.completed_at >= task.started_at
|
|
]
|
|
duration_seconds = time.time() - run.created_at
|
|
budget = ((run.request_body.get("orchestration_plan") or {}).get("budget") or {})
|
|
duration_budget = budget.get("duration_seconds") or budget.get("max_duration_seconds")
|
|
budget_ratio = duration_seconds / float(duration_budget) if duration_budget else None
|
|
return {
|
|
"deployment_id": run.deployment_id,
|
|
"swarm_id": run.swarm_id,
|
|
"mode": run.mode,
|
|
"window": window,
|
|
"step": step,
|
|
"status": run.status,
|
|
"runtime_execution_status": run.status,
|
|
"runtime_duration_seconds": duration_seconds,
|
|
"tasks_total": len(tasks),
|
|
"tasks_by_status": status_counts,
|
|
"agents_connected": len(manager.active_connections),
|
|
"agents_registered": len(agents),
|
|
"average_task_duration_seconds": (
|
|
sum(completed_durations) / len(completed_durations)
|
|
if completed_durations else 0
|
|
),
|
|
"budget": {
|
|
"duration_seconds": duration_budget,
|
|
"duration_ratio": budget_ratio,
|
|
},
|
|
}
|
|
|
|
|
|
async def build_runtime_diagnostics(run) -> Dict[str, Any]:
|
|
"""Build the diagnostics response for a swarm run."""
|
|
tasks = await load_runtime_tasks(run)
|
|
callback_attempts = list(run.metadata.get("callback_attempts") or [])
|
|
failures = []
|
|
for task in tasks:
|
|
if task.status.value in {"failed", "blocked"}:
|
|
failures.append({
|
|
"task_id": task.task_id,
|
|
"status": task.status.value,
|
|
"blocked_reason": task.blocked_reason,
|
|
"result": parse_task_result(task),
|
|
})
|
|
return {
|
|
"deployment_id": run.deployment_id,
|
|
"runtime_deployment_id": run.deployment_id,
|
|
"manager_deployment_id": run.manager_deployment_id,
|
|
"swarm_id": run.swarm_id,
|
|
"mode": run.mode,
|
|
"status": run.status,
|
|
"runtime_execution_status": run.status,
|
|
"correlation_id": run.correlation_id,
|
|
"request_headers": run.metadata.get("runtime_headers") or {},
|
|
"callback": {
|
|
"url": run.callback.url,
|
|
"subscribed_events": run.callback.subscribed_events,
|
|
"attempts": callback_attempts,
|
|
},
|
|
"approvals": list(run.approvals.values()),
|
|
"error_context": {
|
|
"task_failures": failures,
|
|
"callback_failures": [item for item in callback_attempts if item.get("status") == "failed"],
|
|
},
|
|
"request_body": run.request_body,
|
|
}
|
|
|
|
|
|
@app.get("/api/agent/health")
|
|
@app.get("/api/agnet/health")
|
|
async def agnet_health():
|
|
"""Agent Manager compatible health endpoint."""
|
|
return await swarm_runtime.health()
|
|
|
|
|
|
@app.post("/api/agent/swarm/deployments")
|
|
@app.post("/api/swarms")
|
|
@app.post("/api/agnet/deployments")
|
|
async def create_swarm(
|
|
request: Request,
|
|
x_correlation_id: Optional[str] = Header(default=None),
|
|
x_idempotency_key: Optional[str] = Header(default=None),
|
|
):
|
|
"""Create a swarm runtime deployment."""
|
|
body = await request.json()
|
|
return await create_swarm_run_from_request(
|
|
request,
|
|
body,
|
|
x_correlation_id,
|
|
x_idempotency_key,
|
|
)
|
|
|
|
|
|
@app.get("/api/agent/swarm/deployments/{deployment_id}")
|
|
@app.get("/api/swarms/{deployment_id}")
|
|
@app.get("/api/agnet/deployments/{deployment_id}")
|
|
async def get_swarm(deployment_id: str, request: Request):
|
|
"""Return a runtime deployment summary."""
|
|
run, auth_error = await get_runtime_run_or_404(request, deployment_id)
|
|
if auth_error:
|
|
return auth_error
|
|
return {"success": True, "data": runtime_deployment_response(run)}
|
|
|
|
|
|
@app.get("/api/agent/swarm/deployments/{deployment_id}/tasks")
|
|
@app.get("/api/swarms/{deployment_id}/tasks")
|
|
@app.get("/api/agnet/deployments/{deployment_id}/tasks")
|
|
async def get_swarm_tasks(deployment_id: str, request: Request):
|
|
"""Return the runtime task graph for a swarm deployment."""
|
|
run, auth_error = await get_runtime_run_or_404(request, deployment_id)
|
|
if auth_error:
|
|
return auth_error
|
|
tasks = [task_payload(task) for task in await load_runtime_tasks(run)]
|
|
return {
|
|
"success": True,
|
|
"data": {
|
|
"deployment_id": run.deployment_id,
|
|
"swarm_id": run.swarm_id,
|
|
"tasks": tasks,
|
|
},
|
|
}
|
|
|
|
|
|
@app.get("/api/agent/swarm/deployments/{deployment_id}/logs")
|
|
@app.get("/api/swarms/{deployment_id}/logs")
|
|
@app.get("/api/agnet/deployments/{deployment_id}/logs")
|
|
async def get_swarm_logs(
|
|
deployment_id: str,
|
|
request: Request,
|
|
limit: int = 100,
|
|
cursor: Optional[str] = None,
|
|
):
|
|
"""Return callback/timeline events as the runtime log stream."""
|
|
run, auth_error = await get_runtime_run_or_404(request, deployment_id)
|
|
if auth_error:
|
|
return auth_error
|
|
logs = await swarm_runtime.list_events(run.swarm_id, limit=limit, cursor=cursor)
|
|
return {"success": True, "data": logs}
|
|
|
|
|
|
@app.get("/api/agent/swarm/deployments/{deployment_id}/events")
|
|
@app.get("/api/swarms/{deployment_id}/events")
|
|
@app.get("/api/agnet/deployments/{deployment_id}/events")
|
|
async def get_swarm_events(
|
|
deployment_id: str,
|
|
request: Request,
|
|
limit: int = 100,
|
|
cursor: Optional[str] = None,
|
|
):
|
|
"""Return raw runtime events for a swarm deployment."""
|
|
run, auth_error = await get_runtime_run_or_404(request, deployment_id)
|
|
if auth_error:
|
|
return auth_error
|
|
events = await swarm_runtime.list_events(run.swarm_id, limit=limit, cursor=cursor)
|
|
return {"success": True, "data": events}
|
|
|
|
|
|
@app.get("/api/agent/swarm/deployments/{deployment_id}/metrics")
|
|
@app.get("/api/swarms/{deployment_id}/metrics")
|
|
@app.get("/api/agnet/deployments/{deployment_id}/metrics")
|
|
async def get_swarm_metrics(
|
|
deployment_id: str,
|
|
request: Request,
|
|
window: str = "15m",
|
|
step: str = "60s",
|
|
):
|
|
"""Return lightweight run/task/agent metrics for Manager UI."""
|
|
run, auth_error = await get_runtime_run_or_404(request, deployment_id)
|
|
if auth_error:
|
|
return auth_error
|
|
return {"success": True, "data": await build_runtime_metrics(run, window, step)}
|
|
|
|
|
|
@app.get("/api/agent/swarm/deployments/{deployment_id}/workflow")
|
|
@app.get("/api/swarms/{deployment_id}/workflow")
|
|
@app.get("/api/agnet/deployments/{deployment_id}/workflow")
|
|
async def get_swarm_workflow(deployment_id: str, request: Request):
|
|
"""Return the workflow-oriented summary of a swarm deployment."""
|
|
run, auth_error = await get_runtime_run_or_404(request, deployment_id)
|
|
if auth_error:
|
|
return auth_error
|
|
return {"success": True, "data": await build_runtime_workflow(run)}
|
|
|
|
|
|
@app.get("/api/agent/swarm/deployments/{deployment_id}/diagnostics")
|
|
@app.get("/api/swarms/{deployment_id}/diagnostics")
|
|
@app.get("/api/agnet/deployments/{deployment_id}/diagnostics")
|
|
async def get_swarm_diagnostics(deployment_id: str, request: Request):
|
|
"""Return runtime diagnostics for a swarm deployment."""
|
|
run, auth_error = await get_runtime_run_or_404(request, deployment_id)
|
|
if auth_error:
|
|
return auth_error
|
|
return {"success": True, "data": await build_runtime_diagnostics(run)}
|
|
|
|
|
|
@app.post("/api/agent/swarm/deployments/{deployment_id}/stop")
|
|
@app.post("/api/swarms/{deployment_id}/stop")
|
|
@app.post("/api/agnet/deployments/{deployment_id}/stop")
|
|
async def stop_swarm(
|
|
deployment_id: str,
|
|
request: Request,
|
|
stop_request: Optional[StopDeploymentRequest] = None,
|
|
):
|
|
"""Stop a runtime deployment."""
|
|
auth_error = await require_runtime_auth(request)
|
|
if auth_error:
|
|
return auth_error
|
|
try:
|
|
body = await request.json() if request.method == "POST" else {}
|
|
except json.JSONDecodeError:
|
|
body = {}
|
|
reason = (
|
|
(stop_request.reason if stop_request else None)
|
|
or body.get("reason")
|
|
or "Heicode Manager requested stop"
|
|
)
|
|
run = await stop_swarm_run(deployment_id, reason)
|
|
return {
|
|
"success": True,
|
|
"data": {
|
|
"deployment_id": run.deployment_id,
|
|
"runtime_deployment_id": run.deployment_id,
|
|
"manager_deployment_id": run.manager_deployment_id,
|
|
"swarm_id": run.swarm_id,
|
|
"status": run.status,
|
|
"runtime_execution_status": run.status,
|
|
},
|
|
}
|
|
|
|
|
|
@app.post("/api/agent/swarm/deployments/{deployment_id}/approvals/{approval_id}")
|
|
@app.post("/api/swarms/{deployment_id}/approvals/{approval_id}")
|
|
@app.post("/api/agnet/deployments/{deployment_id}/approvals/{approval_id}")
|
|
async def decide_swarm_approval(
|
|
deployment_id: str,
|
|
approval_id: str,
|
|
request: Request,
|
|
):
|
|
"""Receive an approval decision from Manager."""
|
|
auth_error = await require_runtime_auth(request)
|
|
if auth_error:
|
|
return auth_error
|
|
decision = await request.json()
|
|
run = await swarm_runtime.record_approval_decision(deployment_id, approval_id, decision)
|
|
if not run:
|
|
raise HTTPException(status_code=404, detail="Swarm run not found")
|
|
|
|
if run.status == "running" and not run.task_ids:
|
|
await create_tasks_for_run(run, run.request_body or {
|
|
"orchestration_plan": {"objective": run.objective},
|
|
})
|
|
|
|
return {
|
|
"success": True,
|
|
"data": {
|
|
"deployment_id": run.deployment_id,
|
|
"runtime_deployment_id": run.deployment_id,
|
|
"manager_deployment_id": run.manager_deployment_id,
|
|
"swarm_id": run.swarm_id,
|
|
"approval_id": approval_id,
|
|
"status": run.status,
|
|
"runtime_execution_status": run.status,
|
|
},
|
|
}
|
|
|
|
|
|
@app.get("/agents")
|
|
async def list_agents():
|
|
"""List all registered agents."""
|
|
agents = await agent_registry.get_all_agents()
|
|
return {"agents": [agent.model_dump() for agent in agents]}
|
|
|
|
|
|
@app.get("/agents/idle")
|
|
async def list_idle_agents():
|
|
"""List all idle agents."""
|
|
agents = await agent_registry.get_idle_agents()
|
|
return {"agents": [agent.model_dump() for agent in agents]}
|
|
|
|
|
|
@app.get("/agents/{agent_id}")
|
|
async def get_agent(agent_id: str):
|
|
"""Get specific agent details."""
|
|
agent = await agent_registry.get_agent(agent_id)
|
|
if not agent:
|
|
raise HTTPException(status_code=404, detail="Agent not found")
|
|
return agent.model_dump()
|
|
|
|
|
|
@app.post("/tasks")
|
|
async def create_task(request: TaskCreateRequest):
|
|
"""Create a new task."""
|
|
task = await task_queue.create_task(
|
|
description=request.description,
|
|
context=request.context,
|
|
max_retries=request.max_retries
|
|
)
|
|
TASKS_CREATED.inc()
|
|
return task_payload(task)
|
|
|
|
|
|
@app.post("/tasks/assign")
|
|
async def assign_task(request: TaskAssignRequest):
|
|
"""Manually assign a task to an agent."""
|
|
task = await task_queue.get_task(request.task_id)
|
|
if not task:
|
|
raise HTTPException(status_code=404, detail="Task not found")
|
|
|
|
success = await task_queue.assign_task(request.task_id, request.agent_id)
|
|
if not success:
|
|
raise HTTPException(status_code=400, detail="Failed to assign task")
|
|
|
|
await task_queue.remove_pending_task(request.task_id)
|
|
run = await swarm_runtime.get_run_for_task(request.task_id)
|
|
if run:
|
|
await swarm_runtime.emit_event(
|
|
run,
|
|
"task.claimed",
|
|
task_id=request.task_id,
|
|
agent_instance_id=request.agent_id,
|
|
payload={
|
|
"task_id": request.task_id,
|
|
"agent_role": task.agent_role,
|
|
"agent_id": request.agent_id,
|
|
},
|
|
)
|
|
|
|
# Notify agent via WebSocket
|
|
await manager.send_message(request.agent_id, {
|
|
"type": "task_assignment",
|
|
"task_id": request.task_id,
|
|
"description": task.description,
|
|
"context": task.context,
|
|
})
|
|
|
|
return {"status": "assigned", "task_id": request.task_id, "agent_id": request.agent_id}
|
|
|
|
|
|
@app.get("/tasks")
|
|
async def list_tasks(status: Optional[str] = None):
|
|
"""List all tasks, optionally filtered by status."""
|
|
task_status = TaskStatus(status) if status else None
|
|
tasks = await task_queue.get_all_tasks(status=task_status)
|
|
return {"tasks": [task_payload(task) for task in tasks]}
|
|
|
|
|
|
@app.get("/tasks/{task_id}")
|
|
async def get_task(task_id: str):
|
|
"""Get specific task details."""
|
|
task = await task_queue.get_task(task_id)
|
|
if not task:
|
|
raise HTTPException(status_code=404, detail="Task not found")
|
|
return task_payload(task)
|
|
|
|
|
|
@app.get("/handoffs")
|
|
async def list_handoffs(agent_id: Optional[str] = None, limit: int = 100):
|
|
"""List handoff history."""
|
|
handoffs = await handoff_manager.get_handoff_history(agent_id=agent_id, limit=limit)
|
|
return {"handoffs": [h.model_dump() for h in handoffs]}
|
|
|
|
|
|
@app.get("/metrics")
|
|
async def metrics():
|
|
"""Prometheus metrics endpoint."""
|
|
# Update gauge metrics before returning
|
|
agents = await agent_registry.get_all_agents()
|
|
AGENTS_ACTIVE.set(len([a for a in agents if a.status == AgentStatus.BUSY]))
|
|
WEBSOCKET_CONNECTIONS.set(len(manager.active_connections))
|
|
|
|
# Update per-agent status
|
|
for agent in agents:
|
|
AGENT_STATUS.labels(agent_id=agent.agent_id, status=agent.status.value).set(1)
|
|
|
|
return PlainTextResponse(generate_latest(), media_type=CONTENT_TYPE_LATEST)
|
|
|
|
|
|
# WebSocket endpoint
|
|
@app.websocket("/ws/{agent_id}")
|
|
async def websocket_endpoint(websocket: WebSocket, agent_id: str):
|
|
"""WebSocket endpoint for agent communication."""
|
|
await manager.connect(agent_id, websocket)
|
|
|
|
try:
|
|
# Wait for registration message
|
|
data = await websocket.receive_json()
|
|
|
|
if data.get("type") == "register":
|
|
capabilities = data.get("capabilities", [])
|
|
await agent_registry.register_agent(agent_id, capabilities)
|
|
record_agent_slots(agent_id, data)
|
|
|
|
await websocket.send_json({
|
|
"type": "registered",
|
|
"agent_id": agent_id,
|
|
"status": "success"
|
|
})
|
|
|
|
# Main message loop
|
|
while True:
|
|
try:
|
|
message = await asyncio.wait_for(
|
|
websocket.receive_json(),
|
|
timeout=15.0 # Expect heartbeat every 15s
|
|
)
|
|
|
|
message_type = message.get("type")
|
|
|
|
if message_type == "heartbeat":
|
|
record_agent_slots(agent_id, message)
|
|
await agent_registry.update_heartbeat(agent_id)
|
|
current_agent = await agent_registry.get_agent(agent_id)
|
|
if current_agent and current_agent.current_task_id:
|
|
current_task = await task_queue.get_task(current_agent.current_task_id)
|
|
run = await swarm_runtime.get_run_for_task(current_agent.current_task_id)
|
|
if run and current_task:
|
|
await swarm_runtime.emit_event(
|
|
run,
|
|
"task.heartbeat",
|
|
task_id=current_agent.current_task_id,
|
|
agent_instance_id=agent_id,
|
|
payload={
|
|
"task_id": current_agent.current_task_id,
|
|
"agent_role": current_task.agent_role,
|
|
"agent_id": agent_id,
|
|
"status": current_task.status.value,
|
|
},
|
|
)
|
|
await maybe_emit_budget_alert(run)
|
|
await websocket.send_json({"type": "heartbeat_ack"})
|
|
|
|
elif message_type == "handoff":
|
|
# Handle handoff request
|
|
target_agent_id = message.get("target_agent_id")
|
|
task_context = message.get("task_context", {})
|
|
|
|
handoff_id = await handoff_manager.initiate_handoff(
|
|
agent_id, target_agent_id, task_context
|
|
)
|
|
|
|
if handoff_id:
|
|
# Notify target agent
|
|
await manager.send_message(target_agent_id, {
|
|
"type": "handoff_request",
|
|
"handoff_id": handoff_id,
|
|
"source_agent_id": agent_id,
|
|
"task_context": task_context
|
|
})
|
|
|
|
await websocket.send_json({
|
|
"type": "handoff_initiated",
|
|
"handoff_id": handoff_id,
|
|
"status": "success"
|
|
})
|
|
else:
|
|
await websocket.send_json({
|
|
"type": "handoff_failed",
|
|
"reason": "Target agent not available"
|
|
})
|
|
|
|
elif message_type == "handoff_request":
|
|
task_id = message.get("task_id")
|
|
subtask = message.get("subtask", {})
|
|
target_capabilities = message.get("target_capabilities", [])
|
|
task = await task_queue.get_task(task_id)
|
|
run = await swarm_runtime.get_run_for_task(task_id)
|
|
child_task = None
|
|
if (
|
|
run and task and handoff_workflow_enabled(task)
|
|
and task.context.get("allow_handoff", False)
|
|
and task.parent_task_id is None
|
|
):
|
|
child_context = {
|
|
**task.context,
|
|
"task_title": subtask.get("description", task.title or task.description[:80]),
|
|
"depends_on": [],
|
|
"agent_role": ",".join(target_capabilities) or "specialist",
|
|
"required_capabilities": target_capabilities,
|
|
"parent_task_id": task.task_id,
|
|
"root_task_id": task.root_task_id or task.task_id,
|
|
"source": "dynamic_handoff",
|
|
"workflow_mode": "multi_agent",
|
|
"allow_handoff": False,
|
|
}
|
|
child_task = await task_queue.create_task(
|
|
title=subtask.get("description", task.title or task.description[:80]),
|
|
description=subtask.get("description", task.description),
|
|
context=child_context,
|
|
max_retries=task.max_retries,
|
|
agent_role=",".join(target_capabilities) or "specialist",
|
|
required_capabilities=target_capabilities,
|
|
depends_on=[],
|
|
parent_task_id=task.task_id,
|
|
root_task_id=task.root_task_id or task.task_id,
|
|
source="dynamic_handoff",
|
|
)
|
|
await task_queue.add_child_task(task.task_id, child_task.task_id)
|
|
await swarm_runtime.attach_task(run, child_task.task_id)
|
|
await swarm_runtime.emit_event(
|
|
run,
|
|
"task.created",
|
|
task_id=child_task.task_id,
|
|
agent_instance_id=agent_id,
|
|
payload={
|
|
"task_id": child_task.task_id,
|
|
"title": child_task.title,
|
|
"description": child_task.description,
|
|
"agent_role": child_task.agent_role,
|
|
"status": child_task.status.value,
|
|
"depends_on": child_task.depends_on,
|
|
"required_capabilities": child_task.required_capabilities,
|
|
"parent_task_id": child_task.parent_task_id,
|
|
"root_task_id": child_task.root_task_id,
|
|
"source": child_task.source,
|
|
"attempt": child_task.retry_count,
|
|
},
|
|
)
|
|
if run and task:
|
|
from_role = task.agent_role
|
|
to_role = ",".join(target_capabilities) or "specialist"
|
|
await swarm_runtime.emit_event(
|
|
run,
|
|
"handoff.requested",
|
|
task_id=task_id,
|
|
agent_instance_id=agent_id,
|
|
payload={
|
|
"task_id": task_id,
|
|
"from_role": from_role,
|
|
"to_role": to_role,
|
|
"summary": subtask.get("description", "Subtask handoff requested"),
|
|
"target_capabilities": target_capabilities,
|
|
"child_task_id": child_task.task_id if child_task else None,
|
|
"parent_task_id": task.task_id,
|
|
},
|
|
)
|
|
await swarm_runtime.emit_event(
|
|
run,
|
|
"timeline.updated",
|
|
task_id=task_id,
|
|
agent_instance_id=agent_id,
|
|
payload={
|
|
"summary": subtask.get("description", "Subtask handoff requested"),
|
|
"task_id": task_id,
|
|
"child_task_id": child_task.task_id if child_task else None,
|
|
"status": "blocked" if child_task else "recorded",
|
|
},
|
|
)
|
|
await websocket.send_json({
|
|
"type": "handoff_response",
|
|
"task_id": task_id,
|
|
"status": "recorded",
|
|
"child_task_id": child_task.task_id if child_task else None,
|
|
})
|
|
|
|
elif message_type == "handoff_accept":
|
|
# Target agent accepts handoff
|
|
handoff_id = message.get("handoff_id")
|
|
handoff = await handoff_manager.get_handoff(handoff_id)
|
|
await handoff_manager.complete_handoff(handoff_id)
|
|
if handoff:
|
|
task_id = handoff.task_context.get("task_id")
|
|
run = await swarm_runtime.get_run_for_task(task_id) if task_id else None
|
|
if run:
|
|
await swarm_runtime.emit_event(
|
|
run,
|
|
"handoff.completed",
|
|
task_id=task_id,
|
|
agent_instance_id=agent_id,
|
|
payload={
|
|
"handoff_id": handoff_id,
|
|
"source_agent_id": handoff.source_agent_id,
|
|
"target_agent_id": handoff.target_agent_id,
|
|
"task_id": task_id,
|
|
# HM requires from_role/to_role; this legacy accept path
|
|
# lacks role context, so fall back to agent ids.
|
|
"from_role": handoff.source_agent_id,
|
|
"to_role": handoff.target_agent_id,
|
|
},
|
|
)
|
|
|
|
await websocket.send_json({
|
|
"type": "handoff_accepted",
|
|
"handoff_id": handoff_id
|
|
})
|
|
|
|
elif message_type == "handoff_reject":
|
|
# Target agent rejects handoff
|
|
handoff_id = message.get("handoff_id")
|
|
reason = message.get("reason", "Rejected by target agent")
|
|
await handoff_manager.fail_handoff(handoff_id, reason)
|
|
|
|
await websocket.send_json({
|
|
"type": "handoff_rejected",
|
|
"handoff_id": handoff_id
|
|
})
|
|
|
|
elif message_type == "task_start":
|
|
# Agent starts working on task
|
|
task_id = message.get("task_id")
|
|
await task_queue.start_task(task_id)
|
|
task = await task_queue.get_task(task_id)
|
|
run = await swarm_runtime.get_run_for_task(task_id)
|
|
if run and task:
|
|
await swarm_runtime.emit_event(
|
|
run,
|
|
"task.running",
|
|
task_id=task_id,
|
|
agent_instance_id=agent_id,
|
|
payload={
|
|
"task_id": task_id,
|
|
"agent_role": task.agent_role,
|
|
"agent_id": agent_id,
|
|
},
|
|
)
|
|
|
|
elif message_type == "blocked_on_handoff":
|
|
task_id = message.get("task_id")
|
|
child_task_id = message.get("child_task_id")
|
|
reason = message.get("reason") or "Waiting on delegated child task"
|
|
task = await task_queue.block_task(task_id, reason=reason, release_agent=True)
|
|
run = await swarm_runtime.get_run_for_task(task_id)
|
|
if run and task:
|
|
await swarm_runtime.emit_event(
|
|
run,
|
|
"task.blocked",
|
|
task_id=task_id,
|
|
agent_instance_id=agent_id,
|
|
payload={
|
|
"task_id": task_id,
|
|
"agent_role": task.agent_role,
|
|
"agent_id": agent_id,
|
|
"child_task_id": child_task_id,
|
|
"reason": reason,
|
|
},
|
|
)
|
|
await swarm_runtime.emit_event(
|
|
run,
|
|
"timeline.updated",
|
|
task_id=task_id,
|
|
agent_instance_id=agent_id,
|
|
payload={
|
|
"summary": reason,
|
|
"task_id": task_id,
|
|
"child_task_id": child_task_id,
|
|
"status": "blocked",
|
|
},
|
|
)
|
|
await refresh_swarm_run_status(run)
|
|
await websocket.send_json({
|
|
"type": "task_blocked_ack",
|
|
"task_id": task_id,
|
|
})
|
|
|
|
elif message_type == "task_complete":
|
|
# Agent completes task
|
|
task_id = message.get("task_id")
|
|
result = message.get("result")
|
|
result_text = json.dumps(result) if isinstance(result, dict) else result
|
|
logger.info(f"Received task_complete for {task_id}, result length: {len(result_text) if result_text else 0}")
|
|
completed = await task_queue.complete_task(task_id, result_text)
|
|
task = await task_queue.get_task(task_id)
|
|
run = await swarm_runtime.get_run_for_task(task_id)
|
|
if completed and task:
|
|
await deposit_pheromone(task, agent_id, success=True, result=result, run=run)
|
|
if completed and run and task:
|
|
await emit_task_completion_events(run, task, agent_id, result)
|
|
await emit_usage_event(run, task, agent_id, result)
|
|
await finalize_parent_after_child(
|
|
run,
|
|
task,
|
|
agent_id=agent_id,
|
|
success=True,
|
|
summary="Delegated child task completed",
|
|
)
|
|
await refresh_swarm_run_status(run)
|
|
|
|
await websocket.send_json({
|
|
"type": "task_completed",
|
|
"task_id": task_id
|
|
})
|
|
|
|
elif message_type == "task_result":
|
|
# Backward-compatible handler for older/simple agents
|
|
task_id = message.get("task_id")
|
|
success = message.get("success")
|
|
if success is None:
|
|
success = message.get("status") == "completed"
|
|
|
|
if success:
|
|
result = message.get("result")
|
|
result_text = json.dumps(result) if isinstance(result, dict) else result
|
|
completed = await task_queue.complete_task(task_id, result_text)
|
|
task = await task_queue.get_task(task_id)
|
|
run = await swarm_runtime.get_run_for_task(task_id)
|
|
if completed and task:
|
|
await deposit_pheromone(task, agent_id, success=True, result=result, run=run)
|
|
if completed and run and task:
|
|
await emit_task_completion_events(run, task, agent_id, result)
|
|
await emit_usage_event(run, task, agent_id, result)
|
|
await finalize_parent_after_child(
|
|
run,
|
|
task,
|
|
agent_id=agent_id,
|
|
success=True,
|
|
summary="Delegated child task completed",
|
|
)
|
|
await refresh_swarm_run_status(run)
|
|
await websocket.send_json({
|
|
"type": "task_completed",
|
|
"task_id": task_id
|
|
})
|
|
else:
|
|
reason = message.get("error") or message.get("reason") or "Unknown error"
|
|
failed = await task_queue.fail_task(task_id, reason)
|
|
task = await task_queue.get_task(task_id)
|
|
run = await swarm_runtime.get_run_for_task(task_id)
|
|
if failed and task:
|
|
await deposit_pheromone(task, agent_id, success=False, run=run)
|
|
if failed and run and task:
|
|
event_type = "task.retried" if task.status == TaskStatus.PENDING else "task.failed"
|
|
await swarm_runtime.emit_event(
|
|
run,
|
|
event_type,
|
|
task_id=task_id,
|
|
agent_instance_id=agent_id,
|
|
payload={
|
|
"task_id": task_id,
|
|
"agent_role": task.agent_role,
|
|
"reason": reason,
|
|
"summary": reason,
|
|
"attempt": task.retry_count,
|
|
"parent_task_id": task.parent_task_id,
|
|
"root_task_id": task.root_task_id,
|
|
},
|
|
)
|
|
if task.status == TaskStatus.FAILED:
|
|
await finalize_parent_after_child(
|
|
run,
|
|
task,
|
|
agent_id=agent_id,
|
|
success=False,
|
|
summary=reason,
|
|
)
|
|
await refresh_swarm_run_status(run)
|
|
await websocket.send_json({
|
|
"type": "task_failed_ack",
|
|
"task_id": task_id
|
|
})
|
|
|
|
elif message_type == "task_failed":
|
|
# Agent reports task failure
|
|
task_id = message.get("task_id")
|
|
reason = message.get("reason", "Unknown error")
|
|
failed = await task_queue.fail_task(task_id, reason)
|
|
task = await task_queue.get_task(task_id)
|
|
run = await swarm_runtime.get_run_for_task(task_id)
|
|
if failed and task:
|
|
await deposit_pheromone(task, agent_id, success=False, run=run)
|
|
if failed and run and task:
|
|
event_type = "task.retried" if task.status == TaskStatus.PENDING else "task.failed"
|
|
await swarm_runtime.emit_event(
|
|
run,
|
|
event_type,
|
|
task_id=task_id,
|
|
agent_instance_id=agent_id,
|
|
payload={
|
|
"task_id": task_id,
|
|
"agent_role": task.agent_role,
|
|
"reason": reason,
|
|
"summary": reason,
|
|
"attempt": task.retry_count,
|
|
"parent_task_id": task.parent_task_id,
|
|
"root_task_id": task.root_task_id,
|
|
},
|
|
)
|
|
if task.status == TaskStatus.FAILED:
|
|
await finalize_parent_after_child(
|
|
run,
|
|
task,
|
|
agent_id=agent_id,
|
|
success=False,
|
|
summary=reason,
|
|
)
|
|
await refresh_swarm_run_status(run)
|
|
|
|
await websocket.send_json({
|
|
"type": "task_failed_ack",
|
|
"task_id": task_id
|
|
})
|
|
|
|
elif message_type == "status_update":
|
|
# Agent updates its status
|
|
raw_status = message.get("status")
|
|
if raw_status == "handoff":
|
|
raw_status = AgentStatus.HANDOFF_PENDING.value
|
|
status = AgentStatus(raw_status)
|
|
task_id = message.get("task_id")
|
|
await agent_registry.update_status(agent_id, status, task_id)
|
|
|
|
elif message_type == "task_accepted":
|
|
# Merged agent confirms it took (or de-duplicated) an assignment.
|
|
record_agent_slots(agent_id, message)
|
|
if message.get("status") == "duplicate":
|
|
logger.info(
|
|
f"Agent {agent_id} reported duplicate task {message.get('task_id')}"
|
|
)
|
|
else:
|
|
logger.debug(
|
|
f"Agent {agent_id} accepted task {message.get('task_id')}"
|
|
)
|
|
|
|
elif message_type == "task_rejected":
|
|
# Agent was at capacity; return the task to the pending queue for another agent.
|
|
record_agent_slots(agent_id, message)
|
|
rejected_task_id = message.get("task_id")
|
|
reject_reason = message.get("reason", "rejected")
|
|
logger.info(
|
|
f"Agent {agent_id} rejected task {rejected_task_id} ({reject_reason}); requeuing"
|
|
)
|
|
released = await task_queue.release_task(rejected_task_id, agent_id=agent_id)
|
|
if released:
|
|
released_task = await task_queue.get_task(rejected_task_id)
|
|
released_run = await swarm_runtime.get_run_for_task(rejected_task_id)
|
|
if released_task and released_run:
|
|
await swarm_runtime.emit_event(
|
|
released_run,
|
|
"task.released",
|
|
task_id=rejected_task_id,
|
|
agent_instance_id=agent_id,
|
|
payload={
|
|
"task_id": rejected_task_id,
|
|
"agent_role": released_task.agent_role,
|
|
"reason": reject_reason,
|
|
},
|
|
)
|
|
|
|
elif message_type == "peer_message":
|
|
# Route peer collaboration messages between agents; stamp the sender so the
|
|
# recipient knows where to send its reply.
|
|
target_agent_id = message.get("target_agent_id")
|
|
delivered = False
|
|
if target_agent_id:
|
|
delivered = await manager.send_message(
|
|
target_agent_id,
|
|
{**message, "from_agent_id": agent_id},
|
|
)
|
|
else:
|
|
logger.warning(f"peer_message from {agent_id} missing target_agent_id")
|
|
# Record internal communication telemetry (feeds benchmark S_communication).
|
|
# Not emitted to the Manager event stream — it is not a registered HM event.
|
|
peer_task_id = message.get("task_id")
|
|
if peer_task_id:
|
|
peer_run = await swarm_runtime.get_run_for_task(peer_task_id)
|
|
if peer_run:
|
|
await swarm_runtime.record_peer_message(
|
|
peer_run,
|
|
correlation_id=message.get("correlation_id"),
|
|
is_reply=bool(message.get("is_reply")),
|
|
delivered=delivered,
|
|
)
|
|
|
|
else:
|
|
logger.warning(f"Unknown message type from {agent_id}: {message_type}")
|
|
|
|
except asyncio.TimeoutError:
|
|
# No message received within timeout - check heartbeat
|
|
agent = await agent_registry.get_agent(agent_id)
|
|
if agent:
|
|
import time
|
|
if time.time() - agent.last_heartbeat > 30:
|
|
logger.warning(f"Agent {agent_id} heartbeat timeout")
|
|
break
|
|
continue
|
|
|
|
else:
|
|
logger.error(f"Agent {agent_id} did not send registration message")
|
|
await websocket.close(code=1008, reason="Registration required")
|
|
|
|
except WebSocketDisconnect:
|
|
logger.info(f"Agent {agent_id} disconnected")
|
|
except Exception as e:
|
|
logger.error(f"Error in WebSocket handler for agent {agent_id}: {e}")
|
|
finally:
|
|
# Cleanup
|
|
manager.disconnect(agent_id)
|
|
AGENT_SLOTS.pop(agent_id, None)
|
|
await agent_registry.deregister_agent(agent_id)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import uvicorn
|
|
uvicorn.run(app, host="0.0.0.0", port=8000)
|