nisse/nisse-core/tests/test_calls.py

"""The call record (architecture L7): every model call, as sent and answered,
with where its request came from, for the operator's study."""

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

import httpx
import pytest

from nisse_core.log.canonical import (
    AssistantText,
    CanonicalEvent,
    Compaction,
    StoredEvent,
    ToolCall,
    ToolResult,
    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, ContextOverflow
from nisse_core.models.llm import LiteLLMClient, LLMError
from nisse_core.oplog import OpLog
from nisse_core.state.store import StateStore

TOKEN = "sk-secret-litellm-token"
FINISH = tool("finish", "finish the work", summary="what was done")


def stored(events: list[CanonicalEvent], start: int = 1) -> list[StoredEvent]:
    return [
        StoredEvent(id=start + i, ts="t", event=e, work_item="w", actor="driver")
        for i, e in enumerate(events)
    ]


def reply(
    content: str = "",
    *,
    tool_call: str | 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 tool_call:
        message["tool_calls"] = [
            {
                "id": "x",
                "type": "function",
                "function": {"name": tool_call, "arguments": '{"summary": "done"}'},
            }
        ]
    data: dict[str, Any] = {
        "model": "local",
        "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: answers each request in turn and keeps what it got."""

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

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


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


@pytest.fixture
def recorder(tmp_path: Path) -> calls.CallRecorder:
    return calls.CallRecorder(tmp_path / "state.calls.db")


def read(recorder: calls.CallRecorder) -> calls.CallReader:
    return calls.CallReader(recorder.path)


# --- what each call was --------------------------------------------------------


def test_a_call_is_recorded_exactly_as_sent_with_its_reasoning(
    recorder: calls.CallRecorder,
) -> None:
    server = Server(
        (
            200,
            reply(
                "Finishing.",
                tool_call="finish",
                reasoning="The tests pass, so the work is done.",
                timings={"cache_n": 90, "prompt_n": 30},
            ),
        )
    )
    client = budgeted(server, recorder)
    client.complete(
        stored([UserMessage(text="go", source="t")]), [FINISH], "local", system="sys"
    )

    reader = read(recorder)
    [call] = reader.calls()
    # The request rebuilt from the record is the body the server got.
    assert reader.request(call) == server.bodies[0]
    response = reader.response(call)
    assert (
        response["choices"][0]["message"]["reasoning_content"]
        == "The tests pass, so the work is done."
    )
    assert call["outcome"] == "answered"
    assert (call["tokens_in"], call["tokens_out"]) == (120, 7)
    assert (call["cache_n"], call["prompt_n"]) == (90, 30)
    assert call["route"] == "chat" and call["status"] == 200


def test_the_token_is_nowhere_in_the_record(recorder: calls.CallRecorder) -> None:
    client = budgeted(Server(), recorder)
    client.complete(stored([UserMessage(text="go", source="t")]), None, "local")
    recorder.close()

    raw = recorder.path.read_bytes()
    conn = sqlite3.connect(recorder.path)
    blobs = b"".join(
        zlib.decompress(d) for (d,) in conn.execute("SELECT data FROM blobs")
    )
    rows = json.dumps(conn.execute("SELECT * FROM calls").fetchall()).encode()
    conn.close()
    assert TOKEN.encode() not in raw + blobs + rows


def test_a_failed_call_is_recorded_with_its_error(recorder: calls.CallRecorder) -> None:
    server = Server((400, {"error": {"message": "bad request: no such model"}}))
    client = budgeted(server, recorder)
    with pytest.raises(LLMError):
        client.complete(stored([UserMessage(text="go", source="t")]), None, "local")

    reader = read(recorder)
    [call] = reader.calls()
    assert call["outcome"] == "failed"
    assert "LLMError" in call["error"]
    assert call["status"] == 400
    assert reader.request(call) == server.bodies[0]
    assert reader.response(call) == {"error": {"message": "bad request: no such model"}}


def test_a_request_too_big_to_send_is_recorded_without_one(
    recorder: calls.CallRecorder,
) -> None:
    server = Server()
    client = budgeted(server, recorder, window=4096)
    with pytest.raises(ContextOverflow):
        client.complete(
            stored([UserMessage(text="go", source="t")]),
            None,
            "local",
            system="s" * 20_000,
        )

    [call] = read(recorder).calls()
    assert server.bodies == []
    assert call["outcome"] == "failed" and "ContextOverflow" in call["error"]
    assert call["request"] is None


def test_a_growing_conversation_is_stored_once(recorder: calls.CallRecorder) -> None:
    client = budgeted(Server(), recorder)
    events: list[CanonicalEvent] = [
        UserMessage(text="message 0 " + "x" * 2000, source="t")
    ]
    for n in range(1, 6):
        client.complete(stored(events), None, "local")
        events.append(AssistantText(text=f"reply {n} " + "y" * 2000, model="local"))
        events.append(UserMessage(text=f"message {n} " + "z" * 2000, source="t"))

    conn = sqlite3.connect(recorder.path)
    stored_chars = sum(
        len(zlib.decompress(d)) for (d,) in conn.execute("SELECT data FROM blobs")
    )
    conn.close()
    # Five calls resend 1 + 3 + 5 + 7 + 9 messages of about 2 kB; stored
    # once each, that's 9 messages (as sent, and as events), not 25.
    assert stored_chars < 2 * 9 * 2100 + 20_000


# --- where each request came from ---------------------------------------------


def test_compaction_shortening_and_dropping_are_recorded(
    recorder: calls.CallRecorder,
) -> None:
    long = "line\n" * 2000
    events = stored(
        [
            UserMessage(text="start", source="t"),  # 1
            AssistantText(text="old thinking", model="m"),  # 2, compacted
            AssistantText(text="more old thinking", model="m"),  # 3, compacted
            Compaction(summary="notes so far", model="m", replaces=[2, 3]),  # 4
        ]
    )
    for n in range(8):  # turns 5.. : a command and its long output each
        events += stored(
            [
                ToolCall(
                    call_id=f"c{n}", tool="run_command", arguments={"command": "make"}
                ),
                ToolResult(call_id=f"c{n}", result=long, success=True),
            ],
            start=len(events) + 1,
        )
    client = BudgetedClient(
        LiteLLMClient(
            "https://litellm.test", TOKEN, transport=httpx.MockTransport(Server())
        ),
        lambda m: (12_000, 1024),
        keep_recent_turns=2,
        call_recorder=recorder,
    )
    client.complete(events, None, "local")

    [call] = read(recorder).calls()
    assert call["compacted"] == {"2": 4, "3": 4}
    assert call["shortened"], "older results were cut to fit"
    assert all(cut > 0 for cut in call["shortened"].values())
    sent_ids = [entry[0] for entry in call["sent"]]
    assert 4 in sent_ids and 2 not in sent_ids
    assert set(map(int, call["shortened"])) <= set(sent_ids)
    assert not set(call["dropped"]) & set(sent_ids)


def test_runs_nest_and_label_their_calls(recorder: calls.CallRecorder) -> None:
    client = budgeted(Server(), recorder)
    events = stored([UserMessage(text="go", source="t")])
    with calls.run(recorder, work_item="w", role="driver") as outer:
        client.complete(events, None, "local")
        with calls.purpose("consult", role="expert"):
            client.complete(events, None, "local")
        with calls.run(
            recorder, work_item="w", role="helper", kind="delegation"
        ) as inner:
            client.complete(events, None, "local")
            inner.end("handback", "done")
        outer.end("wait", "")

    reader = read(recorder)
    runs = {r["kind"]: r for r in reader.runs()}
    assert runs["delegation"]["parent"] == runs["run"]["id"]
    assert runs["run"]["outcome"] == "wait" and runs["run"]["ended_at"]
    assert runs["delegation"]["outcome"] == "handback"
    turn, consult, helped = reader.calls()
    assert (turn["run"], turn["role"], turn["purpose"]) == (
        runs["run"]["id"],
        "driver",
        None,
    )
    assert (consult["role"], consult["purpose"]) == ("expert", "consult")
    assert (helped["run"], helped["role"]) == (runs["delegation"]["id"], "helper")
    assert {c["work_item"] for c in (turn, consult, helped)} == {"w"}


def test_a_driver_run_is_recorded_with_how_it_ended(
    tmp_path: Path, recorder: calls.CallRecorder
) -> None:
    store = StateStore(tmp_path / "state.db", oplog=OpLog("t", io.StringIO()))
    store.open_work_item("w", title="t", position="RUNNING", cause="test")
    store.transcript("w", "operator").append(UserMessage(text="go", source="t"))
    client = budgeted(Server((200, reply(tool_call="finish"))), recorder)

    outcome = Driver(client, "local").run(
        store.transcript("w", "driver"),
        system="sys",
        tools=[FINISH],
        waiting=("finish",),
        workspace=None,
    )

    reader = read(recorder)
    [run] = reader.runs()
    assert outcome.kind == "wait"
    assert (run["work_item"], run["role"], run["outcome"]) == ("w", "driver", "wait")
    [call] = reader.calls(run=run["id"])
    # Each event sent is the log's, by id, so the page can read it against
    # the state store.
    assert [entry[:3] for entry in call["sent"]] == [[1, "operator", "user_message"]]


# --- the agent is told, and nothing depends on the record ----------------------


def test_the_agent_is_told_exactly_while_it_is_recorded(tmp_path: Path) -> None:
    on, off = Server(), Server()
    budgeted(on, calls.CallRecorder(tmp_path / "a.calls.db")).complete(
        stored([UserMessage(text="go", source="t")]), None, "local", system="sys"
    )
    budgeted(off, None).complete(
        stored([UserMessage(text="go", source="t")]), None, "local", system="sys"
    )
    assert on.bodies[0]["messages"][0]["content"] == f"sys\n\n{calls.NOTICE}"
    assert off.bodies[0]["messages"][0]["content"] == "sys"


def test_a_recorder_that_fails_never_fails_the_call(tmp_path: Path) -> None:
    out = io.StringIO()
    recorder = calls.CallRecorder(tmp_path / "state.calls.db", oplog=OpLog("t", out))
    recorder._conn.close()  # the disk went away, say
    client = budgeted(Server(), recorder)

    first = client.complete(stored([UserMessage(text="go", source="t")]), None, "local")
    second = client.complete(
        stored([UserMessage(text="go", source="t")]), None, "local"
    )

    assert [e.text for e in first + second] == ["ok", "ok"]  # type: ignore[union-attr]
    assert not recorder.active
    logged = [json.loads(line) for line in out.getvalue().splitlines()]
    assert [e["event"] for e in logged].count("calls.record_failed") == 1


def test_the_record_is_new_each_start(tmp_path: Path) -> None:
    path = tmp_path / "state.calls.db"
    path.write_bytes(b"an older layout, or not a database at all")
    Path(f"{path}-wal").write_bytes(b"stale")
    recorder = calls.CallRecorder(path)
    budgeted(Server(), recorder).complete(
        stored([UserMessage(text="go", source="t")]), None, "local"
    )
    assert recorder.active
    assert len(read(recorder).calls()) == 1

    again = calls.CallRecorder(path)
    assert again.active and read(again).calls() == []


def test_off_means_no_file(tmp_path: Path) -> None:
    assert calls.recorder_for(tmp_path / "state.db", False) is None
    assert not list(tmp_path.iterdir())
    recorder = calls.recorder_for(tmp_path / "state.db", True)
    assert recorder is not None and recorder.path == tmp_path / "state.calls.db"


def test_a_reader_refuses_a_layout_it_does_not_know(
    recorder: calls.CallRecorder,
) -> None:
    recorder._conn.execute("UPDATE meta SET value = '99' WHERE key = 'layout_version'")
    with pytest.raises(calls.UnknownLayout, match="layout 99"):
        read(recorder)


def test_the_reader_never_writes(recorder: calls.CallRecorder) -> None:
    budgeted(Server(), recorder).complete(
        stored([UserMessage(text="go", source="t")]), None, "local"
    )
    reader = read(recorder)
    with pytest.raises(sqlite3.OperationalError, match="readonly"):
        reader._conn.execute("DELETE FROM calls")


def test_the_page_reads_one_items_calls_and_when_recording_began(
    recorder: calls.CallRecorder,
) -> None:
    client = budgeted(Server(), recorder)
    with calls.run(recorder, work_item="w", role="driver"):
        client.complete(stored([UserMessage(text="go", source="t")]), None, "local")
    with calls.run(recorder, work_item=None, role="classifier"):
        client.complete(stored([UserMessage(text="hi", source="t")]), None, "local")
    reader = read(recorder)
    [mine] = reader.calls_for("w")
    [agent_level] = reader.calls_for(None)
    assert (mine["work_item"], agent_level["work_item"]) == ("w", None)
    assert reader.call(agent_level["id"])["role"] == "classifier"
    assert reader.call(999) is None
    assert reader.started_at is not None
    assert reader.started_at <= mine["at"]