Agentswarm/scripts/test-merge-smoke.py

"""Smoke tests for the agent_swarm_v4 -> heicode-swarm merge.

Exercises the new/changed code paths with the gated in-memory (fakeredis) fallback:
- redis_client REDIS_FAKE fallback through the rich list API
- task_queue.release_task (capacity-rejection requeue, no retry increment)
- planner static fallback + build_planner_task_specs mapping & dependency filtering
- agent peer-collaboration reply routing and capacity rejection messages
"""
import asyncio
import os
import sys
import types
from pathlib import Path

os.environ["REDIS_FAKE"] = "1"            # gated dev/CI in-memory store
os.environ.pop("OPENAI_API_KEY", None)    # force planner static fallback
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from orchestrator.redis_client import redis_client
from orchestrator.agent_registry import agent_registry, AgentStatus
from orchestrator.task_queue import task_queue, TaskStatus
from orchestrator import main as orch
from orchestrator.planner import planner

# Importing orchestrator.main runs load_dotenv(), which may set a real OPENAI_API_KEY from a
# local .env and give the planner a live client. Force it offline so these checks stay
# hermetic and deterministic (static plan + heuristic review + concatenated synthesis).
planner.client = None

failures = []


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


async def test_redis_fallback():
    await redis_client.connect()
    await redis_client.set("k", "v")
    check("redis_fallback get/set", (await redis_client.get("k")) == "v")
    await redis_client.lpush("q", "a")
    await redis_client.lpush("q", "b")
    check("redis_fallback list ops", (await redis_client.lrange("q", 0, -1)) == ["b", "a"])
    await redis_client.lrem("q", 0, "a")
    check("redis_fallback lrem", (await redis_client.lrange("q", 0, -1)) == ["b"])


async def test_release_task():
    await agent_registry.register_agent("agent-x", ["general"])
    task = await task_queue.create_task(description="do thing", task_id="t-rel-1")
    assigned = await task_queue.assign_task(task.task_id, "agent-x")
    check("assign_task ok", assigned)
    released = await task_queue.release_task(task.task_id, agent_id="agent-x")
    check("release_task ok", released)
    reloaded = await task_queue.get_task(task.task_id)
    agent = await agent_registry.get_agent("agent-x")
    pending = await redis_client.lrange(task_queue.PENDING_QUEUE_KEY, 0, -1)
    check("release sets PENDING", reloaded.status == TaskStatus.PENDING)
    check("release no retry increment", reloaded.retry_count == 0)
    check("release frees agent", agent.status == AgentStatus.IDLE)
    check("release requeues task", task.task_id in pending)


async def test_seeder():
    # Swarm task-creation: an objective is SEEDED as one task (no Master decomposition). Agents
    # grow the graph bottom-up at runtime (handle_task_proposal). Tested e2e in test-workflow-e2e.
    run = types.SimpleNamespace(swarm_id="swarm-test", objective="Build a calculator")
    base_specs = [{"context": {"orchestration_plan": {"x": 1}}}]
    specs = orch.build_seed_task_specs(run, {}, base_specs)
    check("seeder injects exactly one seed task", len(specs) == 1)
    check("seed source is 'seed' (no planner)", specs[0]["source"] == "seed")
    check("seed has no required capabilities (any agent self-selects)", specs[0]["required_capabilities"] == [])
    check("seed preserves base context", specs[0]["context"].get("orchestration_plan") == {"x": 1})
    check("seed marks is_seed + objective", specs[0]["context"].get("is_seed") is True
          and specs[0]["context"].get("objective") == "Build a calculator")
    # Manager-provided agents are honored (Manager-first; seeder does not override).
    check("Manager-provided agent breakdown bypasses the seeder",
          orch._manager_provided_agents({"orchestration_plan": {"agents": [{"role": "impl"}]}}) is True)


async def test_agent_peer_routing():
    from agent.main import Agent

    a = Agent(orchestrator_url="ws://localhost:8000", agent_id="agent-peer", capabilities=["python"])

    sent = []

    async def fake_send(payload):
        sent.append(payload)

    a.safe_send = fake_send

    # Inbound peer query (no matching waiter) -> agent answers with a reply.
    await a.handle_peer_message({
        "type": "peer_message",
        "from_agent_id": "agent-impl",
        "task_id": "t1",
        "content": "",  # empty -> deterministic fallback reply (no model needed)
        "correlation_id": "corr-1",
        "is_reply": False,
    })
    await asyncio.sleep(0.05)  # answer_peer_query is scheduled as a background task
    check("peer query produces a reply", len(sent) == 1 and sent[0]["is_reply"] is True)
    check("peer reply targets requester", sent[0]["target_agent_id"] == "agent-impl")

    # Inbound reply resolves an outstanding waiter (the requester side).
    loop = asyncio.get_running_loop()
    waiter = loop.create_future()
    a.peer_waiters["corr-2"] = waiter
    await a.handle_peer_message({
        "type": "peer_message",
        "from_agent_id": "agent-impl",
        "task_id": "t1",
        "content": "here is guidance",
        "correlation_id": "corr-2",
        "is_reply": True,
    })
    check("peer reply resolves waiter", waiter.done() and waiter.result()["content"] == "here is guidance")

    # Capacity rejection: fill active_tasks to the limit, then a new assignment is rejected.
    sent.clear()
    a.active_tasks = {f"t{i}": None for i in range(a.MAX_CONCURRENT_TASKS)}
    await a.handle_task_assignment({"task_id": "overflow", "description": "x", "context": {}})
    check("over-capacity assignment is rejected", len(sent) == 1 and sent[0]["type"] == "task_rejected")


async def test_review_and_synthesis():
    # planner.review heuristic rejects conflicting test frameworks.
    conflicting = {
        "swarm-r-testing": {"result": {"subtasks": [{"summary": "use pytest", "changes": "pytest suite"}]}},
        "swarm-r-doc": {"result": {"subtasks": [{"summary": "docs say unittest", "changes": "unittest examples"}]}},
    }
    verdict = await planner.review("obj", [], conflicting)
    check("review rejects conflicting frameworks", verdict["accepted"] is False and verdict["retry_tasks"])

    aligned = {"swarm-r-impl": {"result": {"subtasks": [{"summary": "clean implementation", "changes": "added add()"}]}}}
    verdict2 = await planner.review("obj", [], aligned)
    check("review accepts aligned results", verdict2["accepted"] is True)

    # synthesize falls back to a deterministic concatenation when no model is configured.
    synth = await planner.synthesize("obj", aligned)
    check("synthesize produces a non-empty response", isinstance(synth, str) and "implementation" in synth)


async def test_review_cycle():
    # Swarm review is peer cross-review (no single-critic Master): >=2 reviewers, one rejects →
    # reopen the flagged task. (Full coverage in test-swarm-cross-review.py.)
    orig_save, orig_emit = orch.swarm_runtime.save_run, orch.swarm_runtime.emit_event

    async def noop(*a, **k):
        return None

    orch.swarm_runtime.save_run = noop
    orch.swarm_runtime.emit_event = noop

    task = await task_queue.create_task(description="impl", task_id="t-rev-1")
    await task_queue.complete_task(task.task_id, '{"summary": "did impl"}')
    completed = await task_queue.get_task(task.task_id)
    run = types.SimpleNamespace(
        swarm_id="swarm-rev", deployment_id="dep-rev", manager_deployment_id="mgr-rev",
        objective="obj", task_ids=["t-rev-1"], status="completed",
        metadata={"reviews": [
            {"verdict": "pass", "reviewer_agent_id": "r1"},
            {"verdict": "fail", "reviewer_agent_id": "r2",
             "recommended_rework": ["t-rev-1"], "summary": "implementation incorrect"},
        ]},
    )
    try:
        reopened = await orch.run_cross_review(run, [completed])
        check("cross-review reopens on reviewer split (safety-biased reject)", reopened is True)
        reloaded = await task_queue.get_task("t-rev-1")
        check("reopened task is PENDING again", reloaded.status == TaskStatus.PENDING)
        check("review cycle counter incremented", run.metadata.get("review_cycles") == 1)
        check("run set back to running", run.status == "running")
        check("disagreement recorded", (run.metadata.get("cross_review") or {}).get("disagreement") is True)
    finally:
        orch.swarm_runtime.save_run = orig_save
        orch.swarm_runtime.emit_event = orig_emit


async def test_dispatch_context():
    # A dependency that completed should appear as a dependency artifact; a connected peer
    # assigned to another task on the run should appear as a peer agent.
    orch.manager.active_connections["peer-conn"] = object()
    try:
        dep = await task_queue.create_task(description="upstream", task_id="dc-dep")
        await task_queue.complete_task(dep.task_id, '{"summary": "upstream done", "files_modified": ["a.py"]}')
        peer = await task_queue.create_task(description="peer work", task_id="dc-peer", agent_role="implementation")
        await task_queue.assign_task(peer.task_id, "peer-conn") if False else None
        # Manually mark the peer task as owned by the connected agent.
        peer.assigned_agent_id = "peer-conn"
        await task_queue._save_task(peer)

        consumer = await task_queue.create_task(
            description="downstream", task_id="dc-main", depends_on=["dc-dep"], agent_role="testing"
        )
        run = types.SimpleNamespace(swarm_id="swarm-dc", objective="ship it", task_ids=["dc-dep", "dc-peer", "dc-main"])
        ctx = await orch.build_dispatch_context(run, consumer)
        check("dispatch context injects dependency_artifacts", len(ctx.get("dependency_artifacts", [])) == 1)
        check("dependency artifact carries summary", ctx["dependency_artifacts"][0]["summary"] == "upstream done")
        check("dispatch context injects peer_agents", any(p["agent_id"] == "peer-conn" for p in ctx.get("peer_agents", [])))
        check("dispatch context carries run goal", ctx.get("run_goal") == "ship it")
        check("dispatch context sets specialist_role from agent_role", ctx.get("specialist_role") == "testing")
    finally:
        orch.manager.active_connections.pop("peer-conn", None)


async def test_agent_peer_shares_summary():
    from agent.main import Agent

    a = Agent(orchestrator_url="ws://localhost:8000", agent_id="agent-impl2", capabilities=["python"])
    a.last_summary = "implemented add() that raises ValueError on bad input"
    sent = []

    async def fake_send(payload):
        sent.append(payload)

    a.safe_send = fake_send

    # Fallback path: no query content -> cached summary (no model needed).
    await a.answer_peer_query({"from_agent_id": "agent-test", "correlation_id": "c", "is_reply": False})
    check("peer reply falls back to last summary", "implemented add()" in sent[0]["content"])

    # Substantive path: a stubbed executor returns a grounded, structured reply.
    class _FakeExec:
        async def peer_reply(self, *, query, capabilities, last_summary, max_tokens=None):
            return {"content": f"grounded answer to: {query}", "stance": "agree",
                    "evidence": "src", "refs": ["a.py"]}

    a._peer_executor = _FakeExec()
    sent.clear()
    await a.answer_peer_query({
        "from_agent_id": "agent-test", "correlation_id": "c2", "task_id": "t",
        "content": "does add() raise ValueError?", "is_reply": False,
    })
    check("substantive peer reply uses executor content",
          sent[0]["content"] == "grounded answer to: does add() raise ValueError?"
          and sent[0]["stance"] == "agree" and sent[0]["refs"] == ["a.py"])


async def main():
    await test_redis_fallback()
    await test_release_task()
    await test_seeder()
    await test_agent_peer_routing()
    await test_review_and_synthesis()
    await test_review_cycle()
    await test_dispatch_context()
    await test_agent_peer_shares_summary()
    print()
    if failures:
        print(f"{len(failures)} check(s) FAILED: {failures}")
        sys.exit(1)
    print("all merge smoke checks passed")


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