Agentswarm/scripts/test-swarm-competition.py

"""Integration test for decentralized-rework P4: task competition (#8).

Agents bid for a task; a deterministic arbitrator picks the winner and assigns it. Agents can
yield a task back, and request takeover of a held task (only winning decisively reassigns it).
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_TASK_COMPETITION=1 python scripts/test-swarm-competition.py
"""
import asyncio
import os
import sys
from pathlib import Path

os.environ["REDIS_FAKE"] = "1"
os.environ["ENABLE_TASK_COMPETITION"] = "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, TaskStatus
from orchestrator.agent_registry import agent_registry
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 mk_task(run, tid, caps=None):
    t = await task_queue.create_task(task_id=f"{run.swarm_id}-{tid}", description=tid,
                                     agent_role="implementation",
                                     required_capabilities=caps or ["python"], enqueue=True)
    await swarm_runtime.attach_task(run, t.task_id)
    return t


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

    body = {"mode": "swarm", "requirement": {"objective": "competition test"},
            "callback": {"url": "http://127.0.0.1:9/cb", "subscribed_events": []},
            "metadata": {"manager_deployment_id": "m-comp"}}
    run, _ = await swarm_runtime.get_or_create_run(body=body, idempotency_key=None, correlation_id="cc")

    for aid in ("bid-A", "bid-B", "yielder", "incumbent", "requester"):
        await agent_registry.register_agent(aid, ["python"])

    # --- 1. two agents bid → arbitrate → stronger wins + gets assigned ---
    t1 = await mk_task(run, "t1")
    await orch.handle_task_bid("bid-A", {"task_id": t1.task_id, "confidence": 0.9, "capabilities": ["python"]})
    r = await orch.handle_task_bid("bid-B", {"task_id": t1.task_id, "confidence": 0.5, "capabilities": ["python"]})
    check("two bids recorded", r.get("bid_count") == 2)
    # arbitrate_and_assign is called by the loop with a freshly-fetched run (bids live in Redis).
    run = await swarm_runtime.get_run(run.swarm_id)
    arb = await orch.arbitrate_and_assign(run, t1.task_id)
    check("arbitration picked higher-confidence bidder A", arb and arb["winner_agent_id"] == "bid-A")
    check("arbitration decisive + losers recorded", arb and arb["decisive"] and arb["losers"] == ["bid-B"])
    t1r = await task_queue.get_task(t1.task_id)
    check("winner assigned the task", t1r.assigned_agent_id == "bid-A" and t1r.status != TaskStatus.PENDING)
    run_r = await swarm_runtime.get_run(run.swarm_id)
    check("bids cleared after arbitration", not (run_r.metadata.get("bids") or {}).get(t1.task_id))
    check("arbitration audit recorded on run", bool(run_r.metadata.get("arbitrations")))

    # --- 2. yield: an assigned task is released back to pending ---
    t2 = await mk_task(run, "t2")
    await task_queue.remove_pending_task(t2.task_id)
    await task_queue.assign_task(t2.task_id, "yielder")
    y = await orch.handle_task_yield("yielder", {"task_id": t2.task_id, "reason": "context too large", "recommend_agent": "bid-A"})
    check("yield released the task", y.get("released") is True)
    t2r = await task_queue.get_task(t2.task_id)
    check("yielded task back to PENDING + agent cleared",
          t2r.status == TaskStatus.PENDING and t2r.assigned_agent_id is None)

    # --- 3. takeover: strong requester wins held task from incumbent ---
    t3 = await mk_task(run, "t3")
    await task_queue.remove_pending_task(t3.task_id)
    await task_queue.assign_task(t3.task_id, "incumbent")
    tk = await orch.handle_task_takeover("requester", {"task_id": t3.task_id, "confidence": 0.95, "capabilities": ["python"]})
    check("decisive requester took over", tk.get("taken_over") is True and tk.get("winner_agent_id") == "requester")
    t3r = await task_queue.get_task(t3.task_id)
    check("task reassigned to requester", t3r.assigned_agent_id == "requester")

    # --- 4. post-cutover: competition is unconditional (swarm default; no enable flag) ---
    os.environ.pop("ENABLE_TASK_COMPETITION", None)
    off = await orch.handle_task_bid("bid-A", {"task_id": t1.task_id, "confidence": 0.9, "capabilities": ["python"]})
    check("bids processed without any flag (swarm default)", off.get("recorded") is True)

    print()
    if failures:
        print(f"{len(failures)} competition (P4) check(s) FAILED: {failures}")
        sys.exit(1)
    print("all swarm competition (P4) checks passed")


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