Agentswarm/agent/main.py

"""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."""

    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)

    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:
        # Compute the dedicated per-task working directory inside the agent workspace.
        return self.workspace_dir / ".agent_tasks" / task_id

    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)
            task_workspace.mkdir(parents=True, exist_ok=True)

            git_enabled = False
            try:
                executor = TaskExecutor(agent_id=self.agent_id, workspace_dir=str(task_workspace))
                git_enabled = await self.workspace_git.is_git_workspace()
                if git_enabled:
                    branch_created = await self.workspace_git.create_result_branch(task_id)
                    if not branch_created:
                        logger.warning("Failed to create task result branch; continuing without git push")
                        git_enabled = False

                result = await asyncio.wait_for(
                    executor.execute_task(
                        task_id=task_id,
                        description=description,
                        context={
                            **context,
                            "workspace_dir": str(task_workspace),
                            "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:
                        commit_sha = await self.workspace_git.commit_changes(
                            message=f"Task {task_id}: {description[:50]}"
                        )
                        if commit_sha:
                            branch_name = await self.workspace_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:
                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())