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,
    ToolResult,
)
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",
    helper="claude-haiku-4-5",
    expert="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 [finish(call_id=f"fin-{len(self.calls)}")]
        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)


def finish(summary: str = "done", call_id: str = "fin") -> ToolCall:
    return ToolCall(call_id=call_id, tool="finish", arguments={"summary": summary})


@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")],
            [finish("wrote the file")],
        ]
    )
    outcome = build(client, store, ws).run(milestone=1)

    assert outcome.status == "completed"
    assert outcome.note == "wrote the file"
    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", "ToolCall", "ToolResult"]


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", "finish"]
    assert "milestone 2" in client.calls[0]["system"]


def test_prose_no_longer_ends_the_milestone(store: StateStore, ws: Workspace) -> None:
    # Only finish ends a milestone. Prose breaks the tool contract: the model
    # is told once, and here it gets back to work.
    client = ScriptedClient(
        [
            [call("write_file", path="a.py", content="x")],
            [AssistantText(text="all steps done", model="m")],
            [finish()],
        ]
    )
    assert build(client, store, ws).run(milestone=1).status == "completed"
    assert "Every reply must be a tool call" in client.calls[2]["system"]


def test_persisting_in_prose_escalates(store: StateStore, ws: Workspace) -> None:
    prose = [AssistantText(text="let me just chat", model="m")]
    client = ScriptedClient([prose, prose])
    outcome = build(client, store, ws).run(milestone=1)
    assert outcome.status == "escalate"
    assert "prose" in outcome.note


def test_finishing_without_any_work_is_refused_once(
    store: StateStore, ws: Workspace
) -> None:
    # A model that finishes on turn one has done nothing; believing it would
    # record an empty milestone as a success.
    client = ScriptedClient(
        [
            [finish("looks done to me", call_id="f1")],
            [call("write_file", path="a.py", content="x")],
            [finish("now it is done", call_id="f2")],
        ]
    )
    outcome = build(client, store, ws).run(milestone=1)
    assert outcome.status == "completed"
    assert (ws.root / "a.py").exists()  # the refusal produced actual work
    [refusal] = [
        e.event
        for e in store.read_events()
        if isinstance(e.event, ToolResult) and e.event.call_id == "f1"
    ]
    assert not refusal.success
    assert "not made any change yet" in refusal.result


def test_persisting_in_an_empty_finish_escalates(
    store: StateStore, ws: Workspace
) -> None:
    client = ScriptedClient([[finish(call_id="f1")], [finish(call_id="f2")]])
    outcome = build(client, store, ws).run(milestone=1)
    assert outcome.status == "escalate"
    assert "without making any change" in outcome.note


def test_calls_after_finish_are_not_run(store: StateStore, ws: Workspace) -> None:
    client = ScriptedClient(
        [
            [call("write_file", path="a.py", content="x")],
            [finish(), call("write_file", call_id="late", path="b.py", content="y")],
        ]
    )
    assert build(client, store, ws).run(milestone=1).status == "completed"
    assert not (ws.root / "b.py").exists()
    [late] = [
        e.event
        for e in store.read_events()
        if isinstance(e.event, ToolResult) and e.event.call_id == "late"
    ]
    assert "finish had already ended the milestone" in late.result


def test_an_invented_tool_counts_toward_escalation(
    store: StateStore, ws: Workspace
) -> None:
    # Measured: a model on a server that doesn't enforce tool_choice will
    # call tools that were never offered ("chat", "respond"). Without this it
    # would burn every turn up to the cap.
    client = ScriptedClient(
        [[call("chat", call_id="x1")], [call("respond_user", call_id="x2")]]
    )
    outcome = build(client, store, ws, failure_limit=2).run(milestone=1)
    assert outcome.status == "escalate"
    assert len(client.calls) == 2
    results = [e.event for e in store.read_events() if isinstance(e.event, ToolResult)]
    assert "no such tool: chat" in results[0].result
    assert "finish" in results[0].result  # it is told what does exist


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")],
            [finish()],
        ]
    )
    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")],
            [finish()],
        ]
    )
    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=".")],
            [finish("looks done to me", call_id="f1")],
            [finish("really, done", call_id="f2")],
        ]
    )
    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"),
            ],
            [finish()],
        ]
    )
    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", "fin"}


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 "ask_expert" in names
    assert names[-1] == "finish"
    # L9: 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")],
            [finish()],
        ]
    )
    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?")],
            [finish("I think it is done", call_id="f1")],
            [finish("really done", call_id="f2")],
        ]
    )
    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")],
            [finish()],
        ]
    )
    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