nisse/nisse-ui/tests/ui_fixtures.py

"""Stores written by the agents' own code, for the page to read (X6): a real
state store and a real call record from a driver run against a fake
LiteLLM, so a change in nisse-core that breaks the page fails here."""

import io
import json
from pathlib import Path
from typing import Any

import httpx

from nisse_core.log.canonical import Fact, UserMessage
from nisse_core.loop.driver import Driver
from nisse_core.loop.tools import tool
from nisse_core.models import calls
from nisse_core.models.context import BudgetedClient
from nisse_core.models.llm import LiteLLMClient
from nisse_core.oplog import OpLog
from nisse_core.state.store import StateStore
from nisse_ui.sources import Source

TOKEN = "sk-secret-litellm-token"
HOSTILE = "<script>alert('x')</script></pre></details>"
READ = tool("read_file", "read a file", path="the file")
POST = tool("post_comment", "post a comment", body="the comment")
FINISH = tool("finish", "finish the work", summary="what was done")


def reply(
    content: str = "",
    *,
    call: tuple[str, dict[str, Any]] | None = None,
    reasoning: str = "",
    timings: dict[str, int] | None = None,
) -> dict[str, Any]:
    message: dict[str, Any] = {"role": "assistant", "content": content}
    if reasoning:
        message["reasoning_content"] = reasoning
    if call:
        message["tool_calls"] = [
            {
                "id": f"x{abs(hash(json.dumps(call)))}",
                "type": "function",
                "function": {"name": call[0], "arguments": json.dumps(call[1])},
            }
        ]
    data: dict[str, Any] = {
        "model": "local-backing",
        "choices": [{"message": message, "finish_reason": "stop"}],
        "usage": {"prompt_tokens": 120, "completion_tokens": 7},
    }
    if timings is not None:
        data["timings"] = timings
    return data


class Server:
    """A fake LiteLLM answering each request in turn."""

    def __init__(self, *answers: dict[str, Any]) -> None:
        self.answers = list(answers)

    def __call__(self, request: httpx.Request) -> httpx.Response:
        data = self.answers.pop(0) if self.answers else reply("ok")
        return httpx.Response(200, json=data)


def client(server: Server, recorder: calls.CallRecorder) -> BudgetedClient:
    inner = LiteLLMClient(
        "https://litellm.test", TOKEN, transport=httpx.MockTransport(server)
    )
    return BudgetedClient(inner, lambda m: (32768, 1024), call_recorder=recorder)


def store_at(root: Path, store_class: type[StateStore] = StateStore) -> StateStore:
    return store_class(root / "state.db", oplog=OpLog("test", io.StringIO()))


def recorded(root: Path) -> tuple[Source, StateStore]:
    """One work item, one driver run of three turns: a tool call whose
    result is hostile text, a prose reply the driver corrects, then the
    finish. And one agent-level call. The store stays open, as a live
    agent's does."""
    root.mkdir(parents=True, exist_ok=True)
    recorder = calls.CallRecorder(calls.path_for(root / "state.db"))
    store = store_at(root)
    store.open_work_item("w1", title="Fix the parser", position="RUNNING", cause="t")
    store.transcript("w1", "orchestrator").append(
        Fact(text="The operator agreed to this restatement: fix the parser.")
    )
    store.transcript("w1", "operator").append(UserMessage(text="go", source="m:1"))
    server = Server(
        reply(
            call=("read_file", {"path": "parser.py"}),
            reasoning="First, read the parser.",
            timings={"cache_n": 0, "prompt_n": 300},
        ),
        reply(
            "I believe the parser is fine now.",
            timings={"cache_n": 280, "prompt_n": 60},
        ),
        reply(
            call=("finish", {"summary": "fixed"}),
            timings={"cache_n": 330, "prompt_n": 20},
        ),
    )
    Driver(client(server, recorder), "local").run(
        store.transcript("w1", "driver"),
        system="You are the driver.",
        tools=[READ, FINISH],
        waiting=("finish",),
        workspace=None,
        handlers={"read_file": lambda args: (f"contents {HOSTILE}", True)},
    )
    store.transcript(None, "operator").append(UserMessage(text="hello", source="m:9"))
    with calls.run(recorder, work_item=None, role="classifier", kind="classification"):
        client(Server(reply("harmless")), recorder).complete(
            store.transcript(None, "operator").read(), None, "local", system="Classify."
        )
    recorder.close()
    return Source("agent", root / "state.db", "agent"), store


def circling(root: Path) -> Source:
    """A run that posts the same comment over and over, as Flash-Next did
    on 2026-10-24, until the finish."""
    root.mkdir(parents=True, exist_ok=True)
    recorder = calls.CallRecorder(calls.path_for(root / "state.db"))
    store = store_at(root)
    store.open_work_item("pr-7", title="Review #7", position="RUNNING", cause="t")
    store.transcript("pr-7", "operator").append(UserMessage(text="review", source="m"))
    answers = [
        reply(call=("post_comment", {"body": f"Duplicate, please disregard ({n})."}))
        for n in range(4)
    ]
    Driver(
        client(Server(*answers, reply(call=("finish", {"summary": "x"}))), recorder),
        "local",
        circle_repeats=0,  # as it ran then: before M7.8 ended such runs
    ).run(
        store.transcript("pr-7", "driver"),
        system="You review.",
        tools=[POST, FINISH],
        waiting=("finish",),
        workspace=None,
        handlers={"post_comment": lambda args: ("posted", True)},
    )
    recorder.close()
    store.close()
    return Source("circling", root / "state.db", "agent")