Agentswarm/orchestrator/test_multi_agent_workflow.py

"""Regression tests for the multi-agent DAG workflow."""
import fnmatch
import os
import unittest
from types import SimpleNamespace
from unittest.mock import patch

from orchestrator import main as orchestrator_main
from orchestrator import swarm_runtime as swarm_runtime_module
from orchestrator import task_queue as task_queue_module


class FakeRedisClient:
    """Tiny in-memory async Redis replacement for task-queue tests."""

    def __init__(self):
        self.kv = {}
        self.lists = {}

    async def set(self, key, value, ex=None):
        self.kv[key] = value

    async def get(self, key):
        return self.kv.get(key)

    async def incr(self, key):
        self.kv[key] = int(self.kv.get(key, 0)) + 1
        return self.kv[key]

    async def delete(self, key):
        self.kv.pop(key, None)
        self.lists.pop(key, None)

    async def keys(self, pattern):
        keys = list(self.kv.keys()) + list(self.lists.keys())
        return sorted([key for key in keys if fnmatch.fnmatch(key, pattern)])

    async def lpush(self, key, *values):
        self.lists.setdefault(key, [])
        for value in values:
            self.lists[key].insert(0, value)

    async def rpush(self, key, *values):
        self.lists.setdefault(key, [])
        self.lists[key].extend(values)

    async def rpop(self, key):
        items = self.lists.get(key, [])
        if not items:
            return None
        return items.pop()

    async def llen(self, key):
        return len(self.lists.get(key, []))

    async def lrange(self, key, start, end):
        items = list(self.lists.get(key, []))
        if end == -1:
            end = len(items) - 1
        return items[start : end + 1]

    async def lrem(self, key, count, value):
        items = self.lists.get(key, [])
        removed = 0
        kept = []
        for item in items:
            if item == value and (count == 0 or removed < count):
                removed += 1
                continue
            kept.append(item)
        self.lists[key] = kept
        return removed


class MultiAgentWorkflowTests(unittest.IsolatedAsyncioTestCase):
    """Covers DAG task creation, dependency dispatch, and parent-child resolution."""

    def setUp(self):
        self.fake_redis = FakeRedisClient()
        self.redis_patch = patch.object(task_queue_module, "redis_client", self.fake_redis)
        self.redis_patch.start()
        self.addCleanup(self.redis_patch.stop)

    async def test_build_task_descriptions_respects_feature_flag(self):
        body = {
            "orchestration_plan": {
                "objective": "Implement a feature",
                "sub_mode": "code",
                "risk_level": "low",
                "budget": {"duration_seconds": 600, "token_limit": 1000},
                "agents": [
                    {
                        "task_id": "backend-1",
                        "role": "backend",
                        "title": "Backend task",
                        "description": "Implement backend changes",
                        "depends_on": [],
                    },
                    {
                        "task_id": "tests-1",
                        "role": "testing",
                        "title": "Testing task",
                        "description": "Add tests",
                        "depends_on": ["backend-1"],
                    },
                ],
            }
        }

        runtime = swarm_runtime_module.SwarmRuntime()
        with patch.dict(os.environ, {"ENABLE_SUBTASK_HANDOFF": "false"}, clear=False):
            tasks = runtime.build_task_descriptions(body)
            self.assertEqual(len(tasks), 1)
            self.assertEqual(tasks[0]["workflow_mode"], "single_agent")

        with patch.dict(os.environ, {"ENABLE_SUBTASK_HANDOFF": "true"}, clear=False):
            tasks = runtime.build_task_descriptions(body)
            self.assertEqual(len(tasks), 2)
            self.assertEqual(tasks[0]["task_id"], "backend-1")
            self.assertEqual(tasks[1]["depends_on"], ["backend-1"])
            self.assertEqual(tasks[1]["workflow_mode"], "multi_agent")

    async def test_ready_pending_task_respects_dependencies_and_capabilities(self):
        parent = await task_queue_module.task_queue.create_task(
            task_id="parent-1",
            title="Backend",
            description="Implement backend",
            agent_role="backend",
            required_capabilities=["backend"],
            source="runtime_bridge",
        )
        child = await task_queue_module.task_queue.create_task(
            task_id="child-1",
            title="Tests",
            description="Add tests",
            agent_role="testing",
            required_capabilities=["testing"],
            depends_on=[parent.task_id],
            root_task_id=parent.task_id,
            source="runtime_bridge",
        )

        task = await task_queue_module.task_queue.get_ready_pending_task(["testing"])
        self.assertIsNone(task, "dependent task should not dispatch before parent completion")

        task = await task_queue_module.task_queue.get_ready_pending_task(["backend"])
        self.assertIsNotNone(task)
        self.assertEqual(task.task_id, parent.task_id)

        await task_queue_module.task_queue.complete_task(parent.task_id, "done")
        task = await task_queue_module.task_queue.get_ready_pending_task(["testing"])
        self.assertIsNotNone(task)
        self.assertEqual(task.task_id, child.task_id)

    async def test_finalize_parent_after_child_success(self):
        event_calls = []

        async def fake_emit_event(*args, **kwargs):
            event_calls.append((args, kwargs))

        runtime_patch = patch.object(orchestrator_main.swarm_runtime, "emit_event", fake_emit_event)
        runtime_patch.start()
        self.addCleanup(runtime_patch.stop)

        parent = await task_queue_module.task_queue.create_task(
            task_id="parent-1",
            description="Parent task",
            title="Parent",
            agent_role="backend",
            source="runtime_bridge",
        )
        child = await task_queue_module.task_queue.create_task(
            task_id="child-1",
            description="Child task",
            title="Child",
            agent_role="testing",
            parent_task_id=parent.task_id,
            root_task_id=parent.task_id,
            source="dynamic_handoff",
        )
        await task_queue_module.task_queue.add_child_task(parent.task_id, child.task_id)
        await task_queue_module.task_queue.block_task(parent.task_id, "Waiting on child")
        await task_queue_module.task_queue.complete_task(child.task_id, "child complete")

        refreshed_child = await task_queue_module.task_queue.get_task(child.task_id)
        run = SimpleNamespace()
        await orchestrator_main.finalize_parent_after_child(
            run,
            refreshed_child,
            agent_id="agent-1",
            success=True,
            summary="Delegated child task completed",
        )

        refreshed_parent = await task_queue_module.task_queue.get_task(parent.task_id)
        self.assertEqual(refreshed_parent.status, task_queue_module.TaskStatus.COMPLETED)
        self.assertGreaterEqual(len(event_calls), 2)

    async def test_finalize_parent_after_child_failure_is_terminal(self):
        async def fake_emit_event(*args, **kwargs):
            return None

        runtime_patch = patch.object(orchestrator_main.swarm_runtime, "emit_event", fake_emit_event)
        runtime_patch.start()
        self.addCleanup(runtime_patch.stop)

        parent = await task_queue_module.task_queue.create_task(
            task_id="parent-2",
            description="Parent task",
            title="Parent",
            agent_role="backend",
            source="runtime_bridge",
            max_retries=3,
        )
        child = await task_queue_module.task_queue.create_task(
            task_id="child-2",
            description="Child task",
            title="Child",
            agent_role="testing",
            parent_task_id=parent.task_id,
            root_task_id=parent.task_id,
            source="dynamic_handoff",
            max_retries=1,
        )
        await task_queue_module.task_queue.add_child_task(parent.task_id, child.task_id)
        await task_queue_module.task_queue.block_task(parent.task_id, "Waiting on child")
        await task_queue_module.task_queue.fail_task(child.task_id, "child failed")

        refreshed_child = await task_queue_module.task_queue.get_task(child.task_id)
        run = SimpleNamespace()
        await orchestrator_main.finalize_parent_after_child(
            run,
            refreshed_child,
            agent_id="agent-1",
            success=False,
            summary="child failed",
        )

        refreshed_parent = await task_queue_module.task_queue.get_task(parent.task_id)
        self.assertEqual(refreshed_parent.status, task_queue_module.TaskStatus.FAILED)

    async def test_build_result_artifact_emits_document_without_git_branch(self):
        run = SimpleNamespace(
            deployment_id="runtime-dep-1",
            manager_deployment_id="dep-manager-1",
            swarm_id="swarm-1",
        )
        task = SimpleNamespace(
            task_id="task-1",
            title="Backend implementation",
            agent_role="backend",
            parent_task_id=None,
            root_task_id="task-1",
            context={},
        )
        result = {
            "summary": "Implemented backend endpoints and startup notes",
            "files_modified": ["backend/app.py", "README.md"],
            "changes": "Added handlers and docs",
            "git_skipped": "Workspace is not a Git checkout",
        }

        artifact = orchestrator_main.build_result_artifact(run, task, result)

        self.assertEqual(artifact["artifact_type"], "deployment_manifest")
        self.assertEqual(artifact["uri"], "runtime://swarm-1/artifacts/task-1")
        self.assertEqual(
            artifact["metadata"]["files_modified"],
            ["backend/app.py", "README.md"],
        )

    async def test_artifact_created_event_copies_artifact_into_payload(self):
        fake_runtime_redis = FakeRedisClient()
        runtime_patch = patch.object(
            swarm_runtime_module,
            "redis_client",
            SimpleNamespace(
                set=fake_runtime_redis.set,
                get=fake_runtime_redis.get,
                keys=fake_runtime_redis.keys,
                rpush=fake_runtime_redis.rpush,
                lrange=fake_runtime_redis.lrange,
            ),
        )
        runtime_patch.start()
        self.addCleanup(runtime_patch.stop)

        runtime = swarm_runtime_module.SwarmRuntime()
        run = swarm_runtime_module.SwarmRun(
            deployment_id="runtime-dep-1",
            swarm_id="swarm-1",
            status="running",
            objective="Test callback payload",
            manager_deployment_id="dep-manager-1",
        )

        await runtime.emit_event(
            run,
            "artifact.created",
            task_id="task-1",
            artifact={
                "artifact_id": "art_task-1",
                "artifact_type": "document",
                "title": "Task result",
                "summary": "Generated a result summary",
                "uri": "runtime://swarm-1/artifacts/task-1",
            },
        )

        logs = await runtime.list_events(run.swarm_id)
        event = logs["events"][-1]
        self.assertEqual(event["payload"]["artifact_id"], "art_task-1")
        self.assertEqual(event["payload"]["artifact_type"], "document")
        self.assertEqual(event["artifact"]["uri"], "runtime://swarm-1/artifacts/task-1")

    async def test_validate_create_request_accepts_new_swarm_shape(self):
        runtime = swarm_runtime_module.SwarmRuntime()
        runtime.validate_create_request({
            "mode": "swarm",
            "conversation_id": "conv-1",
            "requirement": {
                "objective": "Build a swarm-delivered feature",
                "context": ["ticket-123"],
                "attachments": [],
                "constraints": ["keep existing API"],
                "acceptance_criteria": ["tests pass"],
            },
            "model_selection": {
                "type": "primary",
                "primary_model": "gpt-5.4",
            },
            "metadata": {
                "manager_deployment_id": "dep-manager-1",
                "correlation_id": "corr-1",
            },
            "callback": {
                "url": "https://manager.example/api/agent/callbacks/runtime-events",
            },
        })

    async def test_validate_create_request_rejects_missing_objective(self):
        runtime = swarm_runtime_module.SwarmRuntime()
        with self.assertRaises(swarm_runtime_module.RuntimeValidationError):
            runtime.validate_create_request({
                "mode": "swarm",
                "requirement": {},
                "metadata": {"manager_deployment_id": "dep-manager-1"},
                "callback": {"url": "https://manager.example/callback"},
            })

    async def test_get_run_by_identifier_supports_manager_deployment_id(self):
        fake_runtime_redis = FakeRedisClient()
        runtime_patch = patch.object(
            swarm_runtime_module,
            "redis_client",
            SimpleNamespace(
                set=fake_runtime_redis.set,
                get=fake_runtime_redis.get,
                keys=fake_runtime_redis.keys,
                rpush=fake_runtime_redis.rpush,
                lrange=fake_runtime_redis.lrange,
            ),
        )
        runtime_patch.start()
        self.addCleanup(runtime_patch.stop)

        runtime = swarm_runtime_module.SwarmRuntime()
        run = swarm_runtime_module.SwarmRun(
            deployment_id="runtime-dep-1",
            swarm_id="swarm-1",
            status="running",
            objective="Test identifier lookup",
            manager_deployment_id="dep-manager-1",
        )
        await runtime.save_run(run)

        loaded = await runtime.get_run_by_identifier("dep-manager-1")
        self.assertIsNotNone(loaded)
        self.assertEqual(loaded.swarm_id, "swarm-1")

    async def test_runtime_routes_include_agent_swarm_paths(self):
        route_paths = {route.path for route in orchestrator_main.app.routes}
        self.assertIn("/api/agent/health", route_paths)
        self.assertIn("/api/agent/swarm/deployments", route_paths)
        self.assertIn("/api/agent/swarm/deployments/{deployment_id}/events", route_paths)
        self.assertIn("/api/agent/swarm/deployments/{deployment_id}/workflow", route_paths)
        self.assertIn("/api/agent/swarm/deployments/{deployment_id}/diagnostics", route_paths)

    async def test_build_deliverable_fact_marks_summary_only_without_diff(self):
        result = {"summary": "Only a plan document"}
        deliverable = orchestrator_main.build_deliverable_fact(result)
        self.assertFalse(deliverable["has_deliverable"])
        self.assertTrue(deliverable["summary_only"])
        self.assertFalse(deliverable["has_diff"])

    async def test_build_workflow_phases_returns_fixed_phase_order(self):
        run = SimpleNamespace(status="running", approvals={})
        task = SimpleNamespace(
            task_id="task-1",
            title="Backend",
            agent_role="backend",
            status=SimpleNamespace(value="in_progress"),
            started_at=10,
            completed_at=None,
            result=None,
        )
        phases = orchestrator_main.build_workflow_phases(
            run,
            [task],
            [{"event_type": "task.created", "payload": {}}],
        )
        self.assertEqual(
            [phase["name"] for phase in phases],
            ["Plan", "Dispatch", "Execute", "Handoff", "Review", "Deliver"],
        )


if __name__ == "__main__":
    unittest.main()