nisse/nisse-dev/tests/test_stop.py

"""The operator's kill switch, `stop` (spec A4)."""

import io
import threading
from pathlib import Path
from typing import Any

from nisse_dev.cancel import STOPPED_RESULT, Cancel
from nisse_dev.canonical import AssistantText, ToolCall, ToolResult, UserMessage
from nisse_dev.control import (
    AWAIT_ARCH_SIGNOFF,
    AWAIT_INPUT,
    AWAIT_MILESTONE_SIGNOFF,
    Orchestrator,
    PlanningStopped,
    StubAgent,
    SubmitPlan,
)
from nisse_dev.delegation import DelegationResult
from nisse_dev.executor import ExecutorOutcome, LLMExecutor
from nisse_dev.llm import RoleMap
from nisse_dev.oplog import OpLog
from nisse_dev.planner import LLMPlanningAgent
from nisse_dev.state import StateStore
from nisse_dev.tools import ToolOutcome
from nisse_dev.transport import InMemoryTransport
from nisse_dev.worker import SerialWorker

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


def new_store(path: Path) -> StateStore:
    return StateStore(path / "state.db", oplog=OpLog("nisse-dev", io.StringIO()))


def user_messages(store: StateStore) -> list[str]:
    return [
        e.event.text for e in store.read_events() if isinstance(e.event, UserMessage)
    ]


# --- the orchestrator ------------------------------------------------------


class OperatorDuringWork:
    """An executor that hears from the operator mid-milestone, then does what
    the real loop does at its next tool boundary."""

    def __init__(self, cancel: Cancel, says: str) -> None:
        self.cancel = cancel
        self.says = says
        self.transport: InMemoryTransport | None = None

    def run(self, milestone: int) -> ExecutorOutcome:
        assert self.transport is not None
        self.transport.receive("don", self.says)
        if self.cancel.requested:
            return ExecutorOutcome("stopped", "during turn 3")
        return ExecutorOutcome("completed", "done")


def at_start_of_milestone(
    tmp_path: Path, says: str
) -> tuple[StateStore, InMemoryTransport, Cancel]:
    store = new_store(tmp_path)
    transport = InMemoryTransport()
    cancel = Cancel()
    executor = OperatorDuringWork(cancel, says)
    executor.transport = transport
    orch = Orchestrator(
        store,
        transport,
        StubAgent(),
        operators={"don"},
        executor=executor,
        cancel=cancel,
    )
    orch.start()
    for text in ("use demo", "approve", "approve"):
        transport.receive("don", text)
    return store, transport, cancel


def test_stop_halts_a_running_milestone_and_asks_why(tmp_path: Path) -> None:
    store, transport, cancel = at_start_of_milestone(tmp_path, "stop")
    transport.receive("don", "start milestone 1")

    assert store.get_state().state == AWAIT_INPUT
    prompt = transport.outbound[-1]
    assert prompt.kind == "question"
    assert "stopped milestone 1 (during turn 3)" in prompt.text
    assert "Why did you stop it?" in prompt.text
    assert user_messages(store)[-1] == "stop"
    assert not cancel.requested  # the next piece of work is not pre-stopped


def test_the_reason_is_recorded_and_the_plan_stays_approved(tmp_path: Path) -> None:
    store, transport, _ = at_start_of_milestone(tmp_path, "stop")
    transport.receive("don", "start milestone 1")
    transport.receive("don", "it was editing the wrong module")

    assert store.get_state().state == AWAIT_MILESTONE_SIGNOFF
    # The next run renders this, so the model hears why it was stopped.
    assert user_messages(store)[-1] == "it was editing the wrong module"
    assert "start milestone 1" in transport.outbound[-1].text


def test_stop_with_nothing_running_says_so_and_changes_nothing(
    tmp_path: Path,
) -> None:
    store = new_store(tmp_path)
    transport = InMemoryTransport()
    orch = Orchestrator(store, transport, StubAgent(), operators={"don"})
    orch.start()
    transport.receive("don", "use demo")
    transport.receive("don", "stop")

    # At a gate, "stop" is the command, not rejection feedback.
    assert store.get_state().state == AWAIT_ARCH_SIGNOFF
    assert orch._data["revs"]["arch"] == 1
    assert (
        transport.outbound[-1].text == f"nothing to stop (state={AWAIT_ARCH_SIGNOFF})"
    )


def test_stop_from_a_non_operator_is_ignored(tmp_path: Path) -> None:
    _, transport, cancel = at_start_of_milestone(tmp_path, "approve")
    before = len(transport.outbound)
    transport.receive("mallory", "stop")
    assert len(transport.outbound) == before
    assert not cancel.requested


def test_other_messages_during_work_are_answered_not_acted_on(
    tmp_path: Path,
) -> None:
    store, transport, _ = at_start_of_milestone(tmp_path, "approve")
    transport.receive("don", "start milestone 1")

    texts = [m.text for m in transport.outbound]
    assert "busy (EXECUTING); say 'stop' to halt" in texts
    # Recorded, as every operator message is, but not applied to anything.
    assert user_messages(store)[-1] == "approve"


class StopsBeforePlanning:
    def __init__(self, cancel: Cancel) -> None:
        self.cancel = cancel
        self.transport: InMemoryTransport | None = None
        self.answers: list[str | None] = []

    def plan(
        self,
        gate: str,
        project: str,
        feedback: str | None = None,
        answer: str | None = None,
    ) -> SubmitPlan:
        self.answers.append(answer)
        if len(self.answers) == 1:
            assert self.transport is not None
            self.transport.receive("don", "stop")
            if self.cancel.requested:
                raise PlanningStopped("stopped by the operator")
        return SubmitPlan(f"{gate} plan, after: {answer}")


def test_stop_during_planning_resumes_planning_with_the_reason(
    tmp_path: Path,
) -> None:
    store = new_store(tmp_path)
    transport = InMemoryTransport()
    cancel = Cancel()
    planner = StopsBeforePlanning(cancel)
    planner.transport = transport
    orch = Orchestrator(store, transport, planner, operators={"don"}, cancel=cancel)
    orch.start()
    transport.receive("don", "use demo")

    assert store.get_state().state == AWAIT_INPUT
    assert "stopped planning arch" in transport.outbound[-1].text

    transport.receive("don", "keep it to one module")
    assert store.get_state().state == AWAIT_ARCH_SIGNOFF
    assert planner.answers[-1] == "keep it to one module"


# --- on real threads -------------------------------------------------------


class BlocksUntilStopped:
    def __init__(self, cancel: Cancel) -> None:
        self.cancel = cancel
        self.running = threading.Event()

    def run(self, milestone: int) -> ExecutorOutcome:
        self.running.set()
        for _ in range(500):
            if self.cancel.requested:
                return ExecutorOutcome("stopped", "at a tool boundary")
            threading.Event().wait(0.01)
        return ExecutorOutcome("completed", "never stopped")


def test_stop_reaches_work_running_on_the_worker(tmp_path: Path) -> None:
    # The bug this design fixes: the work used to run on the transport's own
    # thread, so no message, stop included, was read until it finished.
    worker = SerialWorker()
    worker.start()
    store = worker.call(lambda: new_store(tmp_path))
    transport = InMemoryTransport()
    cancel = Cancel()
    executor = BlocksUntilStopped(cancel)
    orch = Orchestrator(
        store,
        transport,
        StubAgent(),
        operators={"don"},
        executor=executor,
        worker=worker,
        cancel=cancel,
    )
    worker.call(orch.start)
    for text in ("use demo", "approve", "approve", "start milestone 1"):
        transport.receive("don", text)
    assert executor.running.wait(5)

    transport.receive("don", "status")  # delivered while the milestone runs
    transport.receive("don", "stop")
    worker.call(lambda: None)  # wait for the stop to be recorded

    texts = [m.text for m in transport.outbound]
    assert "busy (EXECUTING); say 'stop' to halt" in texts
    assert worker.call(store.get_state).state == AWAIT_INPUT
    assert "stopped milestone 1 (at a tool boundary)" in texts[-1]
    worker.call(store.close)
    worker.stop(timeout=5)


# --- the loops -------------------------------------------------------------


class ScriptedClient:
    def __init__(self, turns: list[list[Any]], on_call: Any = None) -> None:
        self.turns = turns
        self.on_call = on_call
        self.calls = 0

    def complete(self, events: Any, tools: Any, model: str, system: Any = None):
        self.calls += 1
        if self.on_call:
            self.on_call()
        return self.turns.pop(0) if self.turns else [AssistantText("done", model)]


class StopsAfterFirstCall:
    """A workspace whose first tool call is where the operator says stop."""

    def __init__(self, cancel: Cancel) -> None:
        self.cancel = cancel
        self.ran: list[str] = []

    def describe(self) -> str:
        return "workspace"

    def dispatch(self, tool: str, arguments: dict[str, Any]) -> ToolOutcome:
        self.ran.append(tool)
        self.cancel.set()
        return ToolOutcome("ok", success=True)


def test_executor_stops_before_a_turn_without_calling_the_model(
    tmp_path: Path,
) -> None:
    cancel = Cancel()
    cancel.set()
    client = ScriptedClient([])
    executor = LLMExecutor(
        client, new_store(tmp_path), ROLES, StopsAfterFirstCall(cancel), cancel=cancel
    )
    outcome = executor.run(1)
    assert outcome == ExecutorOutcome("stopped", "before turn 1")
    assert client.calls == 0


def test_executor_stops_between_tool_calls_and_records_the_skipped_ones(
    tmp_path: Path,
) -> None:
    cancel = Cancel()
    store = new_store(tmp_path)
    ws = StopsAfterFirstCall(cancel)
    calls = [
        ToolCall(call_id="a", tool="run_command", arguments={"command": "make"}),
        ToolCall(call_id="b", tool="write_file", arguments={"path": "x"}),
    ]
    executor = LLMExecutor(ScriptedClient([calls]), store, ROLES, ws, cancel=cancel)
    outcome = executor.run(1)

    assert outcome == ExecutorOutcome("stopped", "during turn 1")
    assert ws.ran == ["run_command"]
    results = {
        e.event.call_id: e.event.result
        for e in store.read_events()
        if isinstance(e.event, ToolResult)
    }
    # Every call has a result, so the transcript still renders.
    assert results == {"a": "ok", "b": STOPPED_RESULT}


class StoppedDelegate:
    def run(self, *, brief: str, branch: str) -> DelegationResult:
        return DelegationResult("stopped", "stopped by the operator during turn 2")


def test_a_stopped_delegation_stops_the_milestone(tmp_path: Path) -> None:
    failing = [
        [ToolCall(call_id=f"c{n}", tool="run_command", arguments={"command": "x"})]
        for n in range(2)
    ]

    class Failing:
        def describe(self) -> str:
            return "workspace"

        def dispatch(self, tool: str, arguments: dict[str, Any]) -> ToolOutcome:
            return ToolOutcome("exit 1", success=False)

    executor = LLMExecutor(
        ScriptedClient(failing),
        new_store(tmp_path),
        ROLES,
        Failing(),
        delegate=StoppedDelegate(),
    )
    outcome = executor.run(1)
    assert outcome.status == "stopped"
    assert "during a delegation" in outcome.note


def test_planner_does_not_start_a_turn_once_stopped(tmp_path: Path) -> None:
    cancel = Cancel()
    cancel.set()
    client = ScriptedClient([])
    planner = LLMPlanningAgent(client, new_store(tmp_path), ROLES, cancel=cancel)
    try:
        planner.plan("arch", "demo")
    except PlanningStopped:
        pass
    else:
        raise AssertionError("planning ran after stop")
    assert client.calls == 0