1100 lines
45 KiB
Python
1100 lines
45 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 k8s_manager import K8sManager
|
|
from .agent_client import SwarmAgentClient
|
|
from .artifact_store import store_text_artifact, runtime_artifact_uri, runtime_artifact_download_path
|
|
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
|
|
self.k8s_manager: Optional[K8sManager] = 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()
|
|
|
|
# Ensure all sub-mode agent runtimes exist before dispatching work.
|
|
for agent in swarm_agents:
|
|
try:
|
|
deployment_info = self._ensure_agent_runtime(agent)
|
|
agent.namespace = deployment_info.get("namespace") or agent.namespace
|
|
agent.pod_name = deployment_info.get("pod_name") or agent.pod_name
|
|
agent.service_url = deployment_info.get("service_url") or agent.service_url
|
|
agent.external_ip = deployment_info.get("external_ip")
|
|
self.db.commit()
|
|
except Exception as exc:
|
|
agent.status = SwarmAgentStatus.FAILED
|
|
agent.output = f"Failed to deploy sub-mode agent runtime: {exc}"
|
|
self.db.commit()
|
|
await self._emit_agent_event(
|
|
"agent.crashed",
|
|
agent.agent_id,
|
|
{
|
|
"agent_role": agent.role,
|
|
"status": "failed",
|
|
"error_message": agent.output,
|
|
},
|
|
)
|
|
continue
|
|
|
|
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
|
|
|
|
def _ensure_agent_runtime(self, agent: SwarmAgent) -> Dict[str, Any]:
|
|
"""Ensure the ordinary sub-mode agent Pod and Service exist."""
|
|
if not self.swarm:
|
|
raise ValueError("Swarm not initialized")
|
|
if not self.k8s_manager:
|
|
self.k8s_manager = K8sManager()
|
|
|
|
namespace = self.k8s_manager.create_swarm_namespace(self.swarm_id, agent.role)
|
|
agent_config = {
|
|
"role": agent.role,
|
|
"template": agent.template,
|
|
"model": agent.model,
|
|
"capabilities": agent.capabilities or [],
|
|
"system_prompt": agent.system_prompt,
|
|
"billing_context": {
|
|
**((self.swarm.project_context or {}).get("billing_context") or {}),
|
|
"max_tokens": ((self.swarm.project_context or {}).get("budget") or {}).get("max_tokens"),
|
|
},
|
|
}
|
|
return self.k8s_manager.deploy_swarm_agent(
|
|
self.swarm_id,
|
|
agent.agent_id,
|
|
agent_config,
|
|
namespace,
|
|
)
|
|
|
|
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()
|
|
|
|
artifacts = self._ensure_result_artifacts(result)
|
|
|
|
# Update swarm status
|
|
self.swarm.status = SwarmStatus.COMPLETED
|
|
self.swarm.completed_at = datetime.utcnow()
|
|
self.swarm.progress = 100
|
|
self.swarm.artifacts = artifacts
|
|
self._finalize_running_agents(SwarmAgentStatus.COMPLETED)
|
|
self.db.commit()
|
|
await self._emit_artifacts(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:
|
|
failure_artifacts = self._build_failure_artifacts(str(e))
|
|
self.swarm.status = SwarmStatus.FAILED
|
|
self.swarm.error_message = str(e)
|
|
self.swarm.artifacts = failure_artifacts
|
|
self._finalize_running_agents(SwarmAgentStatus.FAILED, str(e))
|
|
self.db.commit()
|
|
await self._emit_artifacts(failure_artifacts)
|
|
await self._emit_phase("failed", "Swarm execution failed")
|
|
await self._emit_status("failed", {"error": str(e)})
|
|
await self._emit_usage()
|
|
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
|
|
]
|
|
if not executable_agent_records:
|
|
all_agents = self.db.query(SwarmAgent).filter(SwarmAgent.swarm_id == self.swarm_id).all()
|
|
for agent in all_agents:
|
|
await self._emit(
|
|
"task.blocked",
|
|
agent_instance_id=agent.agent_id,
|
|
payload={
|
|
"task_id": f"task_{agent.agent_id}",
|
|
"agent_role": agent.role,
|
|
"status": "blocked",
|
|
"summary": "No executable runtime client was available for this sub-mode agent.",
|
|
"runtime_deployment_id": self.swarm_id,
|
|
},
|
|
)
|
|
raise RuntimeError("No executable sub-mode agents were available")
|
|
|
|
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]] = []
|
|
success_count = 0
|
|
failure_summaries: List[str] = []
|
|
for index, result in enumerate(results_list):
|
|
agent_record = executable_agent_records[index]
|
|
if isinstance(result, Exception):
|
|
failure_summaries.append(f"{agent_record.role}: {result}")
|
|
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),
|
|
}
|
|
)
|
|
success_count += 1
|
|
artifacts.append(self._build_artifact(agent_record, result, index))
|
|
|
|
if failure_summaries:
|
|
raise RuntimeError("Sub-mode agent task failed: " + "; ".join(failure_summaries))
|
|
if success_count == 0:
|
|
raise RuntimeError("All sub-mode agents failed before returning deliverables")
|
|
|
|
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})
|
|
if isinstance(response, dict) and response.get("error"):
|
|
error = response.get("error") or {}
|
|
message_text = error.get("message") if isinstance(error, dict) else str(error)
|
|
error_data = error.get("data") if isinstance(error, dict) else {}
|
|
if isinstance(error_data, dict):
|
|
error_message = SwarmMessage(
|
|
message_id=str(uuid.uuid4()),
|
|
swarm_id=self.swarm_id,
|
|
from_agent_id=client.agent_id,
|
|
to_agent_id=None,
|
|
message_type="error",
|
|
content=message_text or "A2A agent returned an error response",
|
|
message_metadata={
|
|
"newapi_request_id": error_data.get("request_id"),
|
|
"model_status_code": error_data.get("status_code"),
|
|
"model_response_preview": (error_data.get("response_text") or "")[:1000],
|
|
},
|
|
)
|
|
self.db.add(error_message)
|
|
self.db.commit()
|
|
runtime_error = RuntimeError(message_text or "A2A agent returned an error response")
|
|
if isinstance(error_data, dict):
|
|
setattr(runtime_error, "request_id", error_data.get("request_id"))
|
|
raise runtime_error
|
|
usage = self._extract_usage(response)
|
|
model_metadata = self._extract_model_metadata(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,
|
|
"newapi_request_id": model_metadata.get("newapi_request_id"),
|
|
"model_api_format": model_metadata.get("api_format"),
|
|
},
|
|
)
|
|
|
|
# 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={
|
|
"model_usage": usage,
|
|
**model_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,
|
|
"newapi_request_id": model_metadata.get("newapi_request_id"),
|
|
"model_usage": usage,
|
|
},
|
|
)
|
|
|
|
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),
|
|
"newapi_request_id": self._request_id_from_error(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,
|
|
"newapi_request_id": self._request_id_from_error(e),
|
|
},
|
|
)
|
|
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 {}
|
|
usage = self._aggregate_model_usage_from_messages()
|
|
if self.swarm and usage["total_tokens"] > self.swarm.tokens_used:
|
|
self.swarm.tokens_used = usage["total_tokens"]
|
|
self.db.commit()
|
|
model_tokens = self.swarm.tokens_used if self.swarm else usage["total_tokens"]
|
|
if not model_tokens:
|
|
return
|
|
await self._emit(
|
|
"budget.alert",
|
|
payload={
|
|
"model_id": billing_context.get("default_model_id") or "unknown",
|
|
"model_tokens": model_tokens,
|
|
"prompt_tokens": usage["prompt_tokens"],
|
|
"completion_tokens": usage["completion_tokens"],
|
|
"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"),
|
|
},
|
|
},
|
|
)
|
|
|
|
def _aggregate_model_usage_from_messages(self) -> Dict[str, int]:
|
|
"""Aggregate persisted model usage metadata for Runtime usage callbacks."""
|
|
totals = {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}
|
|
messages = (
|
|
self.db.query(SwarmMessage)
|
|
.filter(SwarmMessage.swarm_id == self.swarm_id)
|
|
.all()
|
|
)
|
|
for message in messages:
|
|
usage = (message.message_metadata or {}).get("model_usage") or {}
|
|
if not isinstance(usage, dict):
|
|
continue
|
|
totals["prompt_tokens"] += int(usage.get("prompt_tokens") or usage.get("input_tokens") or 0)
|
|
totals["completion_tokens"] += int(usage.get("completion_tokens") or usage.get("output_tokens") or 0)
|
|
totals["total_tokens"] += int(
|
|
usage.get("total_tokens")
|
|
or (usage.get("prompt_tokens") or usage.get("input_tokens") or 0)
|
|
+ (usage.get("completion_tokens") or usage.get("output_tokens") or 0)
|
|
)
|
|
return totals
|
|
|
|
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._extract_deliverable_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 _extract_deliverable_text(self, response: Any) -> Optional[str]:
|
|
"""Extract complete user-facing artifact text from agent responses."""
|
|
artifact_text = self._find_a2a_artifact_text(response)
|
|
if artifact_text:
|
|
return artifact_text
|
|
return self._find_text(response)
|
|
|
|
def _find_a2a_artifact_text(self, value: Any) -> Optional[str]:
|
|
"""Prefer A2A artifact part text over generic JSON-RPC bookkeeping fields."""
|
|
if isinstance(value, dict):
|
|
parts = value.get("parts")
|
|
if isinstance(parts, list):
|
|
texts = [
|
|
str(part.get("text"))
|
|
for part in parts
|
|
if isinstance(part, dict) and part.get("kind") == "text" and part.get("text")
|
|
]
|
|
if texts:
|
|
return "\n".join(texts)
|
|
|
|
artifacts = value.get("artifacts")
|
|
if isinstance(artifacts, list):
|
|
texts = []
|
|
for artifact in artifacts:
|
|
nested = self._find_a2a_artifact_text(artifact)
|
|
if nested:
|
|
texts.append(nested)
|
|
if texts:
|
|
return "\n\n".join(texts)
|
|
|
|
for key in ("result", "data", "artifact"):
|
|
nested = self._find_a2a_artifact_text(value.get(key))
|
|
if nested:
|
|
return nested
|
|
elif isinstance(value, list):
|
|
texts = []
|
|
for item in value:
|
|
nested = self._find_a2a_artifact_text(item)
|
|
if nested:
|
|
texts.append(nested)
|
|
if texts:
|
|
return "\n\n".join(texts)
|
|
return None
|
|
|
|
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):
|
|
skipped_keys = {"jsonrpc", "id", "messageId", "artifactId", "kind", "role", "code"}
|
|
for key in ("text", "content", "message", "output", "result"):
|
|
if key in value:
|
|
nested = self._find_text(value[key])
|
|
if nested:
|
|
return nested
|
|
for key, nested_value in value.items():
|
|
if key in skipped_keys:
|
|
continue
|
|
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 _extract_model_metadata(self, response: Any) -> Dict[str, Any]:
|
|
"""Extract model request metadata from nested A2A responses."""
|
|
metadata = self._find_model_metadata(response) or {}
|
|
request_id = (
|
|
metadata.get("newapi_request_id")
|
|
or metadata.get("request_id")
|
|
or metadata.get("response_id")
|
|
)
|
|
return {
|
|
"newapi_request_id": request_id,
|
|
"response_id": metadata.get("response_id"),
|
|
"model": metadata.get("model"),
|
|
"api_format": metadata.get("api_format"),
|
|
"endpoint": metadata.get("endpoint"),
|
|
}
|
|
|
|
def _find_model_metadata(self, value: Any) -> Optional[Dict[str, Any]]:
|
|
"""Find nested model metadata emitted by an A2A agent."""
|
|
if isinstance(value, dict):
|
|
if any(key in value for key in ("newapi_request_id", "request_id", "response_id", "api_format")):
|
|
return value
|
|
for nested in value.values():
|
|
found = self._find_model_metadata(nested)
|
|
if found:
|
|
return found
|
|
elif isinstance(value, list):
|
|
for item in value:
|
|
found = self._find_model_metadata(item)
|
|
if found:
|
|
return found
|
|
return None
|
|
|
|
def _request_id_from_error(self, error: Exception) -> Optional[str]:
|
|
"""Best-effort extraction for errors raised after model calls."""
|
|
return getattr(error, "request_id", None)
|
|
|
|
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"
|
|
artifact_id = f"art_{self.swarm_id}_{role}_{index + 1}"
|
|
content = self._extract_deliverable_text(response)
|
|
if not content:
|
|
content = json.dumps(response, ensure_ascii=False, indent=2) if isinstance(response, (dict, list)) else str(response)
|
|
stored = self._store_artifact_content(artifact_id, content)
|
|
metadata = {
|
|
"redacted": True,
|
|
"agent_role": role,
|
|
"runtime_deployment_id": self.swarm_id,
|
|
}
|
|
if stored:
|
|
metadata.update(
|
|
{
|
|
"content_hash": stored.content_hash,
|
|
"download_path": runtime_artifact_download_path(self.swarm_id, artifact_id),
|
|
}
|
|
)
|
|
return {
|
|
"artifact_id": artifact_id,
|
|
"artifact_type": artifact_type,
|
|
"title": f"{role} task delivery",
|
|
"summary": content[:1000],
|
|
"uri": stored.uri if stored else runtime_artifact_uri(self.swarm_id, artifact_id),
|
|
"agent_instance_id": agent.agent_id,
|
|
"mime_type": stored.mime_type if stored else "text/plain",
|
|
"size_bytes": stored.size_bytes if stored else len(content.encode("utf-8")),
|
|
"stage": "development",
|
|
"checkpoint": "artifact_ready",
|
|
"metadata": metadata,
|
|
}
|
|
|
|
def _ensure_result_artifacts(self, result: Dict[str, Any]) -> List[Dict[str, Any]]:
|
|
"""Guarantee every terminal sub-mode run exposes at least one Manager artifact."""
|
|
artifacts = result.get("artifacts") or []
|
|
if artifacts:
|
|
return artifacts
|
|
return [self._build_runtime_summary_artifact(result)]
|
|
|
|
def _build_failure_artifacts(self, error: str) -> List[Dict[str, Any]]:
|
|
"""Build a visible artifact for failed or blocked sub-mode executions."""
|
|
return [
|
|
self._build_runtime_summary_artifact(
|
|
{
|
|
"phases": [],
|
|
"error": error,
|
|
},
|
|
failed=True,
|
|
)
|
|
]
|
|
|
|
def _build_runtime_summary_artifact(
|
|
self,
|
|
result: Optional[Dict[str, Any]] = None,
|
|
*,
|
|
failed: bool = False,
|
|
) -> Dict[str, Any]:
|
|
"""Create a fallback deliverable when agents return no concrete files."""
|
|
result = result or {}
|
|
agents = self.db.query(SwarmAgent).filter(SwarmAgent.swarm_id == self.swarm_id).all()
|
|
agent_summaries = []
|
|
for agent in agents:
|
|
detail = f"{agent.role}:{agent.status.value}"
|
|
if agent.output:
|
|
detail += f" - {agent.output[:300]}"
|
|
agent_summaries.append(detail)
|
|
|
|
error = result.get("error") or (self.swarm.error_message if self.swarm else None)
|
|
if failed:
|
|
title = "Runtime execution failed"
|
|
summary = error or "Runtime failed before returning a concrete artifact."
|
|
artifact_type = "other"
|
|
checkpoint = "failed"
|
|
else:
|
|
title = "Runtime execution summary"
|
|
summary = (
|
|
"Runtime completed but no concrete agent artifact was returned. "
|
|
"Review the sub-mode task objective, agent statuses, and logs for the generated result or failure reason."
|
|
)
|
|
artifact_type = "document"
|
|
checkpoint = "artifact_ready"
|
|
|
|
if agent_summaries:
|
|
summary = f"{summary}\nAgents: " + "; ".join(agent_summaries)
|
|
if self.swarm and self.swarm.task_description:
|
|
summary = f"Task: {self.swarm.task_description}\n{summary}"
|
|
|
|
artifact_id = f"art_{self.swarm_id}_{'failure' if failed else 'summary'}"
|
|
stored = self._store_artifact_content(artifact_id, summary)
|
|
metadata = {
|
|
"redacted": True,
|
|
"source": "agent-manager-sub-mode-runtime",
|
|
"runtime_deployment_id": self.swarm_id,
|
|
"agent_count": len(agents),
|
|
}
|
|
if stored:
|
|
metadata.update(
|
|
{
|
|
"content_hash": stored.content_hash,
|
|
"download_path": runtime_artifact_download_path(self.swarm_id, artifact_id),
|
|
}
|
|
)
|
|
|
|
return {
|
|
"artifact_id": artifact_id,
|
|
"artifact_type": artifact_type,
|
|
"title": title,
|
|
"summary": summary[:2000],
|
|
"uri": stored.uri if stored else runtime_artifact_uri(self.swarm_id, artifact_id),
|
|
"agent_instance_id": None,
|
|
"mime_type": stored.mime_type if stored else "text/plain",
|
|
"size_bytes": stored.size_bytes if stored else len(summary.encode("utf-8")),
|
|
"stage": "development",
|
|
"checkpoint": checkpoint,
|
|
"metadata": metadata,
|
|
}
|
|
|
|
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 in {SwarmAgentStatus.PENDING, SwarmAgentStatus.RUNNING}:
|
|
agent.status = terminal_status
|
|
agent.current_task = None
|
|
if output and not agent.output:
|
|
agent.output = output[:1000]
|
|
|
|
def _store_artifact_content(self, artifact_id: str, content: str):
|
|
"""Persist full artifact content without blocking callback/status projection on storage errors."""
|
|
try:
|
|
return store_text_artifact(self.swarm_id, artifact_id, content, mime_type="text/plain")
|
|
except Exception as e:
|
|
logger.warning("Failed to persist runtime artifact %s: %s", artifact_id, e)
|
|
return None
|