nisse/nisse-dev/src/nisse_dev/loop/driver.py

"""A driver run (L2): one model with tools, looping until it calls a waiting
tool, runs out of turns, or is stopped.

Everything the run does is written to the item's transcript as the driver
(L7), and every model call sees the driver's view of it (L6). Execution tools
go to the build VM (L13); control tools are answered here, and a waiting one
ends the run, so the scheduler can put it to the operator.
"""

from collections.abc import Callable, Mapping
from dataclasses import dataclass
from typing import Any, Literal

from nisse_dev.cancel import Cancel, skip_calls
from nisse_dev.canonical import AssistantText, Fact, ToolCall, ToolResult
from nisse_dev.llm import LLMClient, LLMError
from nisse_dev.loop.tools import WAITING, names
from nisse_dev.state import Transcript
from nisse_dev.tools import Workspace

DEFAULT_MAX_TURNS = 30

# A reply in prose breaks the tool contract; the model is told once.
_PROSE_NUDGE = (
    "Your last reply was prose. Every reply must be a tool call: use one of "
    "your tools, or ask_human if you need the operator."
)


@dataclass(frozen=True)
class RunOutcome:
    # stuck: failed commands in a row hit the failure limit (L9's trigger)
    kind: Literal["wait", "capped", "stuck", "stopped", "failed"]
    tool: str = ""  # the waiting tool called
    arguments: dict[str, Any] | None = None
    note: str = ""


# The waiting outcome a dangerous action ends a run with (P3): the harness
# puts it to the operator, rendered from the action itself.
ASK_ALLOW = "ask_allow"


@dataclass(frozen=True)
class Screening:
    """What the fast path says about an execution call before it runs:
    ``run`` it, ``refuse`` it (the result says why, and the run goes on), or
    ``ask`` the operator (the run ends, waiting)."""

    decision: Literal["run", "refuse", "ask"]
    text: str = ""  # the call's result when not run; for ask, the question
    key: str = ""  # for ask: what an allow would bind to (L3)
    note: str = ""  # for run: what the harness adds to the call's result


class Driver:
    def __init__(
        self,
        client: LLMClient,
        model: str,
        *,
        max_turns: int = DEFAULT_MAX_TURNS,
        cancel: Cancel | None = None,
        failure_limit: int = 0,  # 0: no limit
    ) -> None:
        self._client = client
        self._model = model
        self._max_turns = max_turns
        self._failure_limit = failure_limit
        self._cancel = cancel or Cancel()

    def run(
        self,
        transcript: Transcript,
        *,
        system: str,
        tools: list[dict[str, Any]],
        workspace: Workspace | None,
        handlers: Mapping[str, Callable[[dict[str, Any]], tuple[str, bool]]]
        | None = None,
        checks: Mapping[str, Callable[[dict[str, Any]], str | None]] | None = None,
        model: str | None = None,
        on_turn: Callable[[str], None] | None = None,
        screen: Callable[[ToolCall], Screening] | None = None,
    ) -> RunOutcome:
        """``handlers`` answer control tools that don't end the run (e.g.
        start_milestone): each returns a result and whether it succeeded.
        ``checks`` vet a waiting tool's arguments: a problem is the call's
        failed result and the run goes on, so the model puts it right.
        ``screen`` judges each execution call before it reaches the build VM
        (P4)."""
        handlers = handlers or {}
        checks = checks or {}
        model = model or self._model  # the work item's pinned driver (L9)
        failures: list[str] = []  # the failed commands in a row, latest last
        offered = names(tools)
        nudged = False
        for turn in range(self._max_turns):
            # L2: stop takes effect at the next tool boundary.
            if self._cancel.requested:
                return RunOutcome("stopped", note=f"before turn {turn + 1}")
            try:
                produced = self._client.complete(
                    transcript.view(), tools=tools, model=model, system=system
                )
            except LLMError as exc:
                return RunOutcome("failed", note=f"the driver's call failed: {exc}")
            calls: list[ToolCall] = []
            for event in produced:
                transcript.append(event)
                if isinstance(event, ToolCall):
                    calls.append(event)
            # The model's own words for the turn: its progress, in prose (#42).
            said = " ".join(
                e.text.strip() for e in produced if isinstance(e, AssistantText)
            ).strip()
            if said and calls and on_turn is not None:
                on_turn(said)
            if not calls:
                if nudged:
                    return RunOutcome("failed", note="kept replying in prose")
                nudged = True
                # In the conversation, as the harness: a request must end on
                # the user's side (models refuse assistant prefill), and the
                # nudge is then in the log as what it was.
                transcript.as_actor("orchestrator").append(Fact(text=_PROSE_NUDGE))
                continue
            for index, call in enumerate(calls):
                if self._cancel.requested:
                    skip_calls(transcript, calls[index:])
                    return RunOutcome("stopped", note=f"during turn {turn + 1}")
                problem = (
                    checks[call.tool](dict(call.arguments))
                    if call.tool in checks and call.tool in offered
                    else None
                )
                if problem is not None:
                    transcript.append(
                        ToolResult(call_id=call.call_id, result=problem, success=False)
                    )
                    skip_calls(
                        transcript, calls[index + 1 :], "not run: fix the call above"
                    )
                    break
                if call.tool in WAITING and call.tool in offered:
                    transcript.append(
                        ToolResult(
                            call_id=call.call_id,
                            result="posted; your turn has ended",
                            success=True,
                        )
                    )
                    skip_calls(
                        transcript, calls[index + 1 :], "not run: the turn had ended"
                    )
                    return RunOutcome("wait", call.tool, dict(call.arguments))
                if call.tool in handlers and call.tool in offered:
                    result, ok = handlers[call.tool](dict(call.arguments))
                elif call.tool not in offered or workspace is None:
                    result, ok = (
                        f"no such tool: {call.tool}; the tools are "
                        + ", ".join(sorted(offered)),
                        False,
                    )
                else:
                    screening = screen(call) if screen else Screening("run")
                    if screening.decision == "ask":
                        transcript.append(
                            ToolResult(
                                call_id=call.call_id,
                                result="not run: it needs the operator's agreement, "
                                "and they've been asked; your turn has ended",
                                success=False,
                            )
                        )
                        skip_calls(
                            transcript,
                            calls[index + 1 :],
                            "not run: the turn had ended",
                        )
                        return RunOutcome(
                            "wait",
                            ASK_ALLOW,
                            {"question": screening.text, "key": screening.key},
                        )
                    if screening.decision == "refuse":
                        transcript.append(
                            ToolResult(
                                call_id=call.call_id,
                                result=f"not run: {screening.text}",
                                success=False,
                            )
                        )
                        skip_calls(
                            transcript,
                            calls[index + 1 :],
                            "not run: fix the call above",
                        )
                        break
                    outcome = workspace.dispatch(call.tool, call.arguments)
                    result, ok = outcome.result, outcome.success
                    if screening.note:
                        result = f"{result}\n\n[harness] {screening.note}"
                transcript.append(
                    ToolResult(call_id=call.call_id, result=result, success=ok)
                )
                if call.tool == "run_command":
                    failures = [] if ok else [*failures, _brief(call, result)]
                    if self._failure_limit and len(failures) >= self._failure_limit:
                        skip_calls(
                            transcript, calls[index + 1 :], "not run: you're stuck"
                        )
                        return RunOutcome(
                            "stuck",
                            note=f"{len(failures)} commands failed in a row: "
                            + "; ".join(failures),
                        )
        return RunOutcome("capped", note=f"hit the {self._max_turns}-turn cap")


def _brief(call: ToolCall, result: str) -> str:
    """A failed command and the end of what it said, for a helper's brief."""
    command = str(call.arguments.get("command", ""))[:120]
    return f"`{command}`: {result.strip()[-300:]}"