gongzhiyongandClaude Opus 4.8
62d6fdf78c
取代上一版 Option B(执行器共用仓库根 + 串行 1 任务)。每个任务从本 agent 自己的 clone 切出独立 git worktree(完整仓库内容 + 独立分支/索引),执行器在其中读写真实源码, 任务级 GitOperations 在该 worktree 提交/推送结果分支,完成后回收 worktree。 - 单 agent 可并发多任务(MAX_CONCURRENT_TASKS 恢复 4),互不共用 checkout/index - worktree 置于仓外 /tmp/agent-worktrees(AGENT_WORKTREE_BASE 可配),主 checkout 不受污染 - _git_admin_lock 仅串行 worktree add/remove 等共享 .git plumbing,任务执行仍并行 - 跨 agent 隔离不变:每个 agent 仍各自 clone 一份(各自 pod) - 根治 "Empty workspace: no files detected"(worktree 自带仓库文件,已端到端验证) GitOperations 新增 add_task_worktree / remove_task_worktree。 Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
701 lines
31 KiB
Python
701 lines
31 KiB
Python
"""Agent main loop with WebSocket connection to orchestrator.
|
||
|
||
Merged from agent_swarm_v4:
|
||
- Bounded concurrency, duplicate protection, capacity reporting (available_slots)
|
||
- Serialized websocket sends (safe_send) and reconnect-with-reregister logic
|
||
- Per-task timeout and graceful task cancellation
|
||
- Separate repository root from per-task execution workspace (see git_operations)
|
||
- Peer collaboration: request/reply routing fixed relative to v4 (peer_waiters is now
|
||
initialized and inbound peer replies are routed back to the waiting coroutine; inbound
|
||
peer queries are answered with a lightweight, cost-free acknowledgement)
|
||
- OpenAI-only task executor (see task_executor)
|
||
- Preserves heicode handoff wiring, Manager-facing metrics, and env-var entrypoint
|
||
"""
|
||
import asyncio
|
||
import json
|
||
import logging
|
||
import os
|
||
import signal
|
||
import time
|
||
import uuid
|
||
from pathlib import Path
|
||
from typing import Optional
|
||
|
||
import websockets
|
||
from dotenv import load_dotenv
|
||
from pydantic import BaseModel, Field, ValidationError
|
||
from prometheus_client import Counter, Gauge, Histogram, start_http_server
|
||
from websockets.exceptions import ConnectionClosed
|
||
|
||
from .git_operations import GitOperations
|
||
from .task_executor import TaskExecutor
|
||
|
||
|
||
load_dotenv()
|
||
|
||
|
||
# load_dotenv() above reads a local .env (gitignored) so credentials need not be passed on
|
||
# the command line. These defaults just guarantee the vars exist; an empty OPENAI_API_KEY is
|
||
# still falsy, so TaskExecutor raises a clear error rather than silently using no key.
|
||
# Real deployments supply the key via the environment / secret_ref, NOT a committed file.
|
||
os.environ.setdefault("OPENAI_API_KEY", "")
|
||
os.environ.setdefault("OPENAI_API_BASE", "https://api.openai.com/v1")
|
||
os.environ.setdefault("OPENAI_MODEL", "gpt-4o-mini")
|
||
|
||
logging.basicConfig(
|
||
level=logging.INFO,
|
||
format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
|
||
)
|
||
logger = logging.getLogger(__name__)
|
||
|
||
TASKS_EXECUTED = Counter("agent_tasks_executed_total", "Total tasks executed")
|
||
TASKS_FAILED = Counter("agent_tasks_failed_total", "Total tasks failed")
|
||
TASK_DURATION = Histogram("agent_task_duration_seconds", "Task execution duration")
|
||
HANDOFFS_INITIATED = Counter("agent_handoffs_initiated_total", "Total handoffs initiated")
|
||
AGENT_STATUS = Gauge("agent_status", "Agent status (0=idle, 1=busy, 2=failed)")
|
||
WEBSOCKET_RECONNECTS = Counter("agent_websocket_reconnects_total", "WebSocket reconnection attempts")
|
||
TASKS_REJECTED = Counter("agent_tasks_rejected_total", "Total tasks rejected due to capacity")
|
||
TASKS_DUPLICATE = Counter("agent_tasks_duplicate_total", "Total duplicate task assignments")
|
||
ACTIVE_TASKS = Gauge("agent_active_tasks", "Current number of active tasks")
|
||
|
||
|
||
class TaskAssignment(BaseModel):
|
||
task_id: str
|
||
description: str
|
||
context: dict = Field(default_factory=dict)
|
||
|
||
|
||
class AgentRuntimeDisconnected(RuntimeError):
|
||
"""Raised when the runtime attempts to send without an active websocket."""
|
||
|
||
|
||
class Agent:
|
||
"""Agent that connects to orchestrator and executes tasks."""
|
||
|
||
# An agent may run several tasks at once; each executes in its OWN git worktree cut from this
|
||
# agent's clone (see _execute_assignment / GitOperations.add_task_worktree), so concurrent tasks
|
||
# never share a checkout/index. Swarm parallelism = many agents × this per-agent concurrency.
|
||
MAX_CONCURRENT_TASKS = int(os.getenv("MAX_CONCURRENT_TASKS", "4"))
|
||
TASK_TIMEOUT_SECONDS = int(os.getenv("TASK_TIMEOUT_SECONDS", "60"))
|
||
PEER_REPLY_TIMEOUT_SECONDS = int(os.getenv("PEER_REPLY_TIMEOUT_SECONDS", "20"))
|
||
HEARTBEAT_INTERVAL_SECONDS = 15
|
||
|
||
def __init__(
|
||
self,
|
||
orchestrator_url: str,
|
||
agent_id: Optional[str] = None,
|
||
capabilities: Optional[list[str]] = None,
|
||
workspace_dir: str = "/workspace",
|
||
git_repo_url: Optional[str] = None,
|
||
user_id: Optional[str] = None,
|
||
):
|
||
# Initialize connection state, execution limits, and workspace helpers for this agent runtime.
|
||
self.agent_id = agent_id or f"agent-{uuid.uuid4().hex[:8]}"
|
||
self.orchestrator_url = orchestrator_url.rstrip("/")
|
||
self.capabilities = capabilities or ["general"]
|
||
# The owning user; the orchestrator caps concurrent agents per user (MAX_AGENTS_PER_USER).
|
||
# The agent platform sets this; falls back to env. None => unbound (not subject to the cap).
|
||
self.user_id = user_id or os.getenv("HEICODE_USER_ID")
|
||
self.workspace_dir = Path(workspace_dir)
|
||
self.git_repo_url = git_repo_url
|
||
|
||
self.websocket: Optional[websockets.WebSocketClientProtocol] = None
|
||
self.running = False
|
||
self.current_task_id: Optional[str] = None
|
||
self.send_lock = asyncio.Lock()
|
||
self.task_semaphore = asyncio.Semaphore(self.MAX_CONCURRENT_TASKS)
|
||
self.active_tasks: dict[str, asyncio.Task] = {}
|
||
self.heartbeat_task: Optional[asyncio.Task] = None
|
||
# Outstanding peer-collaboration requests awaiting a reply, keyed by correlation_id.
|
||
self.peer_waiters: dict[str, asyncio.Future] = {}
|
||
# Summary of this agent's most recently completed task, shared when peers consult it.
|
||
self.last_summary: Optional[str] = None
|
||
# Lazily-created executor used to compose substantive peer replies.
|
||
self._peer_executor: Optional[TaskExecutor] = None
|
||
|
||
self.workspace_git = GitOperations(str(self.workspace_dir), self.agent_id)
|
||
# Serializes only the brief shared-repo git plumbing (worktree add/remove + fetch) so
|
||
# concurrent tasks don't race on .git locks. Task EXECUTION stays parallel.
|
||
self._git_admin_lock = asyncio.Lock()
|
||
|
||
def available_slots(self) -> int:
|
||
# Return how many additional tasks this agent can currently accept.
|
||
return max(0, self.MAX_CONCURRENT_TASKS - len(self.active_tasks))
|
||
|
||
def task_workspace(self, task_id: str) -> Path:
|
||
# Per-task git worktree path. Kept OUTSIDE the repo root (default /tmp/agent-worktrees) so
|
||
# the main checkout never sees it as untracked, and namespaced by agent_id to avoid
|
||
# collisions. `git worktree add` requires the leaf to not pre-exist, so we don't mkdir it.
|
||
base = Path(os.getenv("AGENT_WORKTREE_BASE", "/tmp/agent-worktrees"))
|
||
return base / self.agent_id / task_id.replace("/", "-")
|
||
|
||
async def safe_send(self, payload: dict):
|
||
# Serialize and send a websocket message while holding a lock to prevent concurrent writes.
|
||
if not self.websocket:
|
||
raise AgentRuntimeDisconnected("websocket unavailable")
|
||
async with self.send_lock:
|
||
await self.websocket.send(json.dumps(payload))
|
||
|
||
async def connect(self) -> bool:
|
||
# Open a websocket connection to the orchestrator and cache the live socket on success.
|
||
try:
|
||
self.websocket = await websockets.connect(f"{self.orchestrator_url}/ws/{self.agent_id}")
|
||
logger.info(f"Connected to orchestrator at {self.orchestrator_url}")
|
||
return True
|
||
except Exception as e:
|
||
logger.error(f"Failed to connect to orchestrator: {e}")
|
||
return False
|
||
|
||
async def connect_with_retry(self) -> bool:
|
||
# Repeatedly attempt to connect with exponential backoff until connected or shutting down.
|
||
delay = 1
|
||
while self.running:
|
||
if await self.connect():
|
||
return True
|
||
WEBSOCKET_RECONNECTS.inc()
|
||
await asyncio.sleep(min(delay, 60))
|
||
delay = min(delay * 2, 60)
|
||
return False
|
||
|
||
async def register(self):
|
||
# Announce this agent and its current capacity to the orchestrator after connecting.
|
||
try:
|
||
register_msg = {
|
||
"type": "register",
|
||
"agent_id": self.agent_id,
|
||
"capabilities": self.capabilities,
|
||
# Backward-compatible extra fields; older orchestrators ignore them.
|
||
"available_slots": self.available_slots(),
|
||
"active_task_ids": list(self.active_tasks.keys()),
|
||
}
|
||
if self.user_id:
|
||
register_msg["user_id"] = self.user_id # subjects this agent to the per-user cap
|
||
await self.safe_send(register_msg)
|
||
logger.info(f"Registered agent {self.agent_id} with capabilities: {self.capabilities}")
|
||
return True
|
||
except Exception as e:
|
||
logger.error(f"Failed to register: {e}")
|
||
return False
|
||
|
||
async def send_heartbeat(self):
|
||
# Send a periodic liveness and capacity update to the orchestrator.
|
||
await self.safe_send({
|
||
"type": "heartbeat",
|
||
"agent_id": self.agent_id,
|
||
"timestamp": time.time(),
|
||
"active_tasks": len(self.active_tasks),
|
||
"available_slots": self.available_slots(),
|
||
})
|
||
|
||
async def heartbeat_loop(self):
|
||
# Keep sending heartbeat messages until the connection closes or the agent stops running.
|
||
while self.running:
|
||
try:
|
||
await self.send_heartbeat()
|
||
except (ConnectionClosed, AgentRuntimeDisconnected):
|
||
return
|
||
except Exception as e:
|
||
logger.error(f"Failed to send heartbeat: {e}")
|
||
return
|
||
await asyncio.sleep(self.HEARTBEAT_INTERVAL_SECONDS)
|
||
|
||
async def send_status_update(self, status: str, task_id: Optional[str] = None, message: str = ""):
|
||
# Report a human-readable task or agent status transition to the orchestrator.
|
||
try:
|
||
await self.safe_send({
|
||
"type": "status_update",
|
||
"agent_id": self.agent_id,
|
||
"status": status,
|
||
"task_id": task_id,
|
||
"message": message,
|
||
"timestamp": time.time(),
|
||
})
|
||
except Exception as e:
|
||
logger.error(f"Failed to send status update: {e}")
|
||
|
||
async def request_peer_collaboration(
|
||
self,
|
||
task_id: str,
|
||
target_agent_id: str,
|
||
content: str,
|
||
timeout_seconds: float = 20.0,
|
||
):
|
||
# Ask a peer agent (via the orchestrator) for guidance and wait for its reply.
|
||
correlation_id = f"peer-{task_id}-{uuid.uuid4().hex[:8]}"
|
||
loop = asyncio.get_running_loop()
|
||
waiter = loop.create_future()
|
||
self.peer_waiters[correlation_id] = waiter
|
||
try:
|
||
await self.safe_send({
|
||
"type": "peer_message",
|
||
"agent_id": self.agent_id,
|
||
"target_agent_id": target_agent_id,
|
||
"task_id": task_id,
|
||
"content": content,
|
||
"correlation_id": correlation_id,
|
||
"is_reply": False,
|
||
"timestamp": time.time(),
|
||
})
|
||
return await asyncio.wait_for(waiter, timeout=timeout_seconds)
|
||
finally:
|
||
self.peer_waiters.pop(correlation_id, None)
|
||
|
||
@staticmethod
|
||
def _summarize_result(result: dict) -> Optional[str]:
|
||
# Extract a short, human-readable summary from an execution result.
|
||
if not isinstance(result, dict):
|
||
return None
|
||
for subtask_result in reversed(result.get("subtasks", []) or []):
|
||
summary = subtask_result.get("summary")
|
||
if isinstance(summary, str) and summary.strip():
|
||
return summary.strip()
|
||
summary = result.get("summary")
|
||
return summary.strip() if isinstance(summary, str) and summary.strip() else None
|
||
|
||
def _get_peer_executor(self) -> TaskExecutor:
|
||
# Lazily create a TaskExecutor for peer replies (reuses the model client / workspace).
|
||
# Raises if no model key is configured; callers fall back to the cached summary.
|
||
if self._peer_executor is None:
|
||
self._peer_executor = TaskExecutor(agent_id=self.agent_id, workspace_dir=str(self.workspace_dir))
|
||
return self._peer_executor
|
||
|
||
def _peer_fallback_reply(self) -> dict:
|
||
# Cheap, no-LLM reply used when the model is unavailable or errors.
|
||
shared = (
|
||
f"My latest result: {self.last_summary}. "
|
||
if self.last_summary
|
||
else "No completed result yet. "
|
||
)
|
||
content = (
|
||
f"From {self.agent_id} (capabilities: {', '.join(self.capabilities)}). {shared}"
|
||
"Treat implementation artifacts as the source of truth for behavior and exception semantics."
|
||
)
|
||
return {"content": content, "stance": "info", "evidence": self.last_summary or "", "refs": []}
|
||
|
||
async def _build_peer_reply(self, query: str, task_id: Optional[str]) -> dict:
|
||
# Produce a substantive, query-scoped reply grounded in this agent's own work.
|
||
# Falls back to the cached summary if there is no query, no model key, or the call fails.
|
||
if not query:
|
||
return self._peer_fallback_reply()
|
||
try:
|
||
executor = self._get_peer_executor()
|
||
except Exception as e:
|
||
logger.warning(f"peer reply executor unavailable ({e}); using cached summary")
|
||
return self._peer_fallback_reply()
|
||
try:
|
||
return await asyncio.wait_for(
|
||
executor.peer_reply(
|
||
query=query,
|
||
capabilities=self.capabilities,
|
||
last_summary=self.last_summary,
|
||
),
|
||
timeout=self.PEER_REPLY_TIMEOUT_SECONDS,
|
||
)
|
||
except Exception as e:
|
||
logger.warning(f"peer reply LLM failed ({e}); using cached summary")
|
||
return self._peer_fallback_reply()
|
||
|
||
async def answer_peer_query(self, message: dict):
|
||
# Respond to an inbound peer query with a substantive, grounded reply (LLM, with fallback).
|
||
requester = message.get("from_agent_id") or message.get("agent_id")
|
||
correlation_id = message.get("correlation_id")
|
||
if not requester or not correlation_id:
|
||
return
|
||
reply = await self._build_peer_reply((message.get("content") or "").strip(), message.get("task_id"))
|
||
try:
|
||
await self.safe_send({
|
||
"type": "peer_message",
|
||
"agent_id": self.agent_id,
|
||
"target_agent_id": requester,
|
||
"task_id": message.get("task_id"),
|
||
"content": reply.get("content", ""),
|
||
"stance": reply.get("stance"),
|
||
"evidence": reply.get("evidence"),
|
||
"refs": reply.get("refs"),
|
||
"correlation_id": correlation_id,
|
||
"is_reply": True,
|
||
"timestamp": time.time(),
|
||
})
|
||
except Exception as e:
|
||
logger.error(f"Failed to answer peer query: {e}")
|
||
|
||
async def send_task_result(self, task_id: str, success: bool, result: dict):
|
||
# Send a completion or failure payload for a finished task execution.
|
||
try:
|
||
payload = {
|
||
"type": "task_complete" if success else "task_failed",
|
||
"agent_id": self.agent_id,
|
||
"task_id": task_id,
|
||
"timestamp": time.time(),
|
||
}
|
||
if success:
|
||
payload["result"] = result
|
||
else:
|
||
payload["reason"] = result.get("error", "Task failed")
|
||
payload["result"] = result
|
||
await self.safe_send(payload)
|
||
except Exception as e:
|
||
logger.error(f"Failed to send task result: {e}")
|
||
|
||
async def send_blocked_on_handoff(self, task_id: str, result: dict):
|
||
# Report that a task is paused because execution delegated work to a child handoff task.
|
||
try:
|
||
child_task_id = None
|
||
for subtask_result in result.get("subtasks", []):
|
||
if subtask_result.get("status") == "handed_off":
|
||
child_task_id = subtask_result.get("child_task_id")
|
||
break
|
||
await self.safe_send({
|
||
"type": "blocked_on_handoff",
|
||
"agent_id": self.agent_id,
|
||
"task_id": task_id,
|
||
"child_task_id": child_task_id,
|
||
"reason": "Waiting on delegated child task",
|
||
"result": result,
|
||
"timestamp": time.time(),
|
||
})
|
||
except Exception as e:
|
||
logger.error(f"Failed to send blocked_on_handoff: {e}")
|
||
|
||
async def request_handoff(self, task_id: str, subtask: dict, target_capabilities: list[str]):
|
||
# Ask the orchestrator to delegate a discovered subtask to another capable agent.
|
||
try:
|
||
await self.safe_send({
|
||
"type": "handoff_request",
|
||
"agent_id": self.agent_id,
|
||
"task_id": task_id,
|
||
"subtask": subtask,
|
||
"target_capabilities": target_capabilities,
|
||
"timestamp": time.time(),
|
||
})
|
||
HANDOFFS_INITIATED.inc()
|
||
except Exception as e:
|
||
logger.error(f"Failed to request handoff: {e}")
|
||
|
||
async def execute_assignment(self, assignment: TaskAssignment):
|
||
# Execute one accepted assignment, manage workspace/git flow, and publish lifecycle updates.
|
||
async with self.task_semaphore:
|
||
task_id = assignment.task_id
|
||
description = assignment.description
|
||
context = assignment.context or {}
|
||
self.current_task_id = task_id
|
||
ACTIVE_TASKS.set(len(self.active_tasks))
|
||
AGENT_STATUS.set(1)
|
||
|
||
await self.safe_send({
|
||
"type": "task_start",
|
||
"agent_id": self.agent_id,
|
||
"task_id": task_id,
|
||
"timestamp": time.time(),
|
||
})
|
||
await self.send_status_update("busy", task_id, "Starting task execution")
|
||
|
||
start_time = time.time()
|
||
task_workspace = self.task_workspace(task_id)
|
||
|
||
git_enabled = False
|
||
task_git = None
|
||
try:
|
||
# Each task runs in its OWN git worktree cut from this agent's clone, so an agent can
|
||
# run several tasks concurrently without sharing a checkout/index. The worktree holds
|
||
# the full repo contents on a fresh result branch; the executor reads/writes the real
|
||
# source there and the per-task GitOperations (task_git) commits/pushes that branch.
|
||
# This fixes the prior empty per-task subdir → "Empty workspace: no files detected".
|
||
# Falls back to the repo root (no isolated branch) only if worktree creation fails.
|
||
if await self.workspace_git.is_git_workspace():
|
||
async with self._git_admin_lock: # serialize shared-repo git plumbing only
|
||
result_branch = await self.workspace_git.add_task_worktree(str(task_workspace), task_id)
|
||
if result_branch:
|
||
git_enabled = True
|
||
task_git = GitOperations(str(task_workspace), self.agent_id)
|
||
task_git.result_branch = result_branch
|
||
else:
|
||
logger.warning("Failed to create task worktree; executing on repo root without git push")
|
||
exec_dir = str(task_workspace) if git_enabled else str(self.workspace_dir)
|
||
executor = TaskExecutor(agent_id=self.agent_id, workspace_dir=exec_dir)
|
||
|
||
result = await asyncio.wait_for(
|
||
executor.execute_task(
|
||
task_id=task_id,
|
||
description=description,
|
||
context={
|
||
**context,
|
||
"workspace_dir": exec_dir,
|
||
"repo_workspace_dir": str(self.workspace_dir),
|
||
"git_repo_url": self.git_repo_url,
|
||
"agent_id": self.agent_id,
|
||
},
|
||
handoff_callback=self.request_handoff,
|
||
agent_capabilities=self.capabilities,
|
||
peer_collaboration_callback=self.request_peer_collaboration,
|
||
),
|
||
timeout=self.TASK_TIMEOUT_SECONDS,
|
||
)
|
||
|
||
awaiting_handoff = result.get("awaiting_handoff", False)
|
||
if result.get("success"):
|
||
self.last_summary = self._summarize_result(result) or self.last_summary
|
||
if result.get("success") and not awaiting_handoff:
|
||
if git_enabled and task_git:
|
||
commit_sha = await task_git.commit_changes(
|
||
message=f"Task {task_id}: {description[:50]}"
|
||
)
|
||
if commit_sha:
|
||
branch_name = await task_git.push_results()
|
||
result["git_branch"] = branch_name
|
||
result["commit_sha"] = commit_sha
|
||
else:
|
||
result["git_skipped"] = "No workspace changes to commit"
|
||
else:
|
||
result["git_skipped"] = "Workspace is not a Git checkout"
|
||
|
||
if awaiting_handoff:
|
||
await self.send_blocked_on_handoff(task_id, result)
|
||
else:
|
||
await self.send_task_result(task_id, result.get("success", False), result)
|
||
|
||
duration = time.time() - start_time
|
||
TASK_DURATION.observe(duration)
|
||
if result.get("success"):
|
||
TASKS_EXECUTED.inc()
|
||
else:
|
||
TASKS_FAILED.inc()
|
||
|
||
AGENT_STATUS.set(0)
|
||
if awaiting_handoff:
|
||
await self.send_status_update("handoff-pending", task_id, "Waiting for delegated child task")
|
||
await self.send_status_update("idle", None, "Delegated child task created")
|
||
else:
|
||
await self.send_status_update("idle", None, "Task completed")
|
||
|
||
except asyncio.TimeoutError:
|
||
TASKS_FAILED.inc()
|
||
AGENT_STATUS.set(2)
|
||
await self.send_task_result(task_id, False, {"error": "timeout", "success": False})
|
||
await self.send_status_update("idle", None, "Task failed: timeout")
|
||
except asyncio.CancelledError:
|
||
await self.send_task_result(task_id, False, {"error": "cancelled", "success": False})
|
||
await self.send_status_update("idle", None, "Task cancelled")
|
||
raise
|
||
except Exception as e:
|
||
logger.error(f"Error executing task {task_id}: {e}")
|
||
TASKS_FAILED.inc()
|
||
AGENT_STATUS.set(2)
|
||
await self.send_task_result(task_id, False, {"error": str(e), "success": False})
|
||
await self.send_status_update("idle", None, f"Task failed: {e}")
|
||
finally:
|
||
if git_enabled:
|
||
async with self._git_admin_lock:
|
||
await self.workspace_git.remove_task_worktree(str(task_workspace))
|
||
self.active_tasks.pop(task_id, None)
|
||
ACTIVE_TASKS.set(len(self.active_tasks))
|
||
self.current_task_id = None
|
||
if not self.active_tasks:
|
||
AGENT_STATUS.set(0)
|
||
|
||
async def cancel_task(self, task_id: str):
|
||
# Cancel an actively running asyncio task if the orchestrator requests termination.
|
||
task = self.active_tasks.get(task_id)
|
||
if task:
|
||
task.cancel()
|
||
|
||
async def handle_task_assignment(self, message: dict):
|
||
# Validate an incoming assignment, reject duplicates/capacity overflow, and start execution.
|
||
try:
|
||
assignment = TaskAssignment(**message)
|
||
except ValidationError:
|
||
logger.warning(f"invalid task assignment: {message}")
|
||
return
|
||
|
||
if assignment.task_id in self.active_tasks:
|
||
TASKS_DUPLICATE.inc()
|
||
await self.safe_send({
|
||
"type": "task_accepted",
|
||
"task_id": assignment.task_id,
|
||
"status": "duplicate",
|
||
})
|
||
return
|
||
|
||
if len(self.active_tasks) >= self.MAX_CONCURRENT_TASKS:
|
||
TASKS_REJECTED.inc()
|
||
await self.safe_send({
|
||
"type": "task_rejected",
|
||
"task_id": assignment.task_id,
|
||
"reason": "at_capacity",
|
||
"available_slots": self.available_slots(),
|
||
})
|
||
return
|
||
|
||
await self.safe_send({
|
||
"type": "task_accepted",
|
||
"task_id": assignment.task_id,
|
||
"available_slots": self.available_slots() - 1,
|
||
})
|
||
|
||
task = asyncio.create_task(self.execute_assignment(assignment))
|
||
self.active_tasks[assignment.task_id] = task
|
||
ACTIVE_TASKS.set(len(self.active_tasks))
|
||
|
||
def _resolve_peer_reply(self, message: dict) -> bool:
|
||
# Resolve the waiting future for an inbound peer reply; return True if it was a reply.
|
||
correlation_id = message.get("correlation_id")
|
||
if not correlation_id:
|
||
return False
|
||
waiter = self.peer_waiters.get(correlation_id)
|
||
if waiter and not waiter.done():
|
||
waiter.set_result(message)
|
||
return True
|
||
# A correlation we own but already resolved/timed out: treat as handled reply.
|
||
return bool(message.get("is_reply"))
|
||
|
||
async def handle_peer_message(self, message: dict):
|
||
# Route a peer message: resolve our own pending request, or answer an inbound query.
|
||
if message.get("is_reply") or message.get("correlation_id") in self.peer_waiters:
|
||
handled = self._resolve_peer_reply(message)
|
||
if handled:
|
||
return
|
||
logger.info(
|
||
"Received peer query for task %s from %s",
|
||
message.get("task_id"),
|
||
message.get("from_agent_id") or message.get("agent_id"),
|
||
)
|
||
# Answer in the background so a slow (LLM) reply doesn't stall the message loop.
|
||
asyncio.create_task(self.answer_peer_query(message))
|
||
|
||
async def handle_message(self, message: dict):
|
||
# Route each inbound orchestrator message to the appropriate handler.
|
||
msg_type = message.get("type")
|
||
control_messages = {"registered", "heartbeat_ack", "task_completed", "task_failed_ack", "task_blocked_ack"}
|
||
|
||
if msg_type == "task_assignment":
|
||
await self.handle_task_assignment(message)
|
||
elif msg_type == "handoff_response":
|
||
logger.info(f"Handoff accepted for task {message.get('task_id')}")
|
||
elif msg_type == "cancel_task":
|
||
await self.cancel_task(message["task_id"])
|
||
elif msg_type == "ping":
|
||
await self.safe_send({"type": "pong"})
|
||
elif msg_type == "peer_message":
|
||
await self.handle_peer_message(message)
|
||
elif msg_type in control_messages:
|
||
logger.debug(f"Received control message: {msg_type}")
|
||
else:
|
||
logger.warning(f"Unknown message type: {msg_type}")
|
||
|
||
async def message_loop(self):
|
||
# Continuously receive websocket messages and dispatch them until the connection ends.
|
||
async for raw in self.websocket:
|
||
try:
|
||
data = json.loads(raw)
|
||
await self.handle_message(data)
|
||
except json.JSONDecodeError as e:
|
||
logger.error(f"Failed to parse message: {e}")
|
||
except Exception as e:
|
||
logger.error(f"Error handling message: {e}")
|
||
|
||
async def shutdown_active_tasks(self):
|
||
# Cancel and await all active task coroutines during agent shutdown.
|
||
tasks = list(self.active_tasks.values())
|
||
for task in tasks:
|
||
task.cancel()
|
||
if tasks:
|
||
await asyncio.gather(*tasks, return_exceptions=True)
|
||
|
||
def install_signal_handlers(self):
|
||
# Register process signal handlers that trigger a graceful runtime shutdown.
|
||
def _handle_shutdown():
|
||
# Flip the runtime into shutdown mode and close the websocket asynchronously.
|
||
logger.info("shutdown signal received")
|
||
self.running = False
|
||
if self.websocket:
|
||
asyncio.create_task(self.websocket.close())
|
||
|
||
for sig in (signal.SIGINT, signal.SIGTERM):
|
||
try:
|
||
asyncio.get_running_loop().add_signal_handler(sig, _handle_shutdown)
|
||
except NotImplementedError:
|
||
# add_signal_handler is unsupported on Windows event loops; skip gracefully.
|
||
pass
|
||
|
||
async def run(self):
|
||
# Run the full agent lifecycle: startup, registration, message processing, reconnect, and cleanup.
|
||
self.running = True
|
||
self.install_signal_handlers()
|
||
|
||
if os.getenv("METRICS_PORT"):
|
||
start_http_server(int(os.getenv("METRICS_PORT", "9000")))
|
||
|
||
if self.git_repo_url:
|
||
logger.info(f"Cloning workspace from {self.git_repo_url}")
|
||
if not await self.workspace_git.clone_workspace(self.git_repo_url):
|
||
logger.error("Failed to clone workspace, exiting")
|
||
return
|
||
|
||
while self.running:
|
||
connected = await self.connect_with_retry()
|
||
if not connected:
|
||
break
|
||
|
||
if not await self.register():
|
||
logger.error("Failed to register with orchestrator, exiting")
|
||
return
|
||
|
||
self.heartbeat_task = asyncio.create_task(self.heartbeat_loop())
|
||
try:
|
||
await self.message_loop()
|
||
except ConnectionClosed:
|
||
logger.warning("connection closed; reconnecting")
|
||
except Exception as e:
|
||
logger.error(f"Error in message loop: {e}")
|
||
finally:
|
||
if self.heartbeat_task:
|
||
self.heartbeat_task.cancel()
|
||
try:
|
||
await self.heartbeat_task
|
||
except asyncio.CancelledError:
|
||
pass
|
||
self.heartbeat_task = None
|
||
|
||
if self.websocket:
|
||
try:
|
||
await self.safe_send({"type": "deregister", "agent_id": self.agent_id})
|
||
except Exception:
|
||
pass
|
||
try:
|
||
await self.websocket.close()
|
||
except Exception:
|
||
pass
|
||
self.websocket = None
|
||
|
||
if self.running:
|
||
await asyncio.sleep(1)
|
||
|
||
await self.shutdown_active_tasks()
|
||
logger.info("Agent shutdown complete")
|
||
|
||
|
||
async def main():
|
||
# Build an agent instance from environment configuration and start its runtime loop.
|
||
orchestrator_url = os.getenv("ORCHESTRATOR_URL", "ws://localhost:8000")
|
||
agent_id = os.getenv("AGENT_ID")
|
||
capabilities = os.getenv("AGENT_CAPABILITIES", "general").split(",")
|
||
workspace_dir = os.getenv("WORKSPACE_DIR", "/workspace")
|
||
git_repo_url = os.getenv("GIT_REPO_URL")
|
||
|
||
logger.info(f"Starting agent with ID: {agent_id or 'auto-generated'}")
|
||
logger.info(f"Capabilities: {capabilities}")
|
||
logger.info(f"Orchestrator URL: {orchestrator_url}")
|
||
|
||
agent = Agent(
|
||
orchestrator_url=orchestrator_url,
|
||
agent_id=agent_id,
|
||
capabilities=capabilities,
|
||
workspace_dir=workspace_dir,
|
||
git_repo_url=git_repo_url,
|
||
)
|
||
await agent.run()
|
||
|
||
|
||
if __name__ == "__main__":
|
||
asyncio.run(main())
|