784 lines
30 KiB
Python
784 lines
30 KiB
Python
"""Sub-mode runtime orchestrator."""
|
|
|
|
import asyncio
|
|
import json
|
|
import logging
|
|
import uuid
|
|
from datetime import datetime, timedelta
|
|
from typing import Dict, List, Any, Optional
|
|
from sqlalchemy.orm import Session
|
|
|
|
from database import Swarm, SwarmAgent, SwarmMessage, SwarmStatus, SwarmAgentStatus
|
|
from .agent_client import SwarmAgentClient
|
|
from .callback_client import CallbackDeliveryClient
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
PHASE_MAP = {
|
|
"planning": ("requirements", "agent_running"),
|
|
"coding": ("backend", "agent_running"),
|
|
"reviewing": ("review", "ready_for_test"),
|
|
"executing": ("backend", "agent_running"),
|
|
"parallel_execution": ("backend", "agent_running"),
|
|
"completed": ("deploy", "completed"),
|
|
"failed": ("review", "failed"),
|
|
}
|
|
|
|
|
|
class SwarmOrchestrator:
|
|
"""Runtime orchestrator for Heicode sub-mode execution."""
|
|
|
|
def __init__(self, swarm_id: str, db: Session):
|
|
"""
|
|
Initialize orchestrator.
|
|
|
|
Args:
|
|
swarm_id: Swarm ID
|
|
db: Database session
|
|
"""
|
|
self.swarm_id = swarm_id
|
|
self.db = db
|
|
self.agents: Dict[str, SwarmAgentClient] = {}
|
|
self.swarm: Optional[Swarm] = None
|
|
self.callback: Optional[CallbackDeliveryClient] = None
|
|
self.correlation_id: Optional[str] = None
|
|
|
|
async def initialize(self) -> bool:
|
|
"""
|
|
Initialize swarm - create all agent pods.
|
|
|
|
Returns:
|
|
True if successful, False otherwise
|
|
"""
|
|
try:
|
|
# Load swarm from database
|
|
self.swarm = self.db.query(Swarm).filter(
|
|
Swarm.swarm_id == self.swarm_id
|
|
).first()
|
|
|
|
if not self.swarm:
|
|
logger.error(f"Swarm {self.swarm_id} not found")
|
|
return False
|
|
|
|
project_context = self.swarm.project_context or {}
|
|
self.callback = CallbackDeliveryClient(project_context.get("_callback"))
|
|
self.correlation_id = project_context.get("correlation_id") or f"swarm_{self.swarm_id}"
|
|
|
|
# Update status to initializing
|
|
self.swarm.status = SwarmStatus.INITIALIZING
|
|
self.db.commit()
|
|
await self._emit_status("initializing")
|
|
|
|
logger.info(f"Initializing swarm {self.swarm_id}")
|
|
|
|
# Get all agents for this swarm
|
|
swarm_agents = self.db.query(SwarmAgent).filter(
|
|
SwarmAgent.swarm_id == self.swarm_id
|
|
).all()
|
|
|
|
# Wait for all agents to be ready
|
|
for agent in swarm_agents:
|
|
if agent.service_url:
|
|
client = SwarmAgentClient(agent.agent_id, agent.service_url)
|
|
self.agents[agent.agent_id] = client
|
|
|
|
# Update agent status to running
|
|
agent.status = SwarmAgentStatus.RUNNING
|
|
self.db.commit()
|
|
await self._emit_agent_event(
|
|
"agent.started",
|
|
agent.agent_id,
|
|
{
|
|
"agent_role": agent.role,
|
|
"status": "running",
|
|
"service_url": agent.service_url,
|
|
},
|
|
)
|
|
|
|
# Update swarm status to running
|
|
self.swarm.status = SwarmStatus.RUNNING
|
|
self.swarm.phase = "planning"
|
|
self.db.commit()
|
|
await self._emit_status("running")
|
|
await self._emit_phase("planning", "Swarm initialized and planning started")
|
|
await self._emit_usage()
|
|
|
|
logger.info(f"Swarm {self.swarm_id} initialized with {len(self.agents)} agents")
|
|
return True
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error initializing swarm {self.swarm_id}: {e}")
|
|
if self.swarm:
|
|
self.swarm.status = SwarmStatus.FAILED
|
|
self.swarm.error_message = str(e)
|
|
self.db.commit()
|
|
return False
|
|
|
|
async def execute(self) -> Dict[str, Any]:
|
|
"""
|
|
Execute swarm task.
|
|
|
|
Returns:
|
|
Execution results
|
|
"""
|
|
try:
|
|
if not self.swarm:
|
|
raise ValueError("Swarm not initialized")
|
|
|
|
strategy = self.swarm.orchestration_strategy
|
|
|
|
if strategy == "sequential":
|
|
result = await self._execute_sequential()
|
|
elif strategy == "parallel":
|
|
result = await self._execute_parallel()
|
|
else:
|
|
result = await self._execute_hybrid()
|
|
|
|
# Update swarm status
|
|
self.swarm.status = SwarmStatus.COMPLETED
|
|
self.swarm.completed_at = datetime.utcnow()
|
|
self.swarm.progress = 100
|
|
self.swarm.artifacts = result.get("artifacts", [])
|
|
self._finalize_running_agents(SwarmAgentStatus.COMPLETED)
|
|
self.db.commit()
|
|
await self._emit_artifacts(result.get("artifacts", []))
|
|
await self._emit_phase("completed", "Swarm execution completed")
|
|
await self._emit_status("completed")
|
|
await self._emit_usage()
|
|
|
|
return result
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error executing swarm {self.swarm_id}: {e}")
|
|
if self.swarm:
|
|
self.swarm.status = SwarmStatus.FAILED
|
|
self.swarm.error_message = str(e)
|
|
self._finalize_running_agents(SwarmAgentStatus.FAILED, str(e))
|
|
self.db.commit()
|
|
await self._emit_status("failed", {"error": str(e)})
|
|
raise
|
|
|
|
async def _execute_sequential(self) -> Dict[str, Any]:
|
|
"""
|
|
Execute sequential strategy: architect → coder → reviewer.
|
|
|
|
Returns:
|
|
Execution results
|
|
"""
|
|
logger.info(f"Executing sequential strategy for swarm {self.swarm_id}")
|
|
|
|
results = {
|
|
"artifacts": [],
|
|
"phases": []
|
|
}
|
|
|
|
# Phase 1: Architect designs
|
|
architect = self._get_agent_by_role("architect")
|
|
if architect:
|
|
self.swarm.phase = "planning"
|
|
self.swarm.progress = 10
|
|
self.db.commit()
|
|
await self._emit_phase("planning", "Architecture planning started")
|
|
|
|
design = await self._send_task(
|
|
architect,
|
|
f"Design the architecture for: {self.swarm.task_description}"
|
|
)
|
|
results["phases"].append({"phase": "planning", "result": design})
|
|
architect_record = self._get_agent_record(architect.agent_id)
|
|
if architect_record:
|
|
results["artifacts"].append(self._build_artifact(architect_record, design, len(results["artifacts"])))
|
|
|
|
# Phase 2: Coders implement
|
|
coders = self._get_agents_by_role("coder")
|
|
if coders:
|
|
self.swarm.phase = "coding"
|
|
self.swarm.progress = 40
|
|
self.db.commit()
|
|
await self._emit_phase("coding", "Implementation started")
|
|
|
|
code_results = await asyncio.gather(*[
|
|
self._send_task(coder, f"Implement: {self.swarm.task_description}")
|
|
for coder in coders
|
|
])
|
|
results["phases"].append({"phase": "coding", "results": code_results})
|
|
for i, code in enumerate(code_results):
|
|
coder_record = self._get_agent_record(coders[i].agent_id)
|
|
if coder_record:
|
|
results["artifacts"].append(self._build_artifact(coder_record, code, len(results["artifacts"])))
|
|
|
|
# Phase 3: Reviewer reviews
|
|
reviewer = self._get_agent_by_role("reviewer")
|
|
if reviewer:
|
|
self.swarm.phase = "reviewing"
|
|
self.swarm.progress = 80
|
|
self.db.commit()
|
|
await self._emit_phase("reviewing", "Review started")
|
|
|
|
review = await self._send_task(
|
|
reviewer,
|
|
f"Review the implementation: {code_results if coders else 'No code generated'}"
|
|
)
|
|
results["phases"].append({"phase": "reviewing", "result": review})
|
|
reviewer_record = self._get_agent_record(reviewer.agent_id)
|
|
if reviewer_record:
|
|
results["artifacts"].append(self._build_artifact(reviewer_record, review, len(results["artifacts"])))
|
|
|
|
return results
|
|
|
|
async def _execute_parallel(self) -> Dict[str, Any]:
|
|
"""
|
|
Execute parallel strategy: all agents work simultaneously.
|
|
|
|
Returns:
|
|
Execution results
|
|
"""
|
|
logger.info(f"Executing parallel strategy for swarm {self.swarm_id}")
|
|
|
|
self.swarm.phase = "executing"
|
|
self.swarm.progress = 50
|
|
self.db.commit()
|
|
await self._emit_phase("executing", "Parallel execution started")
|
|
|
|
# Send task to all agents in parallel
|
|
executable_agents = list(self.agents.values())
|
|
tasks = [
|
|
self._send_task(client, self._build_agent_task(self._get_agent_record(client.agent_id)))
|
|
for client in executable_agents
|
|
]
|
|
|
|
results_list = await asyncio.gather(*tasks, return_exceptions=True)
|
|
|
|
results = {
|
|
"artifacts": [],
|
|
"phases": [{"phase": "parallel_execution", "results": results_list}]
|
|
}
|
|
|
|
for i, result in enumerate(results_list):
|
|
if not isinstance(result, Exception):
|
|
agent_record = self._get_agent_record(executable_agents[i].agent_id)
|
|
if agent_record:
|
|
results["artifacts"].append(
|
|
self._build_artifact(agent_record, result, i)
|
|
)
|
|
|
|
return results
|
|
|
|
async def _execute_hybrid(self) -> Dict[str, Any]:
|
|
"""
|
|
Execute hybrid strategy: combination of sequential and parallel.
|
|
|
|
Returns:
|
|
Execution results
|
|
"""
|
|
logger.info(f"Executing hybrid strategy for swarm {self.swarm_id}")
|
|
roles = {agent.role for agent in self.db.query(SwarmAgent).filter(SwarmAgent.swarm_id == self.swarm_id).all()}
|
|
project_context = self._callback_context()
|
|
if project_context.get("sub_mode") or not roles.intersection({"architect", "coder", "reviewer"}):
|
|
return await self._execute_sub_mode_agents()
|
|
return await self._execute_sequential()
|
|
|
|
async def _execute_sub_mode_agents(self) -> Dict[str, Any]:
|
|
"""Execute ordinary sub-mode agents using their actual configured roles."""
|
|
logger.info(f"Executing ordinary sub-mode workflow for swarm {self.swarm_id}")
|
|
|
|
self.swarm.phase = "development"
|
|
self.swarm.progress = 30
|
|
self.db.commit()
|
|
await self._emit_phase("executing", "Ordinary sub-mode execution started")
|
|
|
|
executable_agent_records = [
|
|
agent
|
|
for agent in self.db.query(SwarmAgent).filter(SwarmAgent.swarm_id == self.swarm_id).all()
|
|
if agent.agent_id in self.agents
|
|
]
|
|
tasks = [
|
|
self._send_task(self.agents[agent.agent_id], self._build_agent_task(agent))
|
|
for agent in executable_agent_records
|
|
]
|
|
results_list = await asyncio.gather(*tasks, return_exceptions=True)
|
|
|
|
artifacts: List[Dict[str, Any]] = []
|
|
phases: List[Dict[str, Any]] = []
|
|
for index, result in enumerate(results_list):
|
|
agent_record = executable_agent_records[index]
|
|
if isinstance(result, Exception):
|
|
phases.append(
|
|
{
|
|
"phase": "development",
|
|
"agent_id": agent_record.agent_id,
|
|
"role": agent_record.role,
|
|
"status": "failed",
|
|
"error": str(result),
|
|
}
|
|
)
|
|
continue
|
|
|
|
phases.append(
|
|
{
|
|
"phase": "development",
|
|
"agent_id": agent_record.agent_id,
|
|
"role": agent_record.role,
|
|
"status": "completed",
|
|
"summary": self._response_summary(result),
|
|
}
|
|
)
|
|
artifacts.append(self._build_artifact(agent_record, result, index))
|
|
|
|
if not artifacts:
|
|
artifacts.append(
|
|
{
|
|
"artifact_id": f"art_{self.swarm_id}_summary",
|
|
"artifact_type": "document",
|
|
"title": "Runtime execution summary",
|
|
"summary": self.swarm.error_message or "Runtime completed without per-agent artifacts; review swarm logs for details.",
|
|
"uri": f"runtime://{self.swarm_id}/artifacts/summary",
|
|
"agent_instance_id": None,
|
|
"stage": "development",
|
|
"checkpoint": "artifact_ready",
|
|
"metadata": {
|
|
"redacted": True,
|
|
"source": "agent-manager-swarm",
|
|
"runtime_deployment_id": self.swarm_id,
|
|
},
|
|
}
|
|
)
|
|
|
|
self.swarm.progress = 85
|
|
self.db.commit()
|
|
return {"artifacts": artifacts, "phases": phases}
|
|
|
|
async def _send_task(self, client: SwarmAgentClient, task: str) -> Dict[str, Any]:
|
|
"""
|
|
Send task to agent via A2A protocol.
|
|
|
|
Args:
|
|
client: Agent client
|
|
task: Task description
|
|
|
|
Returns:
|
|
Agent response
|
|
"""
|
|
try:
|
|
agent_record = self._get_agent_record(client.agent_id)
|
|
if agent_record:
|
|
agent_record.current_task = task
|
|
agent_record.status = SwarmAgentStatus.RUNNING
|
|
self.db.commit()
|
|
|
|
# Record message to database
|
|
message = SwarmMessage(
|
|
message_id=str(uuid.uuid4()),
|
|
swarm_id=self.swarm_id,
|
|
from_agent_id=None, # From orchestrator
|
|
to_agent_id=client.agent_id,
|
|
message_type="task",
|
|
content=task,
|
|
message_metadata={}
|
|
)
|
|
self.db.add(message)
|
|
self.db.commit()
|
|
|
|
# Update message count
|
|
self.swarm.total_messages += 1
|
|
self.db.commit()
|
|
await self._emit_tool_event(
|
|
"sk_tool.called",
|
|
client.agent_id,
|
|
{
|
|
"tool_name": "agent_task",
|
|
"tool_invocation_id": message.message_id,
|
|
"summary": "Dispatching task to agent",
|
|
"arguments_redacted": True,
|
|
},
|
|
)
|
|
|
|
# Send message to agent
|
|
response = await client.send_message({"text": task})
|
|
usage = self._extract_usage(response)
|
|
if self.swarm and usage["total_tokens"]:
|
|
self.swarm.tokens_used += usage["total_tokens"]
|
|
self.db.commit()
|
|
await self._emit_tool_event(
|
|
"sk_tool.completed",
|
|
client.agent_id,
|
|
{
|
|
"tool_name": "agent_task",
|
|
"tool_invocation_id": message.message_id,
|
|
"summary": "Agent task completed",
|
|
"result_preview": str(response)[:500],
|
|
"model_usage": usage,
|
|
},
|
|
)
|
|
|
|
# Record response
|
|
response_message = SwarmMessage(
|
|
message_id=str(uuid.uuid4()),
|
|
swarm_id=self.swarm_id,
|
|
from_agent_id=client.agent_id,
|
|
to_agent_id=None, # To orchestrator
|
|
message_type="response",
|
|
content=str(response),
|
|
message_metadata={}
|
|
)
|
|
self.db.add(response_message)
|
|
self.db.commit()
|
|
|
|
self.swarm.total_messages += 1
|
|
if agent_record:
|
|
agent_record.status = SwarmAgentStatus.COMPLETED
|
|
agent_record.output = self._response_summary(response)
|
|
agent_record.current_task = None
|
|
self.db.commit()
|
|
await self._emit_agent_event(
|
|
"agent.completed",
|
|
client.agent_id,
|
|
{
|
|
"status": "completed",
|
|
"summary": "Agent task completed",
|
|
},
|
|
)
|
|
await self._emit(
|
|
"task.completed",
|
|
agent_instance_id=client.agent_id,
|
|
payload={
|
|
"task_id": message.message_id,
|
|
"agent_role": agent_record.role if agent_record else None,
|
|
"status": "completed",
|
|
"summary": self._response_summary(response),
|
|
"runtime_deployment_id": self.swarm_id,
|
|
},
|
|
)
|
|
|
|
return response
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error sending task to agent {client.agent_id}: {e}")
|
|
agent_record = self._get_agent_record(client.agent_id)
|
|
if agent_record:
|
|
agent_record.status = SwarmAgentStatus.FAILED
|
|
agent_record.output = str(e)
|
|
agent_record.current_task = None
|
|
self.db.commit()
|
|
await self._emit_tool_event(
|
|
"sk_tool.failed",
|
|
client.agent_id,
|
|
{
|
|
"tool_name": "agent_task",
|
|
"summary": "Agent task failed",
|
|
"error": str(e),
|
|
},
|
|
)
|
|
await self._emit(
|
|
"task.failed",
|
|
agent_instance_id=client.agent_id,
|
|
payload={
|
|
"task_id": message.message_id if "message" in locals() else f"task_{client.agent_id}",
|
|
"agent_role": agent_record.role if agent_record else None,
|
|
"status": "failed",
|
|
"summary": str(e),
|
|
"runtime_deployment_id": self.swarm_id,
|
|
},
|
|
)
|
|
raise
|
|
|
|
def _get_agent_by_role(self, role: str) -> Optional[SwarmAgentClient]:
|
|
"""Get first agent by role"""
|
|
agent = self.db.query(SwarmAgent).filter(
|
|
SwarmAgent.swarm_id == self.swarm_id,
|
|
SwarmAgent.role == role
|
|
).first()
|
|
|
|
if agent and agent.agent_id in self.agents:
|
|
return self.agents[agent.agent_id]
|
|
return None
|
|
|
|
def _get_agents_by_role(self, role: str) -> List[SwarmAgentClient]:
|
|
"""Get all agents by role"""
|
|
agents = self.db.query(SwarmAgent).filter(
|
|
SwarmAgent.swarm_id == self.swarm_id,
|
|
SwarmAgent.role == role
|
|
).all()
|
|
|
|
return [
|
|
self.agents[agent.agent_id]
|
|
for agent in agents
|
|
if agent.agent_id in self.agents
|
|
]
|
|
|
|
async def stop(self, reason: Optional[str] = None):
|
|
"""
|
|
Stop swarm execution.
|
|
|
|
Args:
|
|
reason: Reason for stopping
|
|
"""
|
|
try:
|
|
if self.swarm:
|
|
self.swarm.status = SwarmStatus.STOPPED
|
|
self.swarm.error_message = reason
|
|
self.db.commit()
|
|
await self._emit_status("stopped", {"reason": reason})
|
|
|
|
# Close all agent clients
|
|
for client in self.agents.values():
|
|
await client.close()
|
|
|
|
logger.info(f"Swarm {self.swarm_id} stopped: {reason}")
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error stopping swarm {self.swarm_id}: {e}")
|
|
raise
|
|
|
|
async def cleanup(self):
|
|
"""Cleanup resources"""
|
|
for client in self.agents.values():
|
|
await client.close()
|
|
self.agents.clear()
|
|
|
|
def _callback_context(self) -> Dict[str, Any]:
|
|
"""Return callback context stored on the swarm."""
|
|
if not self.swarm:
|
|
return {}
|
|
return self.swarm.project_context or {}
|
|
|
|
async def _emit_status(self, status: str, extra_payload: Optional[Dict[str, Any]] = None) -> None:
|
|
"""Emit deployment.status_changed callback."""
|
|
payload = {"status": status, **(extra_payload or {})}
|
|
await self._emit("deployment.status_changed", payload=payload)
|
|
|
|
async def _emit_phase(self, internal_phase: str, summary: str) -> None:
|
|
"""Emit phase.changed and timeline.updated callbacks."""
|
|
stage, checkpoint = PHASE_MAP.get(internal_phase, (internal_phase, "agent_running"))
|
|
payload = {
|
|
"stage": stage,
|
|
"phase": stage,
|
|
"checkpoint": checkpoint,
|
|
"progress_pct": self.swarm.progress if self.swarm else 0,
|
|
"summary": summary,
|
|
"internal_phase": internal_phase,
|
|
}
|
|
await self._emit("phase.changed", payload=payload)
|
|
await self._emit(
|
|
"timeline.updated",
|
|
payload={
|
|
"title": summary,
|
|
"summary": summary,
|
|
"stage": stage,
|
|
"checkpoint": checkpoint,
|
|
"progress_pct": self.swarm.progress if self.swarm else 0,
|
|
"severity": "success" if checkpoint == "completed" else "info",
|
|
"next_action": "continue" if checkpoint != "completed" else "stop",
|
|
},
|
|
)
|
|
|
|
async def _emit_agent_event(self, event_type: str, agent_id: str, payload: Dict[str, Any]) -> None:
|
|
"""Emit agent lifecycle callback."""
|
|
await self._emit(event_type, agent_instance_id=agent_id, payload=payload)
|
|
|
|
async def _emit_tool_event(self, event_type: str, agent_id: str, payload: Dict[str, Any]) -> None:
|
|
"""Emit SK/tool callback."""
|
|
await self._emit(event_type, agent_instance_id=agent_id, payload=payload)
|
|
|
|
async def _emit_artifacts(self, artifacts: List[Dict[str, Any]]) -> None:
|
|
"""Emit artifact.created callbacks for generated outputs."""
|
|
for index, artifact in enumerate(artifacts):
|
|
artifact_type = artifact.get("artifact_type") or artifact.get("type") or "other"
|
|
await self._emit(
|
|
"artifact.created",
|
|
agent_instance_id=artifact.get("agent_instance_id"),
|
|
payload={
|
|
"artifact_id": artifact.get("artifact_id") or f"art_{self.swarm_id}_{index}",
|
|
"artifact_type": artifact_type,
|
|
"title": artifact.get("title") or f"{artifact_type} artifact",
|
|
"summary": artifact.get("summary") or str(artifact.get("content", ""))[:300],
|
|
"uri": artifact.get("uri") or f"runtime://{self.swarm_id}/artifacts/{artifact.get('artifact_id') or index}",
|
|
"mime_type": artifact.get("mime_type") or "text/plain",
|
|
"size_bytes": artifact.get("size_bytes"),
|
|
"stage": self._callback_context().get("agile_context", {}).get("stage") or "development",
|
|
"checkpoint": artifact.get("checkpoint") or "artifact_ready",
|
|
"metadata": {
|
|
"redacted": True,
|
|
"source": "agent-manager-swarm",
|
|
"runtime_deployment_id": self.swarm_id,
|
|
**(artifact.get("metadata") or {}),
|
|
},
|
|
},
|
|
)
|
|
|
|
async def _emit_usage(self) -> None:
|
|
"""Emit a minimal usage/cost event for Manager attribution."""
|
|
context = self._callback_context()
|
|
budget = context.get("budget") or {}
|
|
billing_context = context.get("billing_context") or {}
|
|
await self._emit(
|
|
"budget.alert",
|
|
payload={
|
|
"model_id": billing_context.get("default_model_id") or "unknown",
|
|
"model_tokens": self.swarm.tokens_used if self.swarm else 0,
|
|
"prompt_tokens": 0,
|
|
"completion_tokens": 0,
|
|
"model_cost_usd": 0,
|
|
"runtime_seconds": 0,
|
|
"cpu_core_seconds": 0,
|
|
"memory_mb_seconds": 0,
|
|
"billing_source": billing_context.get("provider") or "newapi",
|
|
"consumed_usd": 0,
|
|
"max_cost_usd": budget.get("max_cost_usd"),
|
|
"threshold_pct": budget.get("alert_threshold_pct", 80),
|
|
"severity": "info",
|
|
"budget": {
|
|
"max_tokens": budget.get("max_tokens"),
|
|
"max_cost_usd": budget.get("max_cost_usd"),
|
|
"consumed_usd": 0,
|
|
"remaining_usd": budget.get("max_cost_usd"),
|
|
},
|
|
},
|
|
)
|
|
|
|
async def _emit(
|
|
self,
|
|
event_type: str,
|
|
*,
|
|
agent_instance_id: Optional[str] = None,
|
|
payload: Optional[Dict[str, Any]] = None,
|
|
) -> None:
|
|
"""Emit callback if delivery is configured."""
|
|
if not self.callback or not self.swarm:
|
|
return
|
|
callback_context = self._callback_context()
|
|
await self.callback.emit(
|
|
event_type,
|
|
callback_context.get("manager_deployment_id")
|
|
or callback_context.get("heicode_deployment_id")
|
|
or self.swarm.swarm_id,
|
|
swarm_id=self.swarm.swarm_id,
|
|
agent_instance_id=agent_instance_id,
|
|
correlation_id=self.correlation_id,
|
|
payload=payload,
|
|
)
|
|
|
|
def _get_agent_record(self, agent_id: str) -> Optional[SwarmAgent]:
|
|
"""Load a swarm agent record by runtime agent ID."""
|
|
return (
|
|
self.db.query(SwarmAgent)
|
|
.filter(SwarmAgent.swarm_id == self.swarm_id, SwarmAgent.agent_id == agent_id)
|
|
.first()
|
|
)
|
|
|
|
def _build_agent_task(self, agent: Optional[SwarmAgent]) -> str:
|
|
"""Create a role-aware task prompt for ordinary sub-mode agents."""
|
|
if not agent or not self.swarm:
|
|
return self.swarm.task_description if self.swarm else ""
|
|
project_context = self._callback_context()
|
|
prompt_parts = [
|
|
f"Task objective: {self.swarm.task_description}",
|
|
f"Your role: {agent.role}",
|
|
"Return a concrete deliverable summary suitable for Heicode Manager artifacts.",
|
|
]
|
|
if project_context.get("repo_url"):
|
|
prompt_parts.append(f"Repository: {project_context['repo_url']}")
|
|
if project_context.get("branch"):
|
|
prompt_parts.append(f"Branch: {project_context['branch']}")
|
|
if project_context.get("agile_context"):
|
|
prompt_parts.append(
|
|
"Agile context: "
|
|
+ json.dumps(project_context["agile_context"], ensure_ascii=False, sort_keys=True)
|
|
)
|
|
return "\n".join(prompt_parts)
|
|
|
|
def _response_summary(self, response: Any) -> str:
|
|
"""Return a readable response summary for logs, artifacts, and callbacks."""
|
|
text = self._find_text(response)
|
|
if text:
|
|
return text[:1000]
|
|
if isinstance(response, (dict, list)):
|
|
return json.dumps(response, ensure_ascii=False)[:1000]
|
|
return str(response)[:1000]
|
|
|
|
def _find_text(self, value: Any) -> Optional[str]:
|
|
"""Recursively extract the first meaningful text payload from agent responses."""
|
|
if value is None:
|
|
return None
|
|
if isinstance(value, str):
|
|
cleaned = value.strip()
|
|
return cleaned or None
|
|
if isinstance(value, dict):
|
|
for key in ("text", "content", "message", "output", "result"):
|
|
if key in value:
|
|
nested = self._find_text(value[key])
|
|
if nested:
|
|
return nested
|
|
for nested_value in value.values():
|
|
nested = self._find_text(nested_value)
|
|
if nested:
|
|
return nested
|
|
return None
|
|
if isinstance(value, list):
|
|
for item in value:
|
|
nested = self._find_text(item)
|
|
if nested:
|
|
return nested
|
|
return None
|
|
|
|
def _extract_usage(self, response: Any) -> Dict[str, int]:
|
|
"""Extract best-effort token usage from nested agent responses."""
|
|
usage = self._find_usage_dict(response) or {}
|
|
prompt_tokens = int(usage.get("prompt_tokens") or usage.get("input_tokens") or 0)
|
|
completion_tokens = int(usage.get("completion_tokens") or usage.get("output_tokens") or 0)
|
|
total_tokens = int(usage.get("total_tokens") or prompt_tokens + completion_tokens)
|
|
return {
|
|
"prompt_tokens": prompt_tokens,
|
|
"completion_tokens": completion_tokens,
|
|
"total_tokens": total_tokens,
|
|
}
|
|
|
|
def _find_usage_dict(self, value: Any) -> Optional[Dict[str, Any]]:
|
|
"""Find a nested usage-like dict containing token counters."""
|
|
if isinstance(value, dict):
|
|
keys = set(value.keys())
|
|
if keys.intersection({"prompt_tokens", "completion_tokens", "total_tokens", "input_tokens", "output_tokens"}):
|
|
return value
|
|
for nested in value.values():
|
|
usage = self._find_usage_dict(nested)
|
|
if usage:
|
|
return usage
|
|
elif isinstance(value, list):
|
|
for item in value:
|
|
usage = self._find_usage_dict(item)
|
|
if usage:
|
|
return usage
|
|
return None
|
|
|
|
def _build_artifact(self, agent: SwarmAgent, response: Any, index: int) -> Dict[str, Any]:
|
|
"""Convert an agent response into a Heicode-visible artifact record."""
|
|
role = agent.role or "worker"
|
|
artifact_type = "code_patch" if role in {"backend", "frontend", "coder", "engineer"} else "document"
|
|
return {
|
|
"artifact_id": f"art_{self.swarm_id}_{role}_{index + 1}",
|
|
"artifact_type": artifact_type,
|
|
"title": f"{role} task delivery",
|
|
"summary": self._response_summary(response),
|
|
"uri": f"runtime://{self.swarm_id}/artifacts/{role}-{index + 1}",
|
|
"agent_instance_id": agent.agent_id,
|
|
"mime_type": "text/plain",
|
|
"stage": "development",
|
|
"checkpoint": "artifact_ready",
|
|
"metadata": {
|
|
"redacted": True,
|
|
"agent_role": role,
|
|
"runtime_deployment_id": self.swarm_id,
|
|
},
|
|
}
|
|
|
|
def _finalize_running_agents(self, terminal_status: SwarmAgentStatus, output: Optional[str] = None) -> None:
|
|
"""Ensure swarm-level completion/failure matches per-agent terminal states."""
|
|
agents = self.db.query(SwarmAgent).filter(SwarmAgent.swarm_id == self.swarm_id).all()
|
|
for agent in agents:
|
|
if agent.status == SwarmAgentStatus.RUNNING:
|
|
agent.status = terminal_status
|
|
agent.current_task = None
|
|
if output and not agent.output:
|
|
agent.output = output[:1000]
|