nisse/nisse-dev/tests/test_executor.py

import io
from pathlib import Path
from typing import Any

import pytest

from nisse_dev.canonical import AssistantText, CanonicalEvent, StoredEvent, ToolCall
from nisse_dev.executor import LLMExecutor
from nisse_dev.llm import LLMError, RoleMap
from nisse_dev.oplog import OpLog
from nisse_dev.state import StateStore
from nisse_dev.tools import Workspace

ROLES = RoleMap(
    executor="llama3.2",
    planner="claude-haiku-4-5",
    delegate="claude-haiku-4-5",
    consultant="claude-haiku-4-5",
)


class ScriptedClient:
    def __init__(self, turns: list[list[CanonicalEvent]]) -> None:
        self._turns = turns
        self.calls: list[dict[str, Any]] = []

    def complete(
        self,
        events: list[StoredEvent],
        tools: list[dict[str, Any]] | None,
        model: str,
        system: str | None = None,
    ) -> list[CanonicalEvent]:
        self.calls.append({"tools": tools, "model": model, "system": system})
        if not self._turns:
            return [AssistantText(text="done", model=model)]
        return self._turns.pop(0)


class FailingClient:
    def complete(self, *args: Any, **kwargs: Any) -> list[CanonicalEvent]:
        raise LLMError("proxy exploded")


def call(tool: str, call_id: str = "c1", **arguments: Any) -> ToolCall:
    return ToolCall(call_id=call_id, tool=tool, arguments=arguments)


@pytest.fixture
def store(tmp_path: Path) -> StateStore:
    return StateStore(tmp_path / "state.db", oplog=OpLog("nisse-dev", io.StringIO()))


@pytest.fixture
def ws(tmp_path: Path) -> Workspace:
    return Workspace(tmp_path / "work", timeout_s=10)


def build(client: Any, store: StateStore, ws: Workspace, **kw: Any) -> LLMExecutor:
    return LLMExecutor(client, store, ROLES, ws, **kw)


def test_tool_calls_are_executed_and_results_logged(
    store: StateStore, ws: Workspace
) -> None:
    client = ScriptedClient(
        [
            [call("write_file", path="app.py", content="print(1)\n")],
            [AssistantText(text="wrote the file", model="llama3.2")],
        ]
    )
    outcome = build(client, store, ws).run(milestone=1)

    assert outcome.status == "completed"
    assert (ws.root / "app.py").read_text() == "print(1)\n"
    kinds = [type(e.event).__name__ for e in store.read_events()]
    assert kinds == ["ToolCall", "ToolResult", "AssistantText"]


def test_uses_executor_role_and_reduced_manifest(
    store: StateStore, ws: Workspace
) -> None:
    client = ScriptedClient([])
    build(client, store, ws).run(milestone=2)
    assert client.calls[0]["model"] == "llama3.2"  # never the planner
    names = [t["function"]["name"] for t in client.calls[0]["tools"]]
    assert names == ["read_file", "write_file", "list_dir", "run_command"]
    assert "milestone 2" in client.calls[0]["system"]


def test_no_tool_call_after_work_ends_the_milestone(
    store: StateStore, ws: Workspace
) -> None:
    client = ScriptedClient(
        [
            [call("write_file", path="a.py", content="x")],
            [AssistantText(text="all steps done", model="m")],
        ]
    )
    assert build(client, store, ws).run(milestone=1).status == "completed"


def test_completion_claimed_without_any_work_is_challenged(
    store: StateStore, ws: Workspace
) -> None:
    # A model that answers with prose on turn one has done nothing; believing
    # it would record an empty milestone as a success.
    client = ScriptedClient(
        [
            [AssistantText(text="looks done to me", model="m")],
            [call("write_file", path="a.py", content="x")],
            [AssistantText(text="now it is done", model="m")],
        ]
    )
    outcome = build(client, store, ws).run(milestone=1)
    assert outcome.status == "completed"
    assert (ws.root / "a.py").exists()  # the nudge produced actual work
    assert "not made any change yet" in client.calls[1]["system"]


def test_persisting_in_an_empty_claim_escalates(
    store: StateStore, ws: Workspace
) -> None:
    prose = [AssistantText(text="nothing to do", model="m")]
    client = ScriptedClient([prose, prose])
    outcome = build(client, store, ws).run(milestone=1)
    assert outcome.status == "escalate"
    assert "without making any change" in outcome.note


def test_two_consecutive_command_failures_escalate(
    store: StateStore, ws: Workspace
) -> None:
    client = ScriptedClient(
        [
            [call("run_command", command="exit 1")],
            [call("run_command", command="exit 1")],
            [call("run_command", command="exit 1")],
        ]
    )
    outcome = build(client, store, ws, failure_limit=2).run(milestone=1)
    assert outcome.status == "escalate"
    assert "2 consecutive command failures" in outcome.note
    assert len(client.calls) == 2  # stopped, did not keep trying


def test_a_success_resets_the_failure_count(store: StateStore, ws: Workspace) -> None:
    client = ScriptedClient(
        [
            [call("run_command", command="exit 1")],
            [call("run_command", command="true")],
            [call("run_command", command="exit 1")],
            [AssistantText(text="done", model="m")],
        ]
    )
    assert build(client, store, ws, failure_limit=2).run(milestone=1).status == (
        "completed"
    )


def test_failed_reads_do_not_count_toward_escalation(
    store: StateStore, ws: Workspace
) -> None:
    # Two failed reads in a row must not trip the failure limit: looking for
    # a file that is not there is the model orienting itself, not failing.
    client = ScriptedClient(
        [
            [call("read_file", path="missing-a.txt")],
            [call("read_file", path="missing-b.txt")],
            [call("write_file", path="a.py", content="x")],
            [AssistantText(text="done", model="m")],
        ]
    )
    assert build(client, store, ws, failure_limit=2).run(milestone=1).status == (
        "completed"
    )


def test_reading_alone_is_not_doing_the_work(store: StateStore, ws: Workspace) -> None:
    # Exploration is not evidence a milestone was carried out, so a
    # completion claim backed only by reads is still challenged.
    client = ScriptedClient(
        [
            [call("list_dir", path=".")],
            [AssistantText(text="looks done to me", model="m")],
            [AssistantText(text="really, done", model="m")],
        ]
    )
    outcome = build(client, store, ws).run(milestone=1)
    assert outcome.status == "escalate"


def test_turn_cap_is_enforced(store: StateStore, ws: Workspace) -> None:
    forever = [[call("list_dir", path=".")] for _ in range(50)]
    client = ScriptedClient(forever)
    outcome = build(client, store, ws, max_turns=3).run(milestone=1)
    assert outcome.status == "capped"
    assert "3-turn cap" in outcome.note
    assert len(client.calls) == 3


def test_llm_error_escalates_rather_than_crashing(
    store: StateStore, ws: Workspace
) -> None:
    outcome = build(FailingClient(), store, ws).run(milestone=1)
    assert outcome.status == "escalate"
    assert "proxy exploded" in outcome.note


def test_multiple_calls_in_one_turn_all_run(store: StateStore, ws: Workspace) -> None:
    client = ScriptedClient(
        [
            [
                call("write_file", call_id="a", path="one.txt", content="1"),
                call("write_file", call_id="b", path="two.txt", content="2"),
            ],
            [AssistantText(text="done", model="m")],
        ]
    )
    build(client, store, ws).run(milestone=1)
    assert (ws.root / "one.txt").exists() and (ws.root / "two.txt").exists()
    results = [
        e.event for e in store.read_events() if type(e.event).__name__ == "ToolResult"
    ]
    assert {r.call_id for r in results} == {"a", "b"}


class StubDesk:
    """Stands in for ExpertDesk without a model behind it."""

    def __init__(self, answer: str = "use sqlite", error: Exception | None = None):
        self.asked: list[str] = []
        self._answer = answer
        self._error = error

    def ask(self, question: str, **kw: Any) -> Any:
        self.asked.append(question)
        if self._error is not None:
            raise self._error

        class _A:
            def render(self_inner) -> str:
                return self._answer

        return _A()


def test_ask_expert_is_offered_only_when_a_desk_exists(
    store: StateStore, ws: Workspace
) -> None:
    client = ScriptedClient([])
    build(client, store, ws).run(milestone=1)
    assert all(t["function"]["name"] != "ask_expert" for t in client.calls[0]["tools"])

    client2 = ScriptedClient([])
    build(client2, store, ws, expert=StubDesk()).run(milestone=1)
    names = [t["function"]["name"] for t in client2.calls[0]["tools"]]
    assert names[-1] == "ask_expert"
    # A12: in addition to the observable triggers, not instead of them.
    assert "run_command" in names


def test_ask_expert_answer_comes_back_as_a_tool_result(
    store: StateStore, ws: Workspace
) -> None:
    desk = StubDesk(answer="prefer sqlite here")
    client = ScriptedClient(
        [
            [call("ask_expert", question="sqlite or postgres?")],
            [call("write_file", path="a.py", content="x")],
            [AssistantText(text="done", model="m")],
        ]
    )
    build(client, store, ws, expert=desk).run(milestone=1)

    assert desk.asked == ["sqlite or postgres?"]
    results = [
        e.event for e in store.read_events() if type(e.event).__name__ == "ToolResult"
    ]
    assert "prefer sqlite here" in results[0].result


def test_consulting_is_not_doing_the_work(store: StateStore, ws: Workspace) -> None:
    # Asking is not changing anything, so a completion claim backed only by
    # consults is still challenged.
    client = ScriptedClient(
        [
            [call("ask_expert", question="what should I do?")],
            [AssistantText(text="I think it is done", model="m")],
            [AssistantText(text="really done", model="m")],
        ]
    )
    outcome = build(client, store, ws, expert=StubDesk()).run(milestone=1)
    assert outcome.status == "escalate"


def test_exhausted_budget_is_reported_to_the_model_not_raised(
    store: StateStore, ws: Workspace
) -> None:
    from nisse_dev.expert import BudgetExhausted

    desk = StubDesk(error=BudgetExhausted("consult budget of 3 is spent"))
    client = ScriptedClient(
        [
            [call("ask_expert", question="again?")],
            [call("write_file", path="a.py", content="x")],
            [AssistantText(text="done", model="m")],
        ]
    )
    outcome = build(client, store, ws, expert=desk).run(milestone=1)

    assert outcome.status == "completed"  # it got on with the work
    results = [
        e.event for e in store.read_events() if type(e.event).__name__ == "ToolResult"
    ]
    assert not results[0].success
    assert "budget" in results[0].result