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_ends_the_milestone(store: StateStore, ws: Workspace) -> None:
    client = ScriptedClient([[AssistantText(text="all steps done", model="m")]])
    assert build(client, store, ws).run(milestone=1).status == "completed"


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:
    client = ScriptedClient(
        [
            [call("read_file", path="missing-a.txt")],
            [call("read_file", path="missing-b.txt")],
            [AssistantText(text="looked around, done", model="m")],
        ]
    )
    assert build(client, store, ws, failure_limit=2).run(milestone=1).status == (
        "completed"
    )


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"}