# ⚠️ STALE / DO NOT USE FOR CURRENT DEPLOYS. # This ConfigMap embeds an OLD snapshot of the orchestrator source (no planner / review loop / # synthesis). The canonical deploy builds the image from Dockerfile.orchestrator and uses # k8s/orchestrator-deployment.yaml. Regenerate from the live source (e.g. # kubectl create configmap orchestrator-source --from-file=orchestrator/ --dry-run=client -o yaml) # if a source-mounted deployment is ever needed, otherwise prefer the image-based deployment. apiVersion: v1 data: __init__.py: | """Swarm orchestrator package.""" __version__ = "1.0.0" README.md: "# Swarm Orchestrator\n\nFastAPI-based orchestrator for coordinating agent pods in the K8s swarm system.\n\n## Features\n\n- **WebSocket Communication**: Real-time bidirectional communication with agent pods\n- **Agent Registry**: Redis-backed registration and heartbeat tracking\n- **Task Queue**: Task assignment and reassignment with failure recovery\n- **Handoff Mechanism**: Coordinate task handoffs between agents\n- **Failure Recovery**: Automatic detection and recovery from agent failures\n\n## Architecture\n\n### Components\n\n1. **FastAPI Server** (`main.py`)\n - REST API for external clients\n - WebSocket endpoint for agent communication\n - Background failure detection loop\n\n2. **Redis Client** (`redis_client.py`)\n \ - Async Redis connection wrapper\n - Key-value, hash, and list operations\n\n3. **Agent Registry** (`agent_registry.py`)\n - Agent registration and deregistration\n \ - Heartbeat tracking (30s timeout)\n - Status management: idle, busy, handoff-pending, failed\n\n4. **Handoff Manager** (`handoff_manager.py`)\n - Handoff initiation and validation\n - Handoff history tracking\n - Target agent availability checks\n\n5. **Task Queue** (`task_queue.py`)\n - Task creation and assignment\n \ - Retry logic (max 3 retries)\n - Automatic reassignment on agent failure\n\n## WebSocket Protocol\n\n### Agent Registration\n\n```json\n{\n \"type\": \"register\",\n \ \"capabilities\": [\"python\", \"javascript\"]\n}\n```\n\n### Heartbeat\n\n```json\n{\n \ \"type\": \"heartbeat\"\n}\n```\n\nResponse:\n```json\n{\n \"type\": \"heartbeat_ack\"\n}\n```\n\n### Handoff Request\n\n```json\n{\n \"type\": \"handoff\",\n \"target_agent_id\": \"agent-2\",\n \"task_context\": {\n \"task_id\": \"task-123\",\n \"description\": \"Continue implementation\",\n \"files\": [\"src/main.py\"]\n }\n}\n```\n\n### Task Status Updates\n\n```json\n{\n \"type\": \"task_start\",\n \"task_id\": \"task-123\"\n}\n```\n\n```json\n{\n \"type\": \"task_complete\",\n \"task_id\": \"task-123\"\n}\n```\n\n```json\n{\n \"type\": \"task_failed\",\n \"task_id\": \"task-123\",\n \"reason\": \"Compilation error\"\n}\n```\n\n## REST API Endpoints\n\n### Health Check\n- `GET /` - Basic health check\n- `GET /health` - Detailed health status\n\n### Agents\n- `GET /agents` - List all agents\n- `GET /agents/{agent_id}` - Get agent details\n- `GET /agents/idle` - List idle agents\n\n### Tasks\n- `POST /tasks` - Create new task\n- `POST /tasks/assign` - Manually assign task\n- `GET /tasks` - List all tasks (optional `?status=pending`)\n- `GET /tasks/{task_id}` - Get task details\n\n### Handoffs\n- `GET /handoffs` - List handoff history (optional `?agent_id=agent-1`)\n\n## Deployment\n\n### Build Docker Image\n\n```bash\ndocker build -f Dockerfile.orchestrator -t swarm-orchestrator:latest .\n```\n\n### Deploy to Kubernetes\n\n```bash\n# Deploy Redis\nkubectl apply -f k8s/redis-statefulset.yaml\n\n# Deploy orchestrator\nkubectl apply -f k8s/orchestrator-deployment.yaml\n```\n\n### Environment Variables\n\n- `REDIS_HOST` - Redis hostname (default: `redis-service`)\n- `REDIS_PORT` - Redis port (default: `6379`)\n- `REDIS_DB` - Redis database number (default: `0`)\n- `LOG_LEVEL` - Logging level (default: `INFO`)\n\n## Failure Recovery\n\nThe orchestrator implements automatic failure recovery:\n\n1. **Heartbeat Monitoring**: Agents must send heartbeat every 15s\n2. **Timeout Detection**: Agents with no heartbeat for 30s are marked as failed\n3. **Task Reassignment**: Tasks from failed agents are automatically reassigned\n4. **Retry Logic**: Failed tasks are retried up to 3 times\n5. **WebSocket Cleanup**: Failed agent connections are closed\n\n## Redis Data Model\n\n### Agent Registry\n- Key: `agent:{agent_id}`\n- Value: JSON with status, last_heartbeat, capabilities, current_task_id\n\n### Tasks\n- Key: `task:{task_id}`\n- Value: JSON with description, status, assigned_agent_id, retry_count\n\n### Task Queue\n- Key: `queue:pending`\n- Type: List of task_ids\n\n### Handoff History\n- Key: `handoff:{handoff_id}`\n- Value: JSON with source, target, timestamp, context, status\n\n## Development\n\n### Install Dependencies\n\n```bash\npip install -r orchestrator/requirements.txt\n```\n\n### Run Locally\n\n```bash\n# Start Redis\ndocker run -d -p 6379:6379 redis:7-alpine\n\n# Run orchestrator\npython -m uvicorn orchestrator.main:app --reload\n```\n\n### Test WebSocket Connection\n\n```python\nimport asyncio\nimport websockets\nimport json\n\nasync def test_agent():\n uri = \"ws://localhost:8000/ws/test-agent-1\"\n async with websockets.connect(uri) as websocket:\n # Register\n await websocket.send(json.dumps({\n \ \"type\": \"register\",\n \"capabilities\": [\"python\"]\n \ }))\n response = await websocket.recv()\n print(f\"Registration: {response}\")\n \n # Send heartbeat\n while True:\n await websocket.send(json.dumps({\"type\": \"heartbeat\"}))\n response = await websocket.recv()\n print(f\"Heartbeat: {response}\")\n await asyncio.sleep(10)\n\nasyncio.run(test_agent())\n```\n\n## Acceptance Criteria Status\n\n✅ Orchestrator successfully registers new agent pods via WebSocket \n✅ Agent registry persists in Redis and survives orchestrator restart \n✅ Agent A can handoff task to Agent B via orchestrator \n✅ Orchestrator detects and deregisters crashed agents within 60s \n✅ Desktop client can query orchestrator API for agent status \n✅ System recovers from agent crashes without user intervention\n" agent_registry.py: | """Agent registry with Redis-backed state management.""" import json import time import logging from typing import Dict, List, Optional from enum import Enum from pydantic import BaseModel from .redis_client import redis_client logger = logging.getLogger(__name__) class AgentStatus(str, Enum): """Agent status enumeration.""" IDLE = "idle" BUSY = "busy" HANDOFF_PENDING = "handoff-pending" FAILED = "failed" class AgentMetadata(BaseModel): """Agent metadata model.""" agent_id: str status: AgentStatus last_heartbeat: float capabilities: List[str] current_task_id: Optional[str] = None class AgentRegistry: """Manages agent registration and heartbeat tracking.""" HEARTBEAT_TIMEOUT = 30 # seconds AGENT_KEY_PREFIX = "agent:" def __init__(self): pass async def register_agent( self, agent_id: str, capabilities: List[str] ) -> AgentMetadata: """Register a new agent.""" metadata = AgentMetadata( agent_id=agent_id, status=AgentStatus.IDLE, last_heartbeat=time.time(), capabilities=capabilities, ) key = f"{self.AGENT_KEY_PREFIX}{agent_id}" await redis_client.set(key, metadata.model_dump_json()) logger.info(f"Registered agent {agent_id} with capabilities: {capabilities}") return metadata async def deregister_agent(self, agent_id: str): """Deregister an agent.""" key = f"{self.AGENT_KEY_PREFIX}{agent_id}" await redis_client.delete(key) logger.info(f"Deregistered agent {agent_id}") async def update_heartbeat(self, agent_id: str) -> bool: """Update agent heartbeat timestamp.""" key = f"{self.AGENT_KEY_PREFIX}{agent_id}" data = await redis_client.get(key) if not data: logger.warning(f"Agent {agent_id} not found for heartbeat update") return False metadata = AgentMetadata.model_validate_json(data) metadata.last_heartbeat = time.time() await redis_client.set(key, metadata.model_dump_json()) return True async def update_status( self, agent_id: str, status: AgentStatus, task_id: Optional[str] = None ) -> bool: """Update agent status.""" key = f"{self.AGENT_KEY_PREFIX}{agent_id}" data = await redis_client.get(key) if not data: logger.warning(f"Agent {agent_id} not found for status update") return False metadata = AgentMetadata.model_validate_json(data) metadata.status = status metadata.current_task_id = task_id await redis_client.set(key, metadata.model_dump_json()) logger.info(f"Updated agent {agent_id} status to {status}") return True async def get_agent(self, agent_id: str) -> Optional[AgentMetadata]: """Get agent metadata.""" key = f"{self.AGENT_KEY_PREFIX}{agent_id}" data = await redis_client.get(key) if not data: return None return AgentMetadata.model_validate_json(data) async def get_all_agents(self) -> List[AgentMetadata]: """Get all registered agents.""" pattern = f"{self.AGENT_KEY_PREFIX}*" keys = await redis_client.keys(pattern) agents = [] for key in keys: data = await redis_client.get(key) if data: agents.append(AgentMetadata.model_validate_json(data)) return agents async def get_idle_agents(self) -> List[AgentMetadata]: """Get all idle agents.""" all_agents = await self.get_all_agents() return [agent for agent in all_agents if agent.status == AgentStatus.IDLE] async def check_failed_agents(self) -> List[str]: """Check for agents with expired heartbeats and mark as failed.""" current_time = time.time() failed_agents = [] all_agents = await self.get_all_agents() for agent in all_agents: if agent.status == AgentStatus.FAILED: continue time_since_heartbeat = current_time - agent.last_heartbeat if time_since_heartbeat > self.HEARTBEAT_TIMEOUT: await self.update_status(agent.agent_id, AgentStatus.FAILED) failed_agents.append(agent.agent_id) logger.warning( f"Agent {agent.agent_id} marked as failed " f"(no heartbeat for {time_since_heartbeat:.1f}s)" ) return failed_agents # Global agent registry instance agent_registry = AgentRegistry() checkpoint_manager.py: | """ Checkpoint Manager for Task Recovery Provides partial checkpointing and retry strategies for agent tasks. Enables recovery from failures without restarting entire workflows. """ import json import time from typing import Dict, List, Optional, Any from dataclasses import dataclass, asdict from datetime import datetime import redis import logging logger = logging.getLogger(__name__) @dataclass class Checkpoint: """Represents a task checkpoint""" task_id: str agent_id: str checkpoint_id: str timestamp: float phase: str # e.g., "analysis", "implementation", "testing" state: Dict[str, Any] # Serializable state data files_modified: List[str] git_commit: Optional[str] = None metadata: Optional[Dict[str, Any]] = None @dataclass class RetryPolicy: """Retry strategy configuration""" max_retries: int = 3 backoff_multiplier: float = 2.0 initial_delay_seconds: float = 1.0 max_delay_seconds: float = 60.0 retry_on_errors: List[str] = None # Error types to retry class CheckpointManager: """Manages task checkpoints and recovery""" def __init__(self, redis_client: redis.Redis): self.redis = redis_client self.checkpoint_ttl = 86400 * 7 # 7 days def save_checkpoint(self, checkpoint: Checkpoint) -> bool: """Save a checkpoint to Redis""" try: key = f"checkpoint:{checkpoint.task_id}:{checkpoint.checkpoint_id}" data = json.dumps(asdict(checkpoint)) self.redis.setex(key, self.checkpoint_ttl, data) # Add to task's checkpoint list list_key = f"checkpoints:{checkpoint.task_id}" self.redis.lpush(list_key, checkpoint.checkpoint_id) self.redis.expire(list_key, self.checkpoint_ttl) logger.info( f"Saved checkpoint {checkpoint.checkpoint_id} for task {checkpoint.task_id}" ) return True except Exception as e: logger.error(f"Failed to save checkpoint: {e}") return False def get_checkpoint(self, task_id: str, checkpoint_id: str) -> Optional[Checkpoint]: """Retrieve a specific checkpoint""" try: key = f"checkpoint:{task_id}:{checkpoint_id}" data = self.redis.get(key) if not data: return None checkpoint_dict = json.loads(data) return Checkpoint(**checkpoint_dict) except Exception as e: logger.error(f"Failed to retrieve checkpoint: {e}") return None def get_latest_checkpoint(self, task_id: str) -> Optional[Checkpoint]: """Get the most recent checkpoint for a task""" try: list_key = f"checkpoints:{task_id}" checkpoint_ids = self.redis.lrange(list_key, 0, 0) if not checkpoint_ids: return None checkpoint_id = checkpoint_ids[0].decode('utf-8') return self.get_checkpoint(task_id, checkpoint_id) except Exception as e: logger.error(f"Failed to get latest checkpoint: {e}") return None def list_checkpoints(self, task_id: str) -> List[str]: """List all checkpoint IDs for a task""" try: list_key = f"checkpoints:{task_id}" checkpoint_ids = self.redis.lrange(list_key, 0, -1) return [cid.decode('utf-8') for cid in checkpoint_ids] except Exception as e: logger.error(f"Failed to list checkpoints: {e}") return [] def delete_checkpoint(self, task_id: str, checkpoint_id: str) -> bool: """Delete a specific checkpoint""" try: key = f"checkpoint:{task_id}:{checkpoint_id}" self.redis.delete(key) # Remove from list list_key = f"checkpoints:{task_id}" self.redis.lrem(list_key, 0, checkpoint_id) logger.info(f"Deleted checkpoint {checkpoint_id} for task {task_id}") return True except Exception as e: logger.error(f"Failed to delete checkpoint: {e}") return False def cleanup_task_checkpoints(self, task_id: str) -> bool: """Delete all checkpoints for a task""" try: checkpoint_ids = self.list_checkpoints(task_id) for checkpoint_id in checkpoint_ids: self.delete_checkpoint(task_id, checkpoint_id) list_key = f"checkpoints:{task_id}" self.redis.delete(list_key) logger.info(f"Cleaned up all checkpoints for task {task_id}") return True except Exception as e: logger.error(f"Failed to cleanup checkpoints: {e}") return False class RetryManager: """Manages retry logic with exponential backoff""" def __init__(self, redis_client: redis.Redis): self.redis = redis_client def record_attempt(self, task_id: str, agent_id: str, error: Optional[str] = None): """Record a task attempt""" key = f"retry:{task_id}" attempt_data = { "agent_id": agent_id, "timestamp": time.time(), "error": error } self.redis.lpush(key, json.dumps(attempt_data)) self.redis.expire(key, 86400) # 24 hour TTL def get_attempt_count(self, task_id: str) -> int: """Get number of attempts for a task""" key = f"retry:{task_id}" return self.redis.llen(key) def should_retry(self, task_id: str, policy: RetryPolicy) -> bool: """Determine if task should be retried""" attempt_count = self.get_attempt_count(task_id) return attempt_count < policy.max_retries def get_retry_delay(self, task_id: str, policy: RetryPolicy) -> float: """Calculate delay before next retry (exponential backoff)""" attempt_count = self.get_attempt_count(task_id) delay = policy.initial_delay_seconds * (policy.backoff_multiplier ** attempt_count) return min(delay, policy.max_delay_seconds) def clear_attempts(self, task_id: str): """Clear retry history for a task""" key = f"retry:{task_id}" self.redis.delete(key) class RecoveryCoordinator: """Coordinates task recovery from checkpoints""" def __init__(self, checkpoint_manager: CheckpointManager, retry_manager: RetryManager): self.checkpoint_manager = checkpoint_manager self.retry_manager = retry_manager def recover_task( self, task_id: str, retry_policy: Optional[RetryPolicy] = None ) -> Optional[Dict[str, Any]]: """ Attempt to recover a failed task from its latest checkpoint Returns recovery instructions or None if recovery not possible """ if retry_policy is None: retry_policy = RetryPolicy() # Check if we should retry if not self.retry_manager.should_retry(task_id, retry_policy): logger.warning(f"Task {task_id} exceeded max retries") return None # Get latest checkpoint checkpoint = self.checkpoint_manager.get_latest_checkpoint(task_id) if not checkpoint: logger.warning(f"No checkpoint found for task {task_id}") return None # Calculate retry delay delay = self.retry_manager.get_retry_delay(task_id, retry_policy) recovery_plan = { "task_id": task_id, "checkpoint_id": checkpoint.checkpoint_id, "resume_phase": checkpoint.phase, "state": checkpoint.state, "files_modified": checkpoint.files_modified, "git_commit": checkpoint.git_commit, "retry_delay_seconds": delay, "attempt_number": self.retry_manager.get_attempt_count(task_id) + 1 } logger.info( f"Recovery plan created for task {task_id} from checkpoint {checkpoint.checkpoint_id}" ) return recovery_plan def mark_recovery_success(self, task_id: str): """Mark a task as successfully recovered""" self.retry_manager.clear_attempts(task_id) logger.info(f"Task {task_id} recovered successfully") def mark_recovery_failure(self, task_id: str, agent_id: str, error: str): """Record a failed recovery attempt""" self.retry_manager.record_attempt(task_id, agent_id, error) logger.warning(f"Recovery attempt failed for task {task_id}: {error}") handoff_manager.py: | """Handoff manager for task coordination between agents.""" import json import time import uuid import logging from typing import Dict, Optional from pydantic import BaseModel from .redis_client import redis_client from .agent_registry import agent_registry, AgentStatus logger = logging.getLogger(__name__) class HandoffRequest(BaseModel): """Handoff request model.""" type: str = "handoff" source_agent_id: str target_agent_id: str task_context: Dict class HandoffRecord(BaseModel): """Handoff history record.""" handoff_id: str source_agent_id: str target_agent_id: str timestamp: float task_context: Dict status: str # pending, completed, failed class HandoffManager: """Manages task handoffs between agents.""" HANDOFF_KEY_PREFIX = "handoff:" def __init__(self): pass async def initiate_handoff( self, source_agent_id: str, target_agent_id: str, task_context: Dict ) -> Optional[str]: """Initiate a handoff from source to target agent.""" # Validate source agent exists and is busy source_agent = await agent_registry.get_agent(source_agent_id) if not source_agent: logger.error(f"Source agent {source_agent_id} not found") return None # Validate target agent exists and is idle target_agent = await agent_registry.get_agent(target_agent_id) if not target_agent: logger.error(f"Target agent {target_agent_id} not found") return None if target_agent.status != AgentStatus.IDLE: logger.error( f"Target agent {target_agent_id} is not idle (status: {target_agent.status})" ) return None # Create handoff record handoff_id = str(uuid.uuid4()) record = HandoffRecord( handoff_id=handoff_id, source_agent_id=source_agent_id, target_agent_id=target_agent_id, timestamp=time.time(), task_context=task_context, status="pending", ) # Store in Redis key = f"{self.HANDOFF_KEY_PREFIX}{handoff_id}" await redis_client.set(key, record.model_dump_json()) # Update agent statuses await agent_registry.update_status(source_agent_id, AgentStatus.IDLE) await agent_registry.update_status( target_agent_id, AgentStatus.HANDOFF_PENDING, task_context.get("task_id") ) logger.info( f"Handoff {handoff_id} initiated: {source_agent_id} -> {target_agent_id}" ) return handoff_id async def complete_handoff(self, handoff_id: str) -> bool: """Mark handoff as completed.""" key = f"{self.HANDOFF_KEY_PREFIX}{handoff_id}" data = await redis_client.get(key) if not data: logger.error(f"Handoff {handoff_id} not found") return False record = HandoffRecord.model_validate_json(data) record.status = "completed" await redis_client.set(key, record.model_dump_json()) # Update target agent to busy await agent_registry.update_status( record.target_agent_id, AgentStatus.BUSY, record.task_context.get("task_id") ) logger.info(f"Handoff {handoff_id} completed") return True async def fail_handoff(self, handoff_id: str, reason: str) -> bool: """Mark handoff as failed.""" key = f"{self.HANDOFF_KEY_PREFIX}{handoff_id}" data = await redis_client.get(key) if not data: logger.error(f"Handoff {handoff_id} not found") return False record = HandoffRecord.model_validate_json(data) record.status = "failed" await redis_client.set(key, record.model_dump_json()) # Revert target agent to idle await agent_registry.update_status(record.target_agent_id, AgentStatus.IDLE) logger.warning(f"Handoff {handoff_id} failed: {reason}") return True async def get_handoff(self, handoff_id: str) -> Optional[HandoffRecord]: """Get handoff record by ID.""" key = f"{self.HANDOFF_KEY_PREFIX}{handoff_id}" data = await redis_client.get(key) if not data: return None return HandoffRecord.model_validate_json(data) async def get_handoff_history( self, agent_id: Optional[str] = None, limit: int = 100 ) -> list[HandoffRecord]: """Get handoff history, optionally filtered by agent.""" pattern = f"{self.HANDOFF_KEY_PREFIX}*" keys = await redis_client.keys(pattern) records = [] for key in keys[:limit]: data = await redis_client.get(key) if data: record = HandoffRecord.model_validate_json(data) if agent_id is None or ( record.source_agent_id == agent_id or record.target_agent_id == agent_id ): records.append(record) # Sort by timestamp descending records.sort(key=lambda r: r.timestamp, reverse=True) return records[:limit] # Global handoff manager instance handoff_manager = HandoffManager() main.py: | """FastAPI orchestrator with WebSocket support for agent coordination.""" import asyncio import logging import time from contextlib import asynccontextmanager from typing import Dict, List, Optional from fastapi import FastAPI, WebSocket, WebSocketDisconnect, HTTPException from fastapi.responses import JSONResponse, PlainTextResponse from pydantic import BaseModel from prometheus_client import Counter, Gauge, Histogram, generate_latest, CONTENT_TYPE_LATEST from .redis_client import redis_client from .agent_registry import agent_registry, AgentStatus, AgentMetadata from .handoff_manager import handoff_manager, HandoffRequest from .task_queue import task_queue, TaskStatus # Configure logging logging.basicConfig( level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s" ) logger = logging.getLogger(__name__) # Prometheus metrics AGENTS_CREATED = Counter('swarm_agents_created_total', 'Total agents created') AGENTS_FAILED = Counter('swarm_agents_failed_total', 'Total agents failed') AGENTS_ACTIVE = Gauge('swarm_agents_active', 'Currently active agents') TASKS_CREATED = Counter('swarm_tasks_created_total', 'Total tasks created') TASKS_COMPLETED = Counter('swarm_tasks_completed_total', 'Total tasks completed') TASKS_FAILED = Counter('swarm_tasks_failed_total', 'Total tasks failed') HANDOFFS_TOTAL = Counter('swarm_handoffs_total', 'Total handoffs performed') HANDOFF_DURATION = Histogram('swarm_handoff_duration_seconds', 'Handoff latency in seconds') TASK_DURATION = Histogram('swarm_task_duration_seconds', 'Task duration in seconds') AGENT_STATUS = Gauge('swarm_agent_status', 'Agent status by ID', ['agent_id', 'status']) WEBSOCKET_CONNECTIONS = Gauge('swarm_websocket_connections', 'Active WebSocket connections') ERRORS_TOTAL = Counter('swarm_errors_total', 'Total errors', ['error_type']) # WebSocket connection manager class ConnectionManager: """Manages WebSocket connections for agents.""" def __init__(self): self.active_connections: Dict[str, WebSocket] = {} async def connect(self, agent_id: str, websocket: WebSocket): """Accept and store WebSocket connection.""" await websocket.accept() self.active_connections[agent_id] = websocket logger.info(f"Agent {agent_id} connected via WebSocket") def disconnect(self, agent_id: str): """Remove WebSocket connection.""" if agent_id in self.active_connections: del self.active_connections[agent_id] logger.info(f"Agent {agent_id} disconnected") async def send_message(self, agent_id: str, message: dict): """Send message to specific agent.""" if agent_id in self.active_connections: try: await self.active_connections[agent_id].send_json(message) except Exception as e: logger.error(f"Failed to send message to agent {agent_id}: {e}") async def broadcast(self, message: dict): """Broadcast message to all connected agents.""" for agent_id, connection in self.active_connections.items(): try: await connection.send_json(message) except Exception as e: logger.error(f"Failed to broadcast to agent {agent_id}: {e}") manager = ConnectionManager() # Background task for failure detection async def failure_detection_loop(): """Periodically check for failed agents and reassign their tasks.""" while True: try: await asyncio.sleep(10) # Check every 10 seconds failed_agents = await agent_registry.check_failed_agents() for agent_id in failed_agents: # Reassign tasks from failed agent reassigned_tasks = await task_queue.reassign_agent_tasks(agent_id) if reassigned_tasks: logger.info( f"Reassigned {len(reassigned_tasks)} tasks from failed agent {agent_id}" ) # Disconnect WebSocket manager.disconnect(agent_id) except Exception as e: logger.error(f"Error in failure detection loop: {e}") # Lifespan context manager @asynccontextmanager async def lifespan(app: FastAPI): """Startup and shutdown events.""" # Startup logger.info("Starting orchestrator...") await redis_client.connect() # Start background tasks failure_task = asyncio.create_task(failure_detection_loop()) yield # Shutdown logger.info("Shutting down orchestrator...") failure_task.cancel() await redis_client.disconnect() # Create FastAPI app app = FastAPI(title="Swarm Orchestrator", version="1.0.0", lifespan=lifespan) # Pydantic models for API class RegisterRequest(BaseModel): """Agent registration request.""" agent_id: str capabilities: List[str] class TaskCreateRequest(BaseModel): """Task creation request.""" description: str context: Dict = {} max_retries: int = 3 class TaskAssignRequest(BaseModel): """Task assignment request.""" task_id: str agent_id: str # REST API endpoints @app.get("/") async def root(): """Health check endpoint.""" return {"status": "ok", "service": "swarm-orchestrator"} @app.get("/health") async def health(): """Detailed health check.""" try: await redis_client.client.ping() redis_status = "connected" except Exception: redis_status = "disconnected" return { "status": "ok", "redis": redis_status, "active_connections": len(manager.active_connections) } @app.get("/agents") async def list_agents(): """List all registered agents.""" agents = await agent_registry.get_all_agents() return {"agents": [agent.model_dump() for agent in agents]} @app.get("/agents/{agent_id}") async def get_agent(agent_id: str): """Get specific agent details.""" agent = await agent_registry.get_agent(agent_id) if not agent: raise HTTPException(status_code=404, detail="Agent not found") return agent.model_dump() @app.get("/agents/idle") async def list_idle_agents(): """List all idle agents.""" agents = await agent_registry.get_idle_agents() return {"agents": [agent.model_dump() for agent in agents]} @app.post("/tasks") async def create_task(request: TaskCreateRequest): """Create a new task.""" task = await task_queue.create_task( description=request.description, context=request.context, max_retries=request.max_retries ) TASKS_CREATED.inc() return task.model_dump() @app.post("/tasks/assign") async def assign_task(request: TaskAssignRequest): """Manually assign a task to an agent.""" success = await task_queue.assign_task(request.task_id, request.agent_id) if not success: raise HTTPException(status_code=400, detail="Failed to assign task") # Notify agent via WebSocket await manager.send_message(request.agent_id, { "type": "task_assigned", "task_id": request.task_id }) return {"status": "assigned", "task_id": request.task_id, "agent_id": request.agent_id} @app.get("/tasks") async def list_tasks(status: Optional[str] = None): """List all tasks, optionally filtered by status.""" task_status = TaskStatus(status) if status else None tasks = await task_queue.get_all_tasks(status=task_status) return {"tasks": [task.model_dump() for task in tasks]} @app.get("/tasks/{task_id}") async def get_task(task_id: str): """Get specific task details.""" task = await task_queue.get_task(task_id) if not task: raise HTTPException(status_code=404, detail="Task not found") return task.model_dump() @app.get("/handoffs") async def list_handoffs(agent_id: Optional[str] = None, limit: int = 100): """List handoff history.""" handoffs = await handoff_manager.get_handoff_history(agent_id=agent_id, limit=limit) return {"handoffs": [h.model_dump() for h in handoffs]} @app.get("/metrics") async def metrics(): """Prometheus metrics endpoint.""" # Update gauge metrics before returning agents = await agent_registry.get_all_agents() AGENTS_ACTIVE.set(len([a for a in agents if a.status == AgentStatus.BUSY])) WEBSOCKET_CONNECTIONS.set(len(manager.active_connections)) # Update per-agent status for agent in agents: AGENT_STATUS.labels(agent_id=agent.agent_id, status=agent.status.value).set(1) return PlainTextResponse(generate_latest(), media_type=CONTENT_TYPE_LATEST) # WebSocket endpoint @app.websocket("/ws/{agent_id}") async def websocket_endpoint(websocket: WebSocket, agent_id: str): """WebSocket endpoint for agent communication.""" await manager.connect(agent_id, websocket) try: # Wait for registration message data = await websocket.receive_json() if data.get("type") == "register": capabilities = data.get("capabilities", []) await agent_registry.register_agent(agent_id, capabilities) await websocket.send_json({ "type": "registered", "agent_id": agent_id, "status": "success" }) # Main message loop while True: try: message = await asyncio.wait_for( websocket.receive_json(), timeout=15.0 # Expect heartbeat every 15s ) message_type = message.get("type") if message_type == "heartbeat": await agent_registry.update_heartbeat(agent_id) await websocket.send_json({"type": "heartbeat_ack"}) elif message_type == "handoff": # Handle handoff request target_agent_id = message.get("target_agent_id") task_context = message.get("task_context", {}) handoff_id = await handoff_manager.initiate_handoff( agent_id, target_agent_id, task_context ) if handoff_id: # Notify target agent await manager.send_message(target_agent_id, { "type": "handoff_request", "handoff_id": handoff_id, "source_agent_id": agent_id, "task_context": task_context }) await websocket.send_json({ "type": "handoff_initiated", "handoff_id": handoff_id, "status": "success" }) else: await websocket.send_json({ "type": "handoff_failed", "reason": "Target agent not available" }) elif message_type == "handoff_accept": # Target agent accepts handoff handoff_id = message.get("handoff_id") await handoff_manager.complete_handoff(handoff_id) await websocket.send_json({ "type": "handoff_accepted", "handoff_id": handoff_id }) elif message_type == "handoff_reject": # Target agent rejects handoff handoff_id = message.get("handoff_id") reason = message.get("reason", "Rejected by target agent") await handoff_manager.fail_handoff(handoff_id, reason) await websocket.send_json({ "type": "handoff_rejected", "handoff_id": handoff_id }) elif message_type == "task_start": # Agent starts working on task task_id = message.get("task_id") await task_queue.start_task(task_id) elif message_type == "task_complete": # Agent completes task task_id = message.get("task_id") await task_queue.complete_task(task_id) await websocket.send_json({ "type": "task_completed", "task_id": task_id }) elif message_type == "task_failed": # Agent reports task failure task_id = message.get("task_id") reason = message.get("reason", "Unknown error") await task_queue.fail_task(task_id, reason) await websocket.send_json({ "type": "task_failed_ack", "task_id": task_id }) elif message_type == "status_update": # Agent updates its status status = AgentStatus(message.get("status")) task_id = message.get("task_id") await agent_registry.update_status(agent_id, status, task_id) else: logger.warning(f"Unknown message type from {agent_id}: {message_type}") except asyncio.TimeoutError: # No message received within timeout - check heartbeat agent = await agent_registry.get_agent(agent_id) if agent: import time if time.time() - agent.last_heartbeat > 30: logger.warning(f"Agent {agent_id} heartbeat timeout") break continue else: logger.error(f"Agent {agent_id} did not send registration message") await websocket.close(code=1008, reason="Registration required") except WebSocketDisconnect: logger.info(f"Agent {agent_id} disconnected") except Exception as e: logger.error(f"Error in WebSocket handler for agent {agent_id}: {e}") finally: # Cleanup manager.disconnect(agent_id) await agent_registry.deregister_agent(agent_id) if __name__ == "__main__": import uvicorn uvicorn.run(app, host="0.0.0.0", port=8000) model_tuner.py: | """ Complexity Model Tuner Implements feedback loop and retraining for the complexity analysis model. Collects actual vs predicted agent counts and adjusts the model over time. """ import json import time from typing import Dict, List, Optional, Tuple from dataclasses import dataclass, asdict from datetime import datetime, timedelta import redis import logging from collections import defaultdict import statistics logger = logging.getLogger(__name__) @dataclass class ComplexityPrediction: """Record of a complexity prediction""" task_id: str timestamp: float task_description: str predicted_agents: int predicted_complexity: str # "simple", "moderate", "complex" confidence: float model_version: str @dataclass class ComplexityActual: """Actual outcome of a task""" task_id: str timestamp: float actual_agents: int actual_duration_seconds: float success: bool user_feedback: Optional[str] = None # "too_many", "too_few", "just_right" user_rating: Optional[int] = None # 1-5 scale @dataclass class ModelMetrics: """Model performance metrics""" model_version: str total_predictions: int mean_absolute_error: float accuracy_within_1: float # % predictions within ±1 agent accuracy_within_2: float # % predictions within ±2 agents user_satisfaction: float # Average user rating last_updated: float class ComplexityModelTuner: """Manages complexity model feedback and tuning""" def __init__(self, redis_client: redis.Redis): self.redis = redis_client self.current_model_version = "v1.0" self.data_ttl = 86400 * 30 # 30 days def record_prediction(self, prediction: ComplexityPrediction) -> bool: """Record a complexity prediction""" try: key = f"prediction:{prediction.task_id}" data = json.dumps(asdict(prediction)) self.redis.setex(key, self.data_ttl, data) # Add to predictions list list_key = f"predictions:{prediction.model_version}" self.redis.lpush(list_key, prediction.task_id) self.redis.expire(list_key, self.data_ttl) logger.info( f"Recorded prediction for task {prediction.task_id}: " f"{prediction.predicted_agents} agents" ) return True except Exception as e: logger.error(f"Failed to record prediction: {e}") return False def record_actual(self, actual: ComplexityActual) -> bool: """Record actual task outcome""" try: key = f"actual:{actual.task_id}" data = json.dumps(asdict(actual)) self.redis.setex(key, self.data_ttl, data) # Add to actuals list list_key = "actuals:all" self.redis.lpush(list_key, actual.task_id) self.redis.expire(list_key, self.data_ttl) logger.info( f"Recorded actual for task {actual.task_id}: " f"{actual.actual_agents} agents, feedback: {actual.user_feedback}" ) return True except Exception as e: logger.error(f"Failed to record actual: {e}") return False def get_prediction(self, task_id: str) -> Optional[ComplexityPrediction]: """Retrieve a prediction""" try: key = f"prediction:{task_id}" data = self.redis.get(key) if not data: return None return ComplexityPrediction(**json.loads(data)) except Exception as e: logger.error(f"Failed to get prediction: {e}") return None def get_actual(self, task_id: str) -> Optional[ComplexityActual]: """Retrieve actual outcome""" try: key = f"actual:{task_id}" data = self.redis.get(key) if not data: return None return ComplexityActual(**json.loads(data)) except Exception as e: logger.error(f"Failed to get actual: {e}") return None def calculate_metrics( self, model_version: Optional[str] = None, limit: int = 1000 ) -> Optional[ModelMetrics]: """Calculate model performance metrics""" if model_version is None: model_version = self.current_model_version try: # Get predictions for this model version list_key = f"predictions:{model_version}" task_ids = self.redis.lrange(list_key, 0, limit - 1) if not task_ids: logger.warning(f"No predictions found for model {model_version}") return None errors = [] within_1 = 0 within_2 = 0 ratings = [] for task_id_bytes in task_ids: task_id = task_id_bytes.decode('utf-8') prediction = self.get_prediction(task_id) actual = self.get_actual(task_id) if not prediction or not actual: continue # Calculate error error = abs(prediction.predicted_agents - actual.actual_agents) errors.append(error) # Check accuracy thresholds if error <= 1: within_1 += 1 if error <= 2: within_2 += 1 # Collect user ratings if actual.user_rating: ratings.append(actual.user_rating) if not errors: logger.warning(f"No matched predictions/actuals for model {model_version}") return None total = len(errors) mae = statistics.mean(errors) acc_1 = (within_1 / total) * 100 acc_2 = (within_2 / total) * 100 avg_rating = statistics.mean(ratings) if ratings else 0.0 metrics = ModelMetrics( model_version=model_version, total_predictions=total, mean_absolute_error=mae, accuracy_within_1=acc_1, accuracy_within_2=acc_2, user_satisfaction=avg_rating, last_updated=time.time() ) # Cache metrics metrics_key = f"metrics:{model_version}" self.redis.setex(metrics_key, 3600, json.dumps(asdict(metrics))) logger.info( f"Model {model_version} metrics: MAE={mae:.2f}, " f"Acc±1={acc_1:.1f}%, Acc±2={acc_2:.1f}%, " f"Satisfaction={avg_rating:.2f}/5" ) return metrics except Exception as e: logger.error(f"Failed to calculate metrics: {e}") return None def get_error_patterns(self, limit: int = 100) -> Dict[str, List[Tuple[str, int, int]]]: """ Analyze error patterns to identify systematic biases Returns dict with categories: - overestimated: tasks where we predicted too many agents - underestimated: tasks where we predicted too few agents - accurate: tasks where prediction was close """ patterns = { "overestimated": [], "underestimated": [], "accurate": [] } try: list_key = "actuals:all" task_ids = self.redis.lrange(list_key, 0, limit - 1) for task_id_bytes in task_ids: task_id = task_id_bytes.decode('utf-8') prediction = self.get_prediction(task_id) actual = self.get_actual(task_id) if not prediction or not actual: continue error = prediction.predicted_agents - actual.actual_agents entry = ( task_id, prediction.predicted_agents, actual.actual_agents ) if error > 1: patterns["overestimated"].append(entry) elif error < -1: patterns["underestimated"].append(entry) else: patterns["accurate"].append(entry) return patterns except Exception as e: logger.error(f"Failed to analyze error patterns: {e}") return patterns def get_feedback_summary(self, limit: int = 100) -> Dict[str, int]: """Summarize user feedback""" feedback_counts = defaultdict(int) try: list_key = "actuals:all" task_ids = self.redis.lrange(list_key, 0, limit - 1) for task_id_bytes in task_ids: task_id = task_id_bytes.decode('utf-8') actual = self.get_actual(task_id) if actual and actual.user_feedback: feedback_counts[actual.user_feedback] += 1 return dict(feedback_counts) except Exception as e: logger.error(f"Failed to get feedback summary: {e}") return {} def generate_tuning_recommendations(self) -> List[str]: """Generate recommendations for model tuning based on data""" recommendations = [] try: # Get current metrics metrics = self.calculate_metrics() if not metrics: return ["Insufficient data for recommendations"] # Check accuracy if metrics.accuracy_within_1 < 60: recommendations.append( f"Low accuracy ({metrics.accuracy_within_1:.1f}%). " "Consider retraining with more diverse examples." ) # Check error patterns patterns = self.get_error_patterns() overestimated = len(patterns["overestimated"]) underestimated = len(patterns["underestimated"]) total = overestimated + underestimated + len(patterns["accurate"]) if total > 0: over_pct = (overestimated / total) * 100 under_pct = (underestimated / total) * 100 if over_pct > 40: recommendations.append( f"Model overestimates in {over_pct:.1f}% of cases. " "Consider reducing base agent count or adjusting complexity thresholds." ) if under_pct > 40: recommendations.append( f"Model underestimates in {under_pct:.1f}% of cases. " "Consider increasing base agent count or lowering complexity thresholds." ) # Check user satisfaction if metrics.user_satisfaction < 3.5: recommendations.append( f"Low user satisfaction ({metrics.user_satisfaction:.1f}/5). " "Review user feedback and adjust model accordingly." ) # Check feedback feedback = self.get_feedback_summary() if feedback.get("too_many", 0) > feedback.get("too_few", 0) * 2: recommendations.append( "Users frequently report 'too many agents'. " "Consider reducing default agent counts." ) elif feedback.get("too_few", 0) > feedback.get("too_many", 0) * 2: recommendations.append( "Users frequently report 'too few agents'. " "Consider increasing default agent counts." ) if not recommendations: recommendations.append( f"Model performing well (MAE={metrics.mean_absolute_error:.2f}, " f"Acc±1={metrics.accuracy_within_1:.1f}%). Continue monitoring." ) except Exception as e: logger.error(f"Failed to generate recommendations: {e}") recommendations.append(f"Error generating recommendations: {e}") return recommendations def export_training_data(self, limit: int = 1000) -> List[Dict]: """Export prediction/actual pairs for model retraining""" training_data = [] try: list_key = "actuals:all" task_ids = self.redis.lrange(list_key, 0, limit - 1) for task_id_bytes in task_ids: task_id = task_id_bytes.decode('utf-8') prediction = self.get_prediction(task_id) actual = self.get_actual(task_id) if not prediction or not actual: continue training_data.append({ "task_description": prediction.task_description, "predicted_agents": prediction.predicted_agents, "actual_agents": actual.actual_agents, "duration_seconds": actual.actual_duration_seconds, "success": actual.success, "user_feedback": actual.user_feedback, "user_rating": actual.user_rating }) logger.info(f"Exported {len(training_data)} training examples") return training_data except Exception as e: logger.error(f"Failed to export training data: {e}") return [] redis_client.py: | """Redis client for orchestrator state management.""" import os import redis.asyncio as redis from typing import Optional import logging logger = logging.getLogger(__name__) class RedisClient: """Async Redis client wrapper for orchestrator operations.""" def __init__(self): self.client: Optional[redis.Redis] = None self.host = os.getenv("REDIS_HOST", "redis-service") self.port = int(os.getenv("REDIS_PORT", "6379")) self.db = int(os.getenv("REDIS_DB", "0")) async def connect(self): """Establish Redis connection.""" try: self.client = await redis.Redis( host=self.host, port=self.port, db=self.db, decode_responses=True, socket_connect_timeout=5, socket_keepalive=True, ) await self.client.ping() logger.info(f"Connected to Redis at {self.host}:{self.port}") except Exception as e: logger.error(f"Failed to connect to Redis: {e}") raise async def disconnect(self): """Close Redis connection.""" if self.client: await self.client.close() logger.info("Disconnected from Redis") async def set(self, key: str, value: str, ex: Optional[int] = None): """Set key-value pair with optional expiration.""" await self.client.set(key, value, ex=ex) async def get(self, key: str) -> Optional[str]: """Get value by key.""" return await self.client.get(key) async def delete(self, key: str): """Delete key.""" await self.client.delete(key) async def exists(self, key: str) -> bool: """Check if key exists.""" return await self.client.exists(key) > 0 async def hset(self, name: str, key: str, value: str): """Set hash field.""" await self.client.hset(name, key, value) async def hget(self, name: str, key: str) -> Optional[str]: """Get hash field.""" return await self.client.hget(name, key) async def hgetall(self, name: str) -> dict: """Get all hash fields.""" return await self.client.hgetall(name) async def hdel(self, name: str, *keys: str): """Delete hash fields.""" await self.client.hdel(name, *keys) async def keys(self, pattern: str) -> list: """Get keys matching pattern.""" return await self.client.keys(pattern) async def lpush(self, key: str, *values: str): """Push values to list head.""" await self.client.lpush(key, *values) async def rpop(self, key: str) -> Optional[str]: """Pop value from list tail.""" return await self.client.rpop(key) async def llen(self, key: str) -> int: """Get list length.""" return await self.client.llen(key) async def lrange(self, key: str, start: int, end: int) -> list: """Get list range.""" return await self.client.lrange(key, start, end) # Global Redis client instance redis_client = RedisClient() requirements.txt: | fastapi==0.115.0 uvicorn[standard]==0.32.0 websockets==13.1 redis==5.2.0 pydantic==2.9.2 python-dotenv==1.0.1 prometheus-client==0.20.0 opentelemetry-api==1.24.0 opentelemetry-sdk==1.24.0 opentelemetry-exporter-otlp==1.24.0 opentelemetry-instrumentation-redis==0.45b0 opentelemetry-instrumentation-requests==0.45b0 opentelemetry-instrumentation-logging==0.45b0 task_queue.py: | """Task queue management with failure recovery.""" import json import time import uuid import logging from typing import Dict, List, Optional from enum import Enum from pydantic import BaseModel from .redis_client import redis_client from .agent_registry import agent_registry, AgentStatus logger = logging.getLogger(__name__) class TaskStatus(str, Enum): """Task status enumeration.""" PENDING = "pending" ASSIGNED = "assigned" IN_PROGRESS = "in_progress" COMPLETED = "completed" FAILED = "failed" class Task(BaseModel): """Task model.""" task_id: str description: str status: TaskStatus assigned_agent_id: Optional[str] = None created_at: float started_at: Optional[float] = None completed_at: Optional[float] = None retry_count: int = 0 max_retries: int = 3 context: Dict = {} class TaskQueue: """Manages task assignment and reassignment with failure recovery.""" TASK_KEY_PREFIX = "task:" PENDING_QUEUE_KEY = "queue:pending" AGENT_TASK_KEY_PREFIX = "agent_task:" def __init__(self): pass async def create_task( self, description: str, context: Dict = {}, max_retries: int = 3 ) -> Task: """Create a new task and add to pending queue.""" task = Task( task_id=str(uuid.uuid4()), description=description, status=TaskStatus.PENDING, created_at=time.time(), max_retries=max_retries, context=context, ) # Store task key = f"{self.TASK_KEY_PREFIX}{task.task_id}" await redis_client.set(key, task.model_dump_json()) # Add to pending queue await redis_client.lpush(self.PENDING_QUEUE_KEY, task.task_id) logger.info(f"Created task {task.task_id}: {description}") return task async def assign_task(self, task_id: str, agent_id: str) -> bool: """Assign a task to an agent.""" # Get task task = await self.get_task(task_id) if not task: logger.error(f"Task {task_id} not found") return False if task.status not in [TaskStatus.PENDING, TaskStatus.FAILED]: logger.error(f"Task {task_id} cannot be assigned (status: {task.status})") return False # Verify agent is idle agent = await agent_registry.get_agent(agent_id) if not agent or agent.status != AgentStatus.IDLE: logger.error(f"Agent {agent_id} is not available for task assignment") return False # Update task task.status = TaskStatus.ASSIGNED task.assigned_agent_id = agent_id task.started_at = time.time() key = f"{self.TASK_KEY_PREFIX}{task_id}" await redis_client.set(key, task.model_dump_json()) # Track agent's current task agent_task_key = f"{self.AGENT_TASK_KEY_PREFIX}{agent_id}" await redis_client.set(agent_task_key, task_id) # Update agent status await agent_registry.update_status(agent_id, AgentStatus.BUSY, task_id) logger.info(f"Assigned task {task_id} to agent {agent_id}") return True async def start_task(self, task_id: str) -> bool: """Mark task as in progress.""" task = await self.get_task(task_id) if not task: return False task.status = TaskStatus.IN_PROGRESS key = f"{self.TASK_KEY_PREFIX}{task_id}" await redis_client.set(key, task.model_dump_json()) logger.info(f"Task {task_id} started") return True async def complete_task(self, task_id: str) -> bool: """Mark task as completed.""" task = await self.get_task(task_id) if not task: return False task.status = TaskStatus.COMPLETED task.completed_at = time.time() key = f"{self.TASK_KEY_PREFIX}{task_id}" await redis_client.set(key, task.model_dump_json()) # Clear agent's current task if task.assigned_agent_id: agent_task_key = f"{self.AGENT_TASK_KEY_PREFIX}{task.assigned_agent_id}" await redis_client.delete(agent_task_key) # Update agent to idle await agent_registry.update_status(task.assigned_agent_id, AgentStatus.IDLE) logger.info(f"Task {task_id} completed") return True async def fail_task(self, task_id: str, reason: str = "") -> bool: """Mark task as failed and handle retry logic.""" task = await self.get_task(task_id) if not task: return False task.retry_count += 1 # Clear agent's current task if task.assigned_agent_id: agent_task_key = f"{self.AGENT_TASK_KEY_PREFIX}{task.assigned_agent_id}" await redis_client.delete(agent_task_key) # Update agent to idle await agent_registry.update_status(task.assigned_agent_id, AgentStatus.IDLE) # Check if we should retry if task.retry_count < task.max_retries: task.status = TaskStatus.PENDING task.assigned_agent_id = None task.started_at = None # Re-add to pending queue await redis_client.lpush(self.PENDING_QUEUE_KEY, task_id) logger.warning( f"Task {task_id} failed (retry {task.retry_count}/{task.max_retries}): {reason}" ) else: task.status = TaskStatus.FAILED task.completed_at = time.time() logger.error( f"Task {task_id} permanently failed after {task.retry_count} retries: {reason}" ) key = f"{self.TASK_KEY_PREFIX}{task_id}" await redis_client.set(key, task.model_dump_json()) return True async def reassign_agent_tasks(self, failed_agent_id: str) -> List[str]: """Reassign all tasks from a failed agent back to pending queue.""" # Get agent's current task agent_task_key = f"{self.AGENT_TASK_KEY_PREFIX}{failed_agent_id}" task_id = await redis_client.get(agent_task_key) reassigned_tasks = [] if task_id: await self.fail_task(task_id, f"Agent {failed_agent_id} failed") reassigned_tasks.append(task_id) logger.info( f"Reassigned {len(reassigned_tasks)} tasks from failed agent {failed_agent_id}" ) return reassigned_tasks async def get_task(self, task_id: str) -> Optional[Task]: """Get task by ID.""" key = f"{self.TASK_KEY_PREFIX}{task_id}" data = await redis_client.get(key) if not data: return None return Task.model_validate_json(data) async def get_next_pending_task(self) -> Optional[Task]: """Get next pending task from queue.""" task_id = await redis_client.rpop(self.PENDING_QUEUE_KEY) if not task_id: return None return await self.get_task(task_id) async def get_pending_count(self) -> int: """Get count of pending tasks.""" return await redis_client.llen(self.PENDING_QUEUE_KEY) async def get_all_tasks(self, status: Optional[TaskStatus] = None) -> List[Task]: """Get all tasks, optionally filtered by status.""" pattern = f"{self.TASK_KEY_PREFIX}*" keys = await redis_client.keys(pattern) tasks = [] for key in keys: data = await redis_client.get(key) if data: task = Task.model_validate_json(data) if status is None or task.status == status: tasks.append(task) return tasks # Global task queue instance task_queue = TaskQueue() tracing.py: | """ OpenTelemetry Distributed Tracing Setup Provides distributed tracing for agent handoffs and task flows. Enables end-to-end visibility across the swarm system. """ import os import logging from typing import Optional, Dict, Any from contextlib import contextmanager from opentelemetry import trace from opentelemetry.sdk.trace import TracerProvider from opentelemetry.sdk.trace.export import BatchSpanProcessor, ConsoleSpanExporter from opentelemetry.sdk.resources import Resource, SERVICE_NAME, SERVICE_VERSION from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter from opentelemetry.instrumentation.redis import RedisInstrumentor from opentelemetry.instrumentation.requests import RequestsInstrumentor from opentelemetry.instrumentation.logging import LoggingInstrumentor from opentelemetry.trace import Status, StatusCode, SpanKind logger = logging.getLogger(__name__) class SwarmTracer: """Manages distributed tracing for the swarm system""" def __init__( self, service_name: str, service_version: str = "1.0.0", otlp_endpoint: Optional[str] = None, enable_console: bool = False ): self.service_name = service_name self.service_version = service_version self.otlp_endpoint = otlp_endpoint or os.getenv( "OTEL_EXPORTER_OTLP_ENDPOINT", "http://localhost:4317" ) self.enable_console = enable_console self._setup_tracing() def _setup_tracing(self): """Initialize OpenTelemetry tracing""" # Create resource with service information resource = Resource.create({ SERVICE_NAME: self.service_name, SERVICE_VERSION: self.service_version, "deployment.environment": os.getenv("ENVIRONMENT", "production"), "k8s.namespace": os.getenv("K8S_NAMESPACE", "swarm-system"), "k8s.pod.name": os.getenv("HOSTNAME", "unknown"), }) # Create tracer provider provider = TracerProvider(resource=resource) # Add OTLP exporter try: otlp_exporter = OTLPSpanExporter(endpoint=self.otlp_endpoint) provider.add_span_processor(BatchSpanProcessor(otlp_exporter)) logger.info(f"OTLP exporter configured: {self.otlp_endpoint}") except Exception as e: logger.warning(f"Failed to configure OTLP exporter: {e}") # Add console exporter for debugging if self.enable_console: console_exporter = ConsoleSpanExporter() provider.add_span_processor(BatchSpanProcessor(console_exporter)) # Set global tracer provider trace.set_tracer_provider(provider) # Auto-instrument libraries self._instrument_libraries() self.tracer = trace.get_tracer(__name__) logger.info(f"Tracing initialized for service: {self.service_name}") def _instrument_libraries(self): """Auto-instrument common libraries""" try: RedisInstrumentor().instrument() RequestsInstrumentor().instrument() LoggingInstrumentor().instrument() logger.info("Auto-instrumentation enabled") except Exception as e: logger.warning(f"Failed to auto-instrument libraries: {e}") @contextmanager def trace_operation( self, operation_name: str, attributes: Optional[Dict[str, Any]] = None, kind: SpanKind = SpanKind.INTERNAL ): """ Context manager for tracing an operation Usage: with tracer.trace_operation("process_task", {"task_id": "123"}): # do work pass """ with self.tracer.start_as_current_span( operation_name, kind=kind, attributes=attributes or {} ) as span: try: yield span except Exception as e: span.set_status(Status(StatusCode.ERROR, str(e))) span.record_exception(e) raise def trace_task_submission(self, task_id: str, task_description: str): """Trace task submission""" with self.trace_operation( "task.submit", { "task.id": task_id, "task.description": task_description[:100] # Truncate }, kind=SpanKind.PRODUCER ) as span: span.add_event("task_submitted") return span def trace_agent_creation(self, agent_id: str, task_id: str, pod_name: str): """Trace agent pod creation""" with self.trace_operation( "agent.create", { "agent.id": agent_id, "task.id": task_id, "k8s.pod.name": pod_name } ) as span: span.add_event("agent_pod_created") return span def trace_agent_execution(self, agent_id: str, task_id: str): """Trace agent task execution""" with self.trace_operation( "agent.execute", { "agent.id": agent_id, "task.id": task_id } ) as span: span.add_event("agent_started") return span def trace_handoff( self, from_agent_id: str, to_agent_id: str, task_id: str, handoff_reason: str ): """Trace agent handoff""" with self.trace_operation( "agent.handoff", { "handoff.from_agent": from_agent_id, "handoff.to_agent": to_agent_id, "task.id": task_id, "handoff.reason": handoff_reason }, kind=SpanKind.CLIENT ) as span: span.add_event("handoff_initiated") return span def trace_result_aggregation(self, task_id: str, agent_count: int): """Trace result aggregation""" with self.trace_operation( "result.aggregate", { "task.id": task_id, "agent.count": agent_count } ) as span: span.add_event("aggregation_started") return span def add_event(self, name: str, attributes: Optional[Dict[str, Any]] = None): """Add an event to the current span""" span = trace.get_current_span() if span: span.add_event(name, attributes or {}) def set_attribute(self, key: str, value: Any): """Set an attribute on the current span""" span = trace.get_current_span() if span: span.set_attribute(key, value) def record_error(self, error: Exception): """Record an error in the current span""" span = trace.get_current_span() if span: span.set_status(Status(StatusCode.ERROR, str(error))) span.record_exception(error) # Singleton instance _tracer_instance: Optional[SwarmTracer] = None def initialize_tracing( service_name: str, service_version: str = "1.0.0", otlp_endpoint: Optional[str] = None, enable_console: bool = False ) -> SwarmTracer: """Initialize global tracing instance""" global _tracer_instance _tracer_instance = SwarmTracer( service_name=service_name, service_version=service_version, otlp_endpoint=otlp_endpoint, enable_console=enable_console ) return _tracer_instance def get_tracer() -> Optional[SwarmTracer]: """Get the global tracer instance""" return _tracer_instance # Decorator for tracing functions def traced(operation_name: Optional[str] = None, **span_attributes): """ Decorator to automatically trace a function Usage: @traced("my_operation", task_id="123") def my_function(): pass """ def decorator(func): def wrapper(*args, **kwargs): tracer = get_tracer() if not tracer: return func(*args, **kwargs) op_name = operation_name or f"{func.__module__}.{func.__name__}" with tracer.trace_operation(op_name, span_attributes): return func(*args, **kwargs) return wrapper return decorator # Context propagation helpers def inject_trace_context(headers: Dict[str, str]) -> Dict[str, str]: """ Inject trace context into HTTP headers for propagation Usage: headers = inject_trace_context({}) requests.post(url, headers=headers) """ from opentelemetry.propagate import inject inject(headers) return headers def extract_trace_context(headers: Dict[str, str]): """ Extract trace context from HTTP headers Usage: extract_trace_context(request.headers) """ from opentelemetry.propagate import extract return extract(headers) # Example usage patterns """ # In orchestrator/main.py: from orchestrator.tracing import initialize_tracing, get_tracer tracer = initialize_tracing( service_name="swarm-orchestrator", service_version="1.0.0", otlp_endpoint="http://otel-collector:4317" ) # Trace task submission with tracer.trace_task_submission(task_id, description): # Submit task logic pass # In agent/main.py: from orchestrator.tracing import initialize_tracing, get_tracer tracer = initialize_tracing( service_name="swarm-agent", service_version="1.0.0" ) # Trace agent execution with tracer.trace_agent_execution(agent_id, task_id): # Execute task pass # Trace handoff with tracer.trace_handoff(from_agent, to_agent, task_id, reason): # Perform handoff pass # Using decorator @traced("process_subtask", task_id="123") def process_subtask(): pass # Manual span management tracer = get_tracer() with tracer.trace_operation("custom_operation", {"key": "value"}): tracer.add_event("checkpoint_reached") tracer.set_attribute("result_count", 42) """ kind: ConfigMap metadata: creationTimestamp: null name: orchestrator-source namespace: swarm-system