1636 lines
70 KiB
Python
1636 lines
70 KiB
Python
"""Sub-mode runtime orchestrator."""
|
|
|
|
import asyncio
|
|
import json
|
|
import logging
|
|
import os
|
|
import re
|
|
import shutil
|
|
import subprocess
|
|
import tempfile
|
|
import uuid
|
|
from datetime import datetime, timedelta
|
|
from pathlib import Path
|
|
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 api.agnet.vault_client import vault_client
|
|
from .agent_client import SwarmAgentClient
|
|
from .artifact_store import (
|
|
store_project_artifact,
|
|
store_text_artifact,
|
|
runtime_artifact_archive_path,
|
|
runtime_artifact_download_path,
|
|
runtime_artifact_file_path,
|
|
runtime_artifact_manifest_path,
|
|
runtime_artifact_uri,
|
|
)
|
|
from .callback_client import CallbackDeliveryClient
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
_CODE_BLOCK_RE = re.compile(r"```(?P<lang>[a-zA-Z0-9_+-]*)\n(?P<body>.*?)```", re.DOTALL)
|
|
_FILENAME_HINT_RE = re.compile(
|
|
r"^\s*(?:#|//|/\*+|\*|--)?\s*(?P<path>[A-Za-z0-9_.\-/]+\.[A-Za-z0-9_]+)\s*(?:\*/)?\s*$"
|
|
)
|
|
|
|
|
|
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
|
|
self._git_workspace_dir: Optional[Path] = None
|
|
|
|
async def initialize(self) -> bool:
|
|
"""
|
|
Initialize swarm - create all agent pods.
|
|
|
|
Returns:
|
|
True if successful, False otherwise
|
|
"""
|
|
try:
|
|
# Load the internal runtime-run record from the database.
|
|
self.swarm = self.db.query(Swarm).filter(
|
|
Swarm.swarm_id == self.swarm_id
|
|
).first()
|
|
|
|
if not self.swarm:
|
|
logger.error(f"Runtime deployment {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 sub-mode runtime deployment {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)
|
|
artifacts = await self._attach_git_delivery_refs(artifacts)
|
|
|
|
# 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 self._is_retryable_runtime_error(message_text or "") and agent_record:
|
|
retry_task = self._build_retry_task(agent_record, task, message_text or "gateway timeout")
|
|
retry_message = SwarmMessage(
|
|
message_id=str(uuid.uuid4()),
|
|
swarm_id=self.swarm_id,
|
|
from_agent_id=None,
|
|
to_agent_id=client.agent_id,
|
|
message_type="task_retry",
|
|
content=retry_task,
|
|
message_metadata={
|
|
"retry_reason": message_text,
|
|
"retryable": True,
|
|
},
|
|
)
|
|
self.db.add(retry_message)
|
|
self.db.commit()
|
|
self.swarm.total_messages += 1
|
|
agent_record.current_task = retry_task
|
|
self.db.commit()
|
|
await self._emit_tool_event(
|
|
"sk_tool.called",
|
|
client.agent_id,
|
|
{
|
|
"tool_name": "agent_task_retry",
|
|
"tool_invocation_id": retry_message.message_id,
|
|
"summary": "Retrying agent task with a reduced output contract",
|
|
"arguments_redacted": True,
|
|
},
|
|
)
|
|
response = await client.send_message({"text": retry_task})
|
|
if not (isinstance(response, dict) and response.get("error")):
|
|
usage = self._extract_usage(response)
|
|
model_metadata = self._extract_model_metadata(response)
|
|
await self._emit_tool_event(
|
|
"sk_tool.completed",
|
|
client.agent_id,
|
|
{
|
|
"tool_name": "agent_task_retry",
|
|
"tool_invocation_id": retry_message.message_id,
|
|
"summary": "Agent retry 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"),
|
|
},
|
|
)
|
|
|
|
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()
|
|
if self._git_workspace_dir:
|
|
workspace_root = self._git_workspace_dir.parent
|
|
if workspace_root.exists():
|
|
shutil.rmtree(workspace_root, ignore_errors=True)
|
|
self._git_workspace_dir = None
|
|
|
|
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)
|
|
)
|
|
prompt_parts.append(self._role_output_contract(agent.role))
|
|
return "\n".join(prompt_parts)
|
|
|
|
def _role_output_contract(self, role: Optional[str]) -> str:
|
|
"""Return a compact output contract tuned for runtime stability."""
|
|
normalized_role = (role or "worker").lower()
|
|
common = (
|
|
"Output contract: keep the response concise and implementation-oriented. "
|
|
"Prefer a minimal viable skeleton over a full project. "
|
|
"Do not exceed 12 bullets. Do not exceed 120 lines of code total."
|
|
)
|
|
role_contracts = {
|
|
"backend": (
|
|
"Return only: 1) a short API/data-model summary, "
|
|
"2) one compact Python/FastAPI code skeleton, "
|
|
"3) a brief test checklist."
|
|
),
|
|
"frontend": (
|
|
"Return only: 1) a short UI/component summary, "
|
|
"2) one compact React/JSX/CSS skeleton, "
|
|
"3) a brief interaction checklist."
|
|
),
|
|
"reviewer": (
|
|
"Return only: 1) major risks, 2) test cases, 3) release/blocking notes. "
|
|
"No long prose and no large code blocks."
|
|
),
|
|
"architect": (
|
|
"Return only: 1) architecture outline, 2) core modules, 3) key interfaces. "
|
|
"No large code blocks."
|
|
),
|
|
}
|
|
fallback = "Return a short implementation summary and one minimal code skeleton if needed."
|
|
return f"{common} {role_contracts.get(normalized_role, fallback)}"
|
|
|
|
def _build_retry_task(self, agent: Optional[SwarmAgent], original_task: str, error: str) -> str:
|
|
"""Build a smaller retry prompt when the first generation overloads the model gateway."""
|
|
role = agent.role if agent else "worker"
|
|
retry_instructions = (
|
|
"Previous attempt failed at the model gateway. Retry with a much smaller response. "
|
|
"Return only the single most important implementation skeleton for your role. "
|
|
"Limit output to at most 6 bullets and at most 60 lines of code total. "
|
|
"Skip optional explanations, examples, and secondary files."
|
|
)
|
|
return f"{original_task}\nRole: {role}\nRetry reason: {error}\n{retry_instructions}"
|
|
|
|
def _is_retryable_runtime_error(self, message: str) -> bool:
|
|
"""Classify model-gateway errors that benefit from a smaller retry."""
|
|
normalized = (message or "").lower()
|
|
retry_markers = (
|
|
"504 gateway time-out",
|
|
"504 gateway timeout",
|
|
"timed out",
|
|
"timeout",
|
|
"upstream request timeout",
|
|
)
|
|
return any(marker in normalized for marker in retry_markers)
|
|
|
|
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"
|
|
is_structural_code_role = role in {"backend", "frontend", "coder", "engineer", "fullstack"}
|
|
artifact_type = "project_folder" if is_structural_code_role 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)
|
|
if is_structural_code_role:
|
|
files, root_dir = self._project_files_from_content(role, content)
|
|
stored_project = self._store_project_artifact(artifact_id, root_dir, files)
|
|
metadata = {
|
|
"redacted": True,
|
|
"agent_role": role,
|
|
"source_agent_role": role,
|
|
"runtime_deployment_id": self.swarm_id,
|
|
"summary_only": False,
|
|
"artifact_layout": "project_folder",
|
|
"primary_read_path": "manifest",
|
|
"root_dir": stored_project.root_dir if stored_project else root_dir,
|
|
"file_count": stored_project.file_count if stored_project else len(files),
|
|
"directory_count": stored_project.directory_count if stored_project else 1,
|
|
"manifest_uri": runtime_artifact_manifest_path(self.swarm_id, artifact_id),
|
|
"archive_uri": runtime_artifact_archive_path(self.swarm_id, artifact_id),
|
|
"project_revision": 1,
|
|
}
|
|
if stored_project:
|
|
metadata.update(
|
|
{
|
|
"content_hash": stored_project.content_hash,
|
|
"download_path": runtime_artifact_archive_path(self.swarm_id, artifact_id),
|
|
"files_base_uri": runtime_artifact_file_path(self.swarm_id, artifact_id, ""),
|
|
"delivery_ref": {
|
|
"kind": "runtime_artifact",
|
|
"artifact_id": artifact_id,
|
|
"project_revision": 1,
|
|
},
|
|
}
|
|
)
|
|
git_ref = self._git_ref_metadata()
|
|
if git_ref:
|
|
metadata["git_ref"] = git_ref
|
|
return {
|
|
"artifact_id": artifact_id,
|
|
"artifact_type": artifact_type,
|
|
"title": f"{role} task delivery",
|
|
"summary": content[:1000],
|
|
"uri": stored_project.uri if stored_project else runtime_artifact_uri(self.swarm_id, artifact_id),
|
|
"agent_instance_id": agent.agent_id,
|
|
"mime_type": "application/zip",
|
|
"size_bytes": stored_project.archive_path.stat().st_size if stored_project else len(content.encode("utf-8")),
|
|
"stage": "development",
|
|
"checkpoint": "artifact_ready",
|
|
"metadata": metadata,
|
|
}
|
|
|
|
stored = self._store_artifact_content(artifact_id, content)
|
|
metadata = {
|
|
"redacted": True,
|
|
"agent_role": role,
|
|
"source_agent_role": role,
|
|
"runtime_deployment_id": self.swarm_id,
|
|
"summary_only": False,
|
|
"artifact_layout": "single_file_content",
|
|
"primary_read_path": "content",
|
|
}
|
|
if stored:
|
|
metadata.update(
|
|
{
|
|
"content_hash": stored.content_hash,
|
|
"download_path": runtime_artifact_download_path(self.swarm_id, artifact_id),
|
|
}
|
|
)
|
|
git_ref = self._git_ref_metadata()
|
|
if git_ref:
|
|
metadata["git_ref"] = git_ref
|
|
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),
|
|
"synthesized": True,
|
|
"summary_only": True,
|
|
"artifact_layout": "single_file_content",
|
|
"primary_read_path": "content",
|
|
"project_revision": 0,
|
|
}
|
|
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
|
|
|
|
def _store_project_artifact(self, artifact_id: str, root_dir: str, files: Dict[str, str]):
|
|
"""Persist a project-folder artifact without blocking runtime completion."""
|
|
try:
|
|
return store_project_artifact(self.swarm_id, artifact_id, root_dir=root_dir, files=files)
|
|
except Exception as e:
|
|
logger.warning("Failed to persist runtime project artifact %s: %s", artifact_id, e)
|
|
return None
|
|
|
|
def _git_context(self) -> Optional[Dict[str, Any]]:
|
|
"""Return repository context when Runtime should write delivery branches."""
|
|
context = self._callback_context()
|
|
repo_url = context.get("repo_url")
|
|
git_binding_id = context.get("git_binding_id")
|
|
resource_grants = context.get("resource_grants") or []
|
|
git_grant = self._find_git_resource_grant(resource_grants, git_binding_id)
|
|
if not repo_url and git_grant:
|
|
repo_url = (
|
|
git_grant.get("external_ref")
|
|
or git_grant.get("repo_url")
|
|
or (git_grant.get("metadata") or {}).get("repo_url")
|
|
)
|
|
if not repo_url:
|
|
return None
|
|
return {
|
|
"repo_url": repo_url,
|
|
"base_branch": context.get("branch") or "main",
|
|
"git_binding_id": git_binding_id or (git_grant or {}).get("resource_id") or (git_grant or {}).get("grant_id"),
|
|
"git_grant": git_grant,
|
|
}
|
|
|
|
def _find_git_resource_grant(self, resource_grants: List[Dict[str, Any]], git_binding_id: Optional[str]) -> Optional[Dict[str, Any]]:
|
|
"""Select the git resource grant attached to this runtime task."""
|
|
if git_binding_id:
|
|
for grant in resource_grants:
|
|
if git_binding_id in {
|
|
grant.get("resource_id"),
|
|
grant.get("grant_id"),
|
|
grant.get("binding_id"),
|
|
}:
|
|
return grant
|
|
for grant in resource_grants:
|
|
resource_type = (grant.get("resource_type") or grant.get("type") or "").lower()
|
|
if resource_type == "git":
|
|
return grant
|
|
return None
|
|
|
|
async def _git_credentials(self) -> Optional[tuple[str, str]]:
|
|
"""Resolve git credentials from Manager-provided resource grants first, then env fallback."""
|
|
git_context = self._git_context() or {}
|
|
git_grant = git_context.get("git_grant") or {}
|
|
metadata = git_grant.get("metadata") or {}
|
|
username = (
|
|
metadata.get("username")
|
|
or git_grant.get("username")
|
|
or os.getenv("GITEE_USERNAME")
|
|
or ""
|
|
)
|
|
secret_ref = git_grant.get("secret_ref") or git_grant.get("ref")
|
|
if secret_ref:
|
|
secret_value = await vault_client.get_secret(secret_ref)
|
|
parsed = self._parse_git_secret(secret_value, username)
|
|
if parsed:
|
|
return parsed
|
|
|
|
# Compatibility fallback for older deployments that still rely on env injection.
|
|
token = os.getenv("GITEE_TOKEN") or ""
|
|
password = os.getenv("GITEE_PASSWORD") or ""
|
|
if token and token != "your-gitee-token":
|
|
return username or "oauth2", token
|
|
if password and password != "your-gitee-password":
|
|
return username or "git", password
|
|
return None
|
|
|
|
def _parse_git_secret(self, secret_value: Any, default_username: str) -> Optional[tuple[str, str]]:
|
|
"""Parse a git secret payload into username/password credentials."""
|
|
if not secret_value:
|
|
return None
|
|
if isinstance(secret_value, str):
|
|
stripped = secret_value.strip()
|
|
if not stripped:
|
|
return None
|
|
if stripped.startswith("{"):
|
|
try:
|
|
secret_value = json.loads(stripped)
|
|
except Exception:
|
|
return default_username or "oauth2", stripped
|
|
else:
|
|
return default_username or "oauth2", stripped
|
|
|
|
if isinstance(secret_value, dict):
|
|
username = (
|
|
secret_value.get("username")
|
|
or secret_value.get("user")
|
|
or default_username
|
|
or "oauth2"
|
|
)
|
|
password = (
|
|
secret_value.get("token")
|
|
or secret_value.get("password")
|
|
or secret_value.get("pat")
|
|
or secret_value.get("access_token")
|
|
)
|
|
if password:
|
|
return username, str(password)
|
|
return None
|
|
|
|
def _inject_git_credentials(self, repo_url: str, username: str, password: str) -> str:
|
|
"""Inject credentials into an HTTP(S) git URL without persisting them."""
|
|
if "://" not in repo_url:
|
|
return repo_url
|
|
proto, rest = repo_url.split("://", 1)
|
|
if "@" in rest:
|
|
rest = rest.split("@", 1)[1]
|
|
return f"{proto}://{username}:{password}@{rest}"
|
|
|
|
def _run_git(self, args: List[str], cwd: Path, *, timeout: int = 120) -> subprocess.CompletedProcess:
|
|
"""Run a git command and raise on failure."""
|
|
result = subprocess.run(
|
|
args,
|
|
cwd=str(cwd),
|
|
capture_output=True,
|
|
text=True,
|
|
timeout=timeout,
|
|
)
|
|
if result.returncode != 0:
|
|
raise RuntimeError(result.stderr.strip() or result.stdout.strip() or f"git command failed: {' '.join(args)}")
|
|
return result
|
|
|
|
async def _ensure_git_workspace(self) -> Optional[Path]:
|
|
"""Clone the target repository once and reuse it for role and delivery branches."""
|
|
git_context = self._git_context()
|
|
credentials = await self._git_credentials()
|
|
if not git_context or not credentials:
|
|
return None
|
|
if self._git_workspace_dir and self._git_workspace_dir.exists():
|
|
return self._git_workspace_dir
|
|
|
|
username, password = credentials
|
|
repo_url = git_context["repo_url"]
|
|
auth_url = self._inject_git_credentials(repo_url, username, password)
|
|
workspace = Path(tempfile.mkdtemp(prefix=f"swarm_git_{self.swarm_id}_"))
|
|
repo_dir = workspace / "repo"
|
|
self._run_git(["git", "clone", "--branch", git_context["base_branch"], auth_url, str(repo_dir)], workspace, timeout=180)
|
|
self._run_git(["git", "remote", "set-url", "origin", repo_url], repo_dir)
|
|
self._run_git(["git", "config", "user.name", "heicode-agent"], repo_dir)
|
|
self._run_git(["git", "config", "user.email", "bot@heicode.local"], repo_dir)
|
|
self._git_workspace_dir = repo_dir
|
|
return repo_dir
|
|
|
|
def _write_project_files_to_repo(self, repo_dir: Path, files: Dict[str, str]) -> None:
|
|
"""Materialize project artifact files into a git workspace."""
|
|
for relative_path, content in files.items():
|
|
file_path = repo_dir / relative_path
|
|
file_path.parent.mkdir(parents=True, exist_ok=True)
|
|
file_path.write_text(content, encoding="utf-8")
|
|
|
|
def _git_repo_files_from_artifact(self, artifact: Dict[str, Any]) -> Dict[str, str]:
|
|
"""Read back stored project artifact files for git delivery."""
|
|
artifact_id = artifact.get("artifact_id")
|
|
project = store = None
|
|
from .artifact_store import load_runtime_project_artifact
|
|
project = load_runtime_project_artifact(self.swarm_id, artifact_id)
|
|
if not project:
|
|
return {}
|
|
files = {}
|
|
root_dir = project.artifact_dir / project.root_dir
|
|
for file_path in sorted(root_dir.rglob("*")):
|
|
if file_path.is_file():
|
|
rel = file_path.relative_to(root_dir).as_posix()
|
|
files[rel] = file_path.read_text(encoding="utf-8", errors="replace")
|
|
return files
|
|
|
|
async def _attach_git_delivery_refs(self, artifacts: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
|
|
"""Best-effort git branch/commit materialization for project-folder artifacts."""
|
|
repo_dir = await self._ensure_git_workspace()
|
|
if not repo_dir:
|
|
return artifacts
|
|
|
|
git_context = self._git_context() or {}
|
|
credentials = await self._git_credentials()
|
|
if not credentials:
|
|
return artifacts
|
|
username, password = credentials
|
|
origin_url = self._run_git(["git", "remote", "get-url", "origin"], repo_dir).stdout.strip()
|
|
auth_url = self._inject_git_credentials(origin_url, username, password)
|
|
base_branch = git_context["base_branch"]
|
|
|
|
role_branches = []
|
|
updated_artifacts = []
|
|
base_git_ref = self._git_ref_metadata() or {}
|
|
for artifact in artifacts:
|
|
metadata = artifact.get("metadata") or {}
|
|
if metadata.get("artifact_layout") != "project_folder":
|
|
updated_artifacts.append(artifact)
|
|
continue
|
|
|
|
role = metadata.get("source_agent_role") or metadata.get("agent_role") or "worker"
|
|
role_branch = f"agent/{role}/{self.swarm_id}"
|
|
self._run_git(["git", "checkout", base_branch], repo_dir)
|
|
self._run_git(["git", "checkout", "-B", role_branch], repo_dir)
|
|
files = self._git_repo_files_from_artifact(artifact)
|
|
if not files:
|
|
updated_artifacts.append(artifact)
|
|
continue
|
|
self._write_project_files_to_repo(repo_dir, files)
|
|
self._run_git(["git", "add", "-A"], repo_dir)
|
|
commit_message = f"sub-mode {role} delivery for {self.swarm_id}"
|
|
commit_result = subprocess.run(
|
|
["git", "commit", "-m", commit_message],
|
|
cwd=str(repo_dir),
|
|
capture_output=True,
|
|
text=True,
|
|
timeout=120,
|
|
)
|
|
if commit_result.returncode != 0 and "nothing to commit" not in (commit_result.stdout + commit_result.stderr).lower():
|
|
raise RuntimeError(commit_result.stderr.strip() or commit_result.stdout.strip() or "git commit failed")
|
|
self._run_git(["git", "push", auth_url, role_branch], repo_dir, timeout=180)
|
|
commit_sha = self._run_git(["git", "rev-parse", "HEAD"], repo_dir).stdout.strip()
|
|
metadata["git_ref"] = {
|
|
"provider": metadata.get("git_ref", {}).get("provider") or base_git_ref.get("provider") or "git",
|
|
"repo_url": origin_url,
|
|
"base_branch": base_branch,
|
|
"branch": role_branch,
|
|
"commit_sha": commit_sha,
|
|
"git_binding_id": git_context.get("git_binding_id"),
|
|
}
|
|
artifact["metadata"] = metadata
|
|
updated_artifacts.append(artifact)
|
|
role_branches.append(role_branch)
|
|
|
|
if role_branches:
|
|
delivery_branch = f"delivery/{self.swarm_id}"
|
|
self._run_git(["git", "checkout", base_branch], repo_dir)
|
|
self._run_git(["git", "checkout", "-B", delivery_branch], repo_dir)
|
|
for role_branch in role_branches:
|
|
merge_result = subprocess.run(
|
|
["git", "merge", "--no-ff", "--no-edit", role_branch],
|
|
cwd=str(repo_dir),
|
|
capture_output=True,
|
|
text=True,
|
|
timeout=180,
|
|
)
|
|
if merge_result.returncode != 0:
|
|
raise RuntimeError(merge_result.stderr.strip() or merge_result.stdout.strip() or f"git merge failed for {role_branch}")
|
|
self._run_git(["git", "push", auth_url, delivery_branch], repo_dir, timeout=180)
|
|
delivery_sha = self._run_git(["git", "rev-parse", "HEAD"], repo_dir).stdout.strip()
|
|
if self.swarm:
|
|
context = self._callback_context()
|
|
context["delivery_branch"] = delivery_branch
|
|
context["delivery_commit_sha"] = delivery_sha
|
|
self.swarm.project_context = context
|
|
self.db.commit()
|
|
for artifact in updated_artifacts:
|
|
metadata = artifact.get("metadata") or {}
|
|
if metadata.get("artifact_layout") == "project_folder":
|
|
metadata["delivery_ref"] = {
|
|
"kind": "git_branch",
|
|
"branch": delivery_branch,
|
|
"commit_sha": delivery_sha,
|
|
"project_revision": metadata.get("project_revision", 1),
|
|
}
|
|
artifact["metadata"] = metadata
|
|
|
|
return updated_artifacts
|
|
|
|
def _project_files_from_content(self, role: str, content: str) -> tuple[Dict[str, str], str]:
|
|
"""Convert a code-oriented response into a minimal project-folder file set."""
|
|
root_dir = f"{role}-delivery"
|
|
files: Dict[str, str] = {}
|
|
for match in _CODE_BLOCK_RE.finditer(content):
|
|
body = (match.group("body") or "").strip("\n")
|
|
if not body:
|
|
continue
|
|
lines = body.splitlines()
|
|
first_line = lines[0].strip() if lines else ""
|
|
filename_match = _FILENAME_HINT_RE.match(first_line)
|
|
if filename_match:
|
|
file_path = filename_match.group("path")
|
|
file_content = "\n".join(lines[1:]).lstrip("\n")
|
|
else:
|
|
file_path = self._default_file_path_for_role(role, match.group("lang") or "", len(files))
|
|
file_content = body
|
|
files[file_path] = file_content or ""
|
|
|
|
if not files:
|
|
files[self._default_file_path_for_role(role, "", 0)] = content.strip() + "\n"
|
|
|
|
files.setdefault("README.md", self._project_readme(role, content))
|
|
files.setdefault("heicode-artifact.json", json.dumps({"role": role, "artifact_layout": "project_folder"}, ensure_ascii=False, indent=2))
|
|
return files, root_dir
|
|
|
|
def _default_file_path_for_role(self, role: str, language: str, index: int) -> str:
|
|
"""Choose a stable fallback file path when the agent output omits file names."""
|
|
normalized_role = (role or "worker").lower()
|
|
language = (language or "").lower()
|
|
if normalized_role == "backend":
|
|
if language in {"python", "py"}:
|
|
return "backend/app.py" if index == 0 else f"backend/module_{index + 1}.py"
|
|
return "backend/implementation.txt"
|
|
if normalized_role == "frontend":
|
|
if language in {"jsx", "tsx", "javascript", "js", "typescript", "ts"}:
|
|
return "frontend/OrderPage.jsx" if index == 0 else f"frontend/component_{index + 1}.jsx"
|
|
if language == "css":
|
|
return "frontend/styles.css"
|
|
return "frontend/implementation.txt"
|
|
if normalized_role == "reviewer":
|
|
return "review/review.md"
|
|
return f"{normalized_role}/artifact_{index + 1}.txt"
|
|
|
|
def _project_readme(self, role: str, content: str) -> str:
|
|
"""Build a lightweight README for project-folder artifacts."""
|
|
return (
|
|
f"# {role} delivery\n\n"
|
|
"This project-folder artifact was synthesized by the sub-mode runtime from the agent response.\n\n"
|
|
"## Summary\n\n"
|
|
f"{content[:1500].strip()}\n"
|
|
)
|
|
|
|
def _git_ref_metadata(self) -> Optional[Dict[str, Any]]:
|
|
"""Build best-effort git delivery metadata from Manager-provided context."""
|
|
context = self._callback_context()
|
|
repo_url = context.get("repo_url")
|
|
if not repo_url:
|
|
return None
|
|
|
|
provider = "git"
|
|
repo_url_lower = str(repo_url).lower()
|
|
if "github.com" in repo_url_lower:
|
|
provider = "github"
|
|
elif "gitee" in repo_url_lower or "gitea" in repo_url_lower:
|
|
provider = "gitea"
|
|
|
|
return {
|
|
"provider": provider,
|
|
"repo_url": repo_url,
|
|
"base_branch": context.get("branch") or "main",
|
|
"branch": context.get("delivery_branch"),
|
|
"commit_sha": context.get("delivery_commit_sha"),
|
|
"git_binding_id": context.get("git_binding_id"),
|
|
}
|