src/tests/test_conductor.py

"""
Frox AI — Conductor Tests
Run with: pytest tests/test_conductor.py -v

Covers the pure-logic pieces (plan parsing, worker validation) that
don't need torch or a loaded model — the orchestrator/engine-calling
parts are integration-level and better tested against a real server.
"""
from __future__ import annotations

import sys
from pathlib import Path

import pytest

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

from orchestration.conductor import MorphConductor, WorkerSpec, WorkflowStep


class FakeOrchestrator:
    """Stands in for MorphInferenceEngine — returns a canned response per call."""
    def __init__(self, responses):
        self.responses = list(responses)
        self.calls = []

    def generate(self, messages, **kwargs):
        self.calls.append((messages, kwargs))
        return self.responses.pop(0) if self.responses else ""


@pytest.fixture
def conductor_with_workers():
    orch = FakeOrchestrator([])
    workers = [
        WorkerSpec("code", "Coding specialist"),
        WorkerSpec("pro", "Hard reasoning specialist"),
    ]
    return MorphConductor(orchestrator=orch, workers=workers), orch


class TestPlanParsing:
    def test_valid_plan_parses(self, conductor_with_workers):
        conductor, _ = conductor_with_workers
        raw = '{"steps": [{"subtask": "write code", "worker": "code", "access": []}]}'
        plan = conductor._parse_plan(raw)
        assert plan.valid
        assert len(plan.steps) == 1
        assert plan.steps[0].worker == "code"

    def test_plan_wrapped_in_prose_still_parses(self, conductor_with_workers):
        """Models often add commentary around the JSON despite instructions — must still extract it."""
        conductor, _ = conductor_with_workers
        raw = 'Here is my plan:\n{"steps": [{"subtask": "x", "worker": "pro", "access": []}]}\nHope that helps!'
        plan = conductor._parse_plan(raw)
        assert plan.valid
        assert plan.steps[0].worker == "pro"

    def test_multi_step_plan_with_access_list(self, conductor_with_workers):
        conductor, _ = conductor_with_workers
        raw = ('{"steps": ['
               '{"subtask": "draft code", "worker": "code", "access": []},'
               '{"subtask": "verify correctness", "worker": "pro", "access": [0]}'
               ']}')
        plan = conductor._parse_plan(raw)
        assert plan.valid
        assert len(plan.steps) == 2
        assert plan.steps[1].access == [0]

    def test_unparseable_json_is_invalid(self, conductor_with_workers):
        conductor, _ = conductor_with_workers
        plan = conductor._parse_plan("I don't want to make a plan.")
        assert not plan.valid
        assert plan.error is not None

    def test_empty_steps_is_invalid(self, conductor_with_workers):
        conductor, _ = conductor_with_workers
        plan = conductor._parse_plan('{"steps": []}')
        assert not plan.valid

    def test_unknown_worker_is_invalid(self, conductor_with_workers):
        conductor, _ = conductor_with_workers
        raw = '{"steps": [{"subtask": "x", "worker": "nonexistent_worker", "access": []}]}'
        plan = conductor._parse_plan(raw)
        assert not plan.valid
        assert "nonexistent_worker" in plan.error

    def test_self_worker_is_valid(self, conductor_with_workers):
        """'self' (the orchestrator reviewing/refining its own work) is always a valid worker name."""
        conductor, _ = conductor_with_workers
        raw = '{"steps": [{"subtask": "x", "worker": "self", "access": []}]}'
        plan = conductor._parse_plan(raw)
        assert plan.valid

    def test_max_steps_truncates(self, conductor_with_workers):
        conductor, _ = conductor_with_workers
        conductor.max_steps = 2
        raw = json_steps = (
            '{"steps": ['
            '{"subtask": "a", "worker": "code", "access": []},'
            '{"subtask": "b", "worker": "pro", "access": []},'
            '{"subtask": "c", "worker": "code", "access": []}'
            ']}'
        )
        plan = conductor._parse_plan(raw)
        assert plan.valid
        assert len(plan.steps) == 2


class TestFastModeFallback:
    def test_fast_mode_picks_a_configured_worker(self, conductor_with_workers):
        conductor, orch = conductor_with_workers
        orch.responses = ["code"]   # orchestrator's routing decision
        # Patch the worker to avoid actually invoking a model
        conductor.workers["code"].engine = FakeOrchestrator(["worker response"])
        result = conductor.run([{"role": "user", "content": "write a function"}], mode="fast")
        assert result.mode == "fast"
        assert result.trace[0]["worker"] == "code"

    def test_invalid_plan_falls_back_gracefully(self, conductor_with_workers):
        conductor, orch = conductor_with_workers
        # First call: _plan() -> garbage (invalid). Second call: _select_single_worker() -> "pro".
        orch.responses = ["not json at all", "pro"]
        conductor.workers["pro"].engine = FakeOrchestrator(["fallback response"])
        result = conductor.run([{"role": "user", "content": "hard task"}], mode="deep")
        assert result.mode == "fallback"
        assert "note" in result.trace[0]


class TestWorkerSpec:
    def test_invoke_requires_engine_or_call_fn(self):
        worker = WorkerSpec("orphan", "no backend configured")
        with pytest.raises(RuntimeError):
            worker.invoke("do something")

    def test_invoke_uses_call_fn_when_present(self):
        calls = []
        def fake_call(system_prompt, content):
            calls.append((system_prompt, content))
            return "remote response"

        worker = WorkerSpec("remote", "a remote worker", call_fn=fake_call, role_prompt="be helpful")
        result = worker.invoke("do the thing", context="prior context")
        assert result == "remote response"
        assert calls[0][0] == "be helpful"
        assert "prior context" in calls[0][1]
        assert "do the thing" in calls[0][1]