Agentswarm/scripts/stub_agent.py

"""Keyless stub agent for the SWARM workflow test (decentralized flow).

Connects over WebSocket, registers, and drives the swarm flow WITHOUT any LLM (no key needed):
when it self-selects the SEED task it PROPOSES the subtasks (bottom-up decomposition, #7) and then
completes the seed; for every subsequently self-selected subtask it returns a canned result. This
exercises seed → agent-proposed decomposition → self-selected execution → convergence — with no
Master plan and no central dispatch.

Run standalone against a manually started orchestrator:
    set "ORCHESTRATOR_URL=ws://localhost:8000"
    set "AGENT_ID=stub-1"
    set "AGENT_CAPABILITIES=python,code_generation,testing,pytest,technical-writing,general"
    python scripts/stub_agent.py
"""
import asyncio
import json
import os
import time

import websockets

# Subtasks the agent proposes when it perceives the seed (bottom-up decomposition).
SEED_SUBTASKS = [
    {"description": "Implement the core add(a, b) function in calc.py",
     "agent_role": "implementation", "required_capabilities": ["python", "code_generation"]},
    {"description": "Write pytest unit tests covering add(a, b)",
     "agent_role": "testing", "required_capabilities": ["testing", "pytest"]},
    {"description": "Document usage of add(a, b) in the README",
     "agent_role": "documentation", "required_capabilities": ["technical-writing", "general"]},
]


class StubAgent:
    def __init__(self, orchestrator_url: str, agent_id: str, capabilities: list[str]):
        self.orchestrator_url = orchestrator_url.rstrip("/")
        self.agent_id = agent_id
        self.capabilities = capabilities
        self.ws = None
        self.proposed = False   # only decompose the seed once
        self.running = True

    def _result_for(self, task_id: str, description: str) -> dict:
        return {
            "success": True,
            "task_id": task_id,
            "subtasks": [{"status": "completed", "summary": f"completed: {description[:60]}",
                          "changes": description[:60], "files": []}],
            "awaiting_handoff": False,
            "agent_id": self.agent_id,
            "usage": {
                "model_id": "stub", "model_tokens": 0, "prompt_tokens": 0,
                "completion_tokens": 0, "model_cost_usd": 0.0, "runtime_seconds": 0.0,
            },
        }

    async def send(self, payload: dict):
        await self.ws.send(json.dumps(payload))

    async def handle(self, msg: dict):
        msg_type = msg.get("type")
        if msg_type == "task_assignment":
            task_id = msg["task_id"]
            context = msg.get("context") or {}
            await self.send({"type": "task_accepted", "task_id": task_id, "available_slots": 4})
            await self.send({"type": "task_start", "agent_id": self.agent_id, "task_id": task_id, "timestamp": time.time()})
            await asyncio.sleep(0.2)  # simulate work

            # Perceived the seed → propose the subtasks (decentralized decomposition), then finish it.
            if context.get("is_seed") and not self.proposed:
                self.proposed = True
                for sub in SEED_SUBTASKS:
                    await self.send({
                        "type": "task_proposal", "agent_id": self.agent_id,
                        "origin_task_id": task_id, "confidence": 0.9,
                        "reason": "decomposing the seed objective into specialist subtasks",
                        "trigger_event": "task.assignment", **sub,
                    })
                await asyncio.sleep(0.4)  # let the orchestrator review + enqueue the proposals

            await self.send({
                "type": "task_complete", "agent_id": self.agent_id,
                "task_id": task_id, "result": self._result_for(task_id, msg.get("description", "")),
                "timestamp": time.time(),
            })
        elif msg_type == "peer_message" and not msg.get("is_reply"):
            await self.send({
                "type": "peer_message", "agent_id": self.agent_id,
                "target_agent_id": msg.get("from_agent_id"), "task_id": msg.get("task_id"),
                "content": f"stub guidance from {self.agent_id}",
                "correlation_id": msg.get("correlation_id"), "is_reply": True, "timestamp": time.time(),
            })

    async def heartbeat(self):
        while self.running:
            try:
                await self.send({
                    "type": "heartbeat", "agent_id": self.agent_id,
                    "timestamp": time.time(), "active_tasks": 0, "available_slots": 4,
                })
            except Exception:
                return
            await asyncio.sleep(5)

    async def run(self):
        async with websockets.connect(f"{self.orchestrator_url}/ws/{self.agent_id}") as ws:
            self.ws = ws
            await self.send({
                "type": "register", "agent_id": self.agent_id,
                "capabilities": self.capabilities, "available_slots": 4, "active_task_ids": [],
            })
            hb = asyncio.create_task(self.heartbeat())
            try:
                async for raw in ws:
                    await self.handle(json.loads(raw))
            finally:
                self.running = False
                hb.cancel()


async def _main():
    agent = StubAgent(
        os.getenv("ORCHESTRATOR_URL", "ws://localhost:8000"),
        os.getenv("AGENT_ID", "stub-1"),
        os.getenv("AGENT_CAPABILITIES", "python,code_generation,testing,pytest,technical-writing,general").split(","),
    )
    await agent.run()


if __name__ == "__main__":
    asyncio.run(_main())