Agentswarm/scripts/test-swarm-dispatch.py

"""Integration test for decentralized-rework P6: pheromone-driven agent self-selection.

Each idle agent perceives the eligible ready tasks and self-selects the best fit. With capability
and load equal, the differentiator is the pheromone trail (τ): an agent self-selects the role it
has historically succeeded at. Hermetic, no model key.

Run from agent_swarm_v6 (install deps first — needs fakeredis):
    pip install -r orchestrator/requirements.txt -r agent/requirements.txt
    REDIS_FAKE=1 ENABLE_SWARM_DISPATCH=1 python scripts/test-swarm-dispatch.py
"""
import asyncio
import os
import sys
from pathlib import Path

os.environ["REDIS_FAKE"] = "1"
os.environ["ENABLE_SWARM_DISPATCH"] = "1"
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from orchestrator.redis_client import redis_client
from orchestrator import swarm_runtime as sr_mod
from orchestrator.swarm_runtime import swarm_runtime
from orchestrator.task_queue import task_queue
from orchestrator.agent_registry import agent_registry
from orchestrator.decision_engine import decision_engine
from orchestrator import main as orch

failures = []


def check(name, cond):
    print(("PASS" if cond else "FAIL"), "-", name)
    if not cond:
        failures.append(name)


async def _noop(self, *a, **k):
    return None


async def ready_task(run, role):
    t = await task_queue.create_task(task_id=f"{run.swarm_id}-t-{role}", description=role,
                                     agent_role=role, required_capabilities=["python"], enqueue=True)
    await swarm_runtime.attach_task(run, t.task_id)
    return t


async def fresh_run(label):
    body = {"mode": "swarm", "requirement": {"objective": "self-select test"},
            "callback": {"url": "http://127.0.0.1:9/cb", "subscribed_events": []},
            "metadata": {"manager_deployment_id": label}}
    run, _ = await swarm_runtime.get_or_create_run(body=body, idempotency_key=None, correlation_id=label)
    return run


async def main():
    await redis_client.connect()
    sr_mod.SwarmRuntime._post_callback = _noop

    # Build pheromone profiles: X strong at implementation, Y strong at testing.
    for _ in range(10):
        await decision_engine.deposit(agent_role="implementation", agent_id="agent-X", success=True)
        await decision_engine.deposit(agent_role="testing", agent_id="agent-X", success=False)
        await decision_engine.deposit(agent_role="testing", agent_id="agent-Y", success=True)
        await decision_engine.deposit(agent_role="implementation", agent_id="agent-Y", success=False)

    # --- Run A: agent-X with BOTH tasks available self-selects its high-τ role (implementation) ---
    await agent_registry.register_agent("agent-X", ["python"])
    runA = await fresh_run("m-selA")
    await ready_task(runA, "implementation")
    await ready_task(runA, "testing")
    x = next(a for a in await agent_registry.get_idle_agents() if a.agent_id == "agent-X")
    assigned_a = await orch.swarm_dispatch([x])
    check("agent-X self-selected its high-τ role (implementation)",
          assigned_a == [("agent-X", f"{runA.swarm_id}-t-implementation")])
    runA_r = await swarm_runtime.get_run(runA.swarm_id)
    check("self-selection recorded an explainable dispatch decision",
          bool(runA_r.dispatch_decisions)
          and runA_r.dispatch_decisions[-1]["chosen_task_id"] == f"{runA.swarm_id}-t-implementation")

    # --- Run B: agent-Y with BOTH tasks available self-selects its high-τ role (testing) ---
    await agent_registry.register_agent("agent-Y", ["python"])
    runB = await fresh_run("m-selB")
    await ready_task(runB, "implementation")
    await ready_task(runB, "testing")
    y = next(a for a in await agent_registry.get_idle_agents() if a.agent_id == "agent-Y")
    assigned_b = await orch.swarm_dispatch([y])
    check("agent-Y self-selected its high-τ role (testing)",
          assigned_b == [("agent-Y", f"{runB.swarm_id}-t-testing")])

    # --- capability gate still holds: an agent lacking caps self-selects nothing ---
    await agent_registry.register_agent("agent-Z", ["rust"])
    runC = await fresh_run("m-selC")
    await ready_task(runC, "implementation")  # requires python
    z = next(a for a in await agent_registry.get_idle_agents() if a.agent_id == "agent-Z")
    assigned_c = await orch.swarm_dispatch([z])
    check("incapable agent self-selects nothing", assigned_c == [])

    print()
    if failures:
        print(f"{len(failures)} self-selection (P6) check(s) FAILED: {failures}")
        sys.exit(1)
    print("all swarm self-selection (P6) checks passed")


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