STT-Assistant/tests/test_llm.py

"""Tests for the LLM client and chat session (no network)."""

import pytest

from stt_daemon.assistant.llm import (
    BriefEngine,
    ChatSession,
    ConversationEngine,
    LlmClient,
    LlmError,
    LlmResponse,
)
from stt_daemon.config import LlmConfig


def _make_client(responder, cfg=None):
    cfg = cfg or LlmConfig()
    calls: list[dict] = []

    def opener(payload, timeout):
        calls.append(payload)
        return responder(payload)

    return LlmClient(cfg, opener=opener), calls


def _reply(content="", reasoning=None, tool_calls=None, finish="stop"):
    message = {"role": "assistant", "content": content}
    if reasoning is not None:
        message["reasoning_content"] = reasoning
    if tool_calls:
        message["tool_calls"] = tool_calls
    return {"choices": [{"message": message, "finish_reason": finish}]}


def test_reasoning_content_is_ignored():
    client, _ = _make_client(
        lambda _p: _reply(content="It is noon.", reasoning="thinking hard about it")
    )
    response = client.chat([{"role": "user", "content": "hi"}])
    assert response.content == "It is noon."
    assert "thinking hard" not in response.content


def test_cache_prompt_and_endpoint_are_sent():
    client, calls = _make_client(lambda _p: _reply(content="ok"))
    client.chat([{"role": "user", "content": "hi"}])

    assert calls[0]["body"]["cache_prompt"] is True
    assert calls[0]["body"]["stream"] is False
    assert calls[0]["__url__"] == LlmConfig().endpoint


def test_model_is_sent_when_set():
    cfg = LlmConfig(model="gpt-4o-mini")
    client, calls = _make_client(lambda _p: _reply(content="ok"), cfg)

    client.chat([{"role": "user", "content": "hi"}])

    assert calls[0]["body"]["model"] == "gpt-4o-mini"


def test_model_omitted_when_empty():
    client, calls = _make_client(lambda _p: _reply(content="ok"))

    client.chat([{"role": "user", "content": "hi"}])

    assert "model" not in calls[0]["body"]


def test_api_key_env_adds_bearer_header(monkeypatch):
    monkeypatch.setenv("STT_TEST_LLM_KEY", "sekret")
    cfg = LlmConfig(api_key_env="STT_TEST_LLM_KEY")
    client, calls = _make_client(lambda _p: _reply(content="ok"), cfg)

    client.chat([{"role": "user", "content": "hi"}])

    assert calls[0]["headers"]["Authorization"] == "Bearer sekret"
    assert calls[0]["headers"]["Content-Type"] == "application/json"


def test_missing_api_key_env_sends_no_authorization(monkeypatch):
    monkeypatch.delenv("STT_TEST_LLM_KEY", raising=False)
    cfg = LlmConfig(api_key_env="STT_TEST_LLM_KEY")
    client, calls = _make_client(lambda _p: _reply(content="ok"), cfg)

    client.chat([{"role": "user", "content": "hi"}])

    assert "Authorization" not in calls[0]["headers"]


def test_no_api_key_sends_only_content_type(monkeypatch):
    client, calls = _make_client(lambda _p: _reply(content="ok"))

    client.chat([{"role": "user", "content": "hi"}])

    assert calls[0]["headers"] == {"Content-Type": "application/json"}


def test_session_header_and_user_agent_sent_when_set():
    cfg = LlmConfig(
        session_header="x-opencode-session",
        session_id="stt-chat",
        user_agent="stt-daemon/1.0",
    )
    client, calls = _make_client(lambda _p: _reply(content="ok"), cfg)

    client.chat([{"role": "user", "content": "hi"}])

    assert calls[0]["headers"]["x-opencode-session"] == "stt-chat"
    assert calls[0]["headers"]["User-Agent"] == "stt-daemon/1.0"


def test_session_header_omitted_when_id_empty(caplog):
    cfg = LlmConfig(session_header="x-opencode-session", session_id="")
    client, calls = _make_client(lambda _p: _reply(content="ok"), cfg)

    client.chat([{"role": "user", "content": "hi"}])

    assert "x-opencode-session" not in calls[0]["headers"]
    assert "session_id is empty" in caplog.text


def test_session_header_omitted_when_header_empty():
    cfg = LlmConfig(session_header="", session_id="stt-chat")
    client, calls = _make_client(lambda _p: _reply(content="ok"), cfg)

    client.chat([{"role": "user", "content": "hi"}])

    assert "stt-chat" not in calls[0]["headers"].values()
    assert calls[0]["headers"] == {"Content-Type": "application/json"}


def test_reasoning_effort_is_sent_when_set():
    cfg = LlmConfig(reasoning_effort="high")
    client, calls = _make_client(lambda _p: _reply(content="ok"), cfg)

    client.chat([{"role": "user", "content": "hi"}])

    assert calls[0]["body"]["reasoning_effort"] == "high"


def test_reasoning_effort_omitted_when_empty():
    client, calls = _make_client(lambda _p: _reply(content="ok"))

    client.chat([{"role": "user", "content": "hi"}])

    assert "reasoning_effort" not in calls[0]["body"]


def test_extra_body_merged_and_overrides_defaults():
    cfg = LlmConfig(extra_body={"enable_thinking": True, "temperature": 0.1})
    client, calls = _make_client(lambda _p: _reply(content="ok"), cfg)

    client.chat([{"role": "user", "content": "hi"}])

    assert calls[0]["body"]["enable_thinking"] is True
    assert calls[0]["body"]["temperature"] == 0.1


def test_tools_sent_only_when_provided():
    client, calls = _make_client(lambda _p: _reply(content="ok"))
    client.chat([{"role": "user", "content": "hi"}])
    assert "tools" not in calls[0]["body"]

    client.chat(
        [{"role": "user", "content": "hi"}],
        tools=[{"type": "function", "function": {"name": "get_datetime"}}],
    )
    assert calls[1]["body"]["tools"][0]["function"]["name"] == "get_datetime"
    assert calls[1]["body"]["tool_choice"] == "auto"


def test_id_slot_sent_only_when_slot_id_provided():
    client, calls = _make_client(lambda _p: _reply(content="ok"))
    client.chat([{"role": "user", "content": "hi"}])
    assert "id_slot" not in calls[0]["body"]

    client.chat([{"role": "user", "content": "hi"}], slot_id=1)
    assert calls[1]["body"]["id_slot"] == 1


def test_slot_and_on_io_never_passed_to_old_fake_when_none():
    seen: list[dict] = []

    class _OldFake:
        def chat(self, messages, tools=None, max_tokens=None, timeout_s=None):
            seen.append(
                {
                    "messages": messages,
                    "tools": tools,
                    "max_tokens": max_tokens,
                    "timeout_s": timeout_s,
                }
            )
            return {"ok": True}

    fake = _OldFake()
    slot_id = None
    on_io = None
    kwargs: dict = {}
    if slot_id is not None:
        kwargs["slot_id"] = slot_id
    if on_io is not None:
        kwargs["on_io"] = on_io
    fake.chat([{"role": "user", "content": "hi"}], **kwargs)
    assert seen[0] == {
        "messages": [{"role": "user", "content": "hi"}],
        "tools": None,
        "max_tokens": None,
        "timeout_s": None,
    }


def test_tool_calls_are_parsed():
    client, _ = _make_client(
        lambda _p: _reply(
            tool_calls=[
                {
                    "id": "1",
                    "type": "function",
                    "function": {"name": "get_datetime", "arguments": "{}"},
                }
            ],
            finish="tool_calls",
        )
    )
    response = client.chat([{"role": "user", "content": "time?"}], tools=[{"type": "function"}])
    assert len(response.tool_calls) == 1
    assert response.tool_calls[0].name == "get_datetime"
    assert response.finish_reason == "tool_calls"


def test_missing_choices_raises():
    client, _ = _make_client(lambda _p: {"choices": []})
    with pytest.raises(LlmError):
        client.chat([{"role": "user", "content": "hi"}])


def test_chat_session_appends_and_keeps_prefix():
    client, calls = _make_client(lambda _p: _reply(content="hello"))
    session = ChatSession(LlmConfig(), client)

    assert session.send("hi") == "hello"
    session.send("again")

    assert [m["role"] for m in session.history] == [
        "user",
        "assistant",
        "user",
        "assistant",
    ]
    first = calls[0]["body"]["messages"]
    second = calls[1]["body"]["messages"]
    assert second[: len(first)] == first  # append-only, prefix preserved
    assert second[0]["content"] == LlmConfig().system_prompt_chat


def test_chat_session_trims_in_blocks_and_keeps_system_prompt():
    cfg = LlmConfig(
        max_context_tokens=1330,
        max_tokens=1024,
        compaction_safety_margin=256,
        trim_block_turns=1,
        compaction_enabled=False,
    )  # prompt_budget == 50
    client, calls = _make_client(lambda _p: _reply(content="ok"))
    session = ChatSession(cfg, client)

    for _ in range(10):
        session.send("x" * 40)

    assert len(session.history) <= 6
    assert calls[-1]["body"]["messages"][0]["content"] == cfg.system_prompt_chat


def test_chat_session_reset_forgets_history():
    client, _ = _make_client(lambda _p: _reply(content="ok"))
    session = ChatSession(LlmConfig(), client)
    session.send("hi")
    session.reset()
    assert session.history == []


def test_chat_session_drops_reply_when_reset_mid_flight():
    holder: list[ChatSession] = []

    def opener(_payload, _timeout):
        holder[0].reset()
        return _reply(content="late reply")

    holder.append(ChatSession(LlmConfig(), LlmClient(LlmConfig(), opener=opener)))

    assert holder[0].send("hi") == ""
    assert holder[0].history == []


# ---------------------------------------------------------------------------
# Phase C tests
# ---------------------------------------------------------------------------

def _compact_cfg(
    *,
    budget_chars=200,
    trigger_ratio=0.8,
    keep_turns=1,
    compaction_max_tokens=64,
    compaction_timeout_s=5.0,
    compaction_input_tokens=8192,
    compaction_enabled=True,
    trim_block_turns=1,
):
    """Return a small LlmConfig whose prompt_budget == budget_chars // 4 tokens.

    prompt_budget = max_context_tokens - max_tokens - compaction_safety_margin
    We want budget_chars // 4 tokens, so:
        max_context_tokens = (budget_chars // 4) + max_tokens + compaction_safety_margin
    """
    max_tokens              = 16
    compaction_safety_margin = 0
    max_context_tokens      = (budget_chars // 4) + max_tokens + compaction_safety_margin
    return LlmConfig(
        max_context_tokens=max_context_tokens,
        max_tokens=max_tokens,
        compaction_safety_margin=compaction_safety_margin,
        compaction_enabled=compaction_enabled,
        compaction_trigger_ratio=trigger_ratio,
        compaction_keep_turns=keep_turns,
        compaction_max_tokens=compaction_max_tokens,
        compaction_timeout_s=compaction_timeout_s,
        compaction_input_tokens=compaction_input_tokens,
        trim_block_turns=trim_block_turns,
        system_prompt_chat="SYS",
        system_prompt_command="CMD",
        compaction_system_prompt="COMPACT",
    )


def test_compaction_fires_and_window_shrinks():
    """Crossing the trigger compacts: window shrinks, summary present in next prompt."""
    cfg = _compact_cfg(budget_chars=200, trigger_ratio=0.8, keep_turns=1)
    # prompt_budget = 200//4 = 50 tokens; trigger at 40 tokens (0.8 * 50)

    summary_text = "summary of old turns"
    call_index   = 0

    def responder(payload):
        nonlocal call_index
        msgs = payload["body"]["messages"]
        call_index += 1
        # Summarizer calls have the compaction system prompt.
        if msgs[0]["content"] == "COMPACT":
            return _reply(content=summary_text)
        return _reply(content="ok")

    client, calls = _make_client(responder, cfg)
    session = ChatSession(cfg, client)

    # Send enough turns to exceed the trigger (each "x"*40 = 40 chars = 10 tokens).
    # After 4 turns (user+assistant pairs = 8 messages * 10 tokens = 80 tokens > 40 trigger).
    for _ in range(4):
        session.send("x" * 40)

    # After compaction the window must be smaller than the full history.
    assert len(session.history) <= 4  # keep_turns=1 → keep=2 messages

    # The summary must be set.
    assert session.summary != ""

    # The next chat call must include the summary as a second system message,
    # with the chat system prompt FIRST.
    session.send("after compaction")
    chat_calls = [c for c in calls if c["body"]["messages"][0]["content"] == "SYS"]
    last_chat  = chat_calls[-1]["body"]["messages"]
    assert last_chat[0]["role"] == "system"
    assert last_chat[0]["content"] == "SYS"
    assert last_chat[1]["role"] == "system"
    assert "Earlier conversation summary:" in last_chat[1]["content"]
    assert summary_text in last_chat[1]["content"]


def test_summary_never_appended_to_history():
    """The summary text must not appear as a history entry."""
    cfg = _compact_cfg(budget_chars=200, trigger_ratio=0.8, keep_turns=1)
    summary_text = "SUMMARY_SENTINEL"

    def responder(payload):
        msgs = payload["body"]["messages"]
        if msgs[0]["content"] == "COMPACT":
            return _reply(content=summary_text)
        return _reply(content="ok")

    client, _ = _make_client(responder, cfg)
    session   = ChatSession(cfg, client)

    for _ in range(4):
        session.send("x" * 40)

    for msg in session.history:
        assert summary_text not in msg["content"]


def test_summarizer_failure_falls_back_to_trim():
    """On summarizer failure the session falls back to _trim() and stays within budget."""
    cfg = _compact_cfg(budget_chars=200, trigger_ratio=0.8, keep_turns=1)

    def responder(payload):
        msgs = payload["body"]["messages"]
        if msgs[0]["content"] == "COMPACT":
            return _reply(content="", finish="stop")  # empty → failure
        return _reply(content="ok")

    client, _ = _make_client(responder, cfg)
    session   = ChatSession(cfg, client)

    for _ in range(4):
        session.send("x" * 40)

    # Summary must remain empty (fallback used trim, not compaction).
    assert session.summary == ""
    # History must be bounded within the budget.
    assert session._estimate_tokens() <= cfg.prompt_budget


def test_summarizer_length_finish_reason_falls_back_to_trim():
    """finish_reason == 'length' on the summarizer is treated as failure."""
    cfg = _compact_cfg(budget_chars=200, trigger_ratio=0.8, keep_turns=1)

    def responder(payload):
        msgs = payload["body"]["messages"]
        if msgs[0]["content"] == "COMPACT":
            return _reply(content="partial summary", finish="length")
        return _reply(content="ok")

    client, _ = _make_client(responder, cfg)
    session   = ChatSession(cfg, client)

    for _ in range(4):
        session.send("x" * 40)

    assert session.summary == ""
    assert session._estimate_tokens() <= cfg.prompt_budget


def test_summarizer_llmerror_falls_back_to_trim():
    """LlmError from the summarizer falls back to _trim()."""
    cfg = _compact_cfg(budget_chars=200, trigger_ratio=0.8, keep_turns=1)

    def responder(payload):
        msgs = payload["body"]["messages"]
        if msgs[0]["content"] == "COMPACT":
            raise LlmError("network error")
        return _reply(content="ok")

    client, _ = _make_client(responder, cfg)
    session   = ChatSession(cfg, client)

    for _ in range(4):
        session.send("x" * 40)

    assert session.summary == ""
    assert session._estimate_tokens() <= cfg.prompt_budget


def test_reset_clears_summary():
    """reset() clears both history and summary."""
    cfg = _compact_cfg(budget_chars=200, trigger_ratio=0.8, keep_turns=1)
    summary_text = "old summary"

    def responder(payload):
        msgs = payload["body"]["messages"]
        if msgs[0]["content"] == "COMPACT":
            return _reply(content=summary_text)
        return _reply(content="ok")

    client, _ = _make_client(responder, cfg)
    session   = ChatSession(cfg, client)

    for _ in range(4):
        session.send("x" * 40)

    assert session.summary != ""

    session.reset()
    assert session.history == []
    assert session.summary == ""


def test_compaction_disabled_never_calls_summarizer():
    """compaction_enabled=False never calls the summarizer and still trims."""
    cfg = _compact_cfg(
        budget_chars=200,
        trigger_ratio=0.8,
        keep_turns=1,
        compaction_enabled=False,
    )

    summarizer_called = False

    def responder(payload):
        nonlocal summarizer_called
        msgs = payload["body"]["messages"]
        if msgs[0]["content"] == "COMPACT":
            summarizer_called = True
            return _reply(content="should not happen")
        return _reply(content="ok")

    client, _ = _make_client(responder, cfg)
    session   = ChatSession(cfg, client)

    for _ in range(6):
        session.send("x" * 40)

    assert not summarizer_called
    assert session.summary == ""
    # History must still be bounded (trim ran).
    assert session._estimate_tokens() <= cfg.prompt_budget


def test_nothing_to_evict_skips_summarizer():
    """When len(history) <= keep, compaction is skipped and trim runs instead."""
    # keep_turns=10 → keep=20 messages; we only send 2 turns (4 messages).
    cfg = _compact_cfg(budget_chars=200, trigger_ratio=0.8, keep_turns=10)

    summarizer_called = False

    def responder(payload):
        nonlocal summarizer_called
        msgs = payload["body"]["messages"]
        if msgs[0]["content"] == "COMPACT":
            summarizer_called = True
            return _reply(content="should not happen")
        return _reply(content="ok")

    client, _ = _make_client(responder, cfg)
    session   = ChatSession(cfg, client)

    # Push over the trigger but keep history shorter than keep=20.
    for _ in range(4):
        session.send("x" * 40)

    assert not summarizer_called


def test_chunk_and_fold_large_evicted():
    """Large evicted list triggers chunk-and-fold (more than one summarizer call).

    budget_chars=160 → prompt_budget=40 tokens, trigger at 32 tokens (0.8 * 40).
    Each turn: user='x'*40 (40 chars) + assistant='ok' (2 chars) = ~10 tokens.
    Compaction fires at turn 4 with 6 evicted messages.
    compaction_input_tokens=1 → char_limit=4 chars per chunk, so each 40-char
    or 2-char message is its own chunk → 6 summarizer calls for 6 chunks.
    """
    cfg = _compact_cfg(
        budget_chars=160,
        trigger_ratio=0.8,
        keep_turns=1,
        compaction_input_tokens=1,  # 1 token * 4 = 4 chars per chunk → each msg is a chunk
    )

    summarizer_calls = []

    def responder(payload):
        msgs = payload["body"]["messages"]
        if msgs[0]["content"] == "COMPACT":
            summarizer_calls.append(payload)
            return _reply(content="folded")
        return _reply(content="ok")

    client, _ = _make_client(responder, cfg)
    session   = ChatSession(cfg, client)

    # 4 turns → 8 messages → fires compaction with 6 evicted (keep=2).
    for _ in range(4):
        session.send("x" * 40)

    # With compaction_input_tokens=1 (4 chars/chunk) and 40-char messages,
    # each message is its own chunk → multiple summarizer calls.
    assert len(summarizer_calls) > 1


def test_chat_per_request_max_tokens_override():
    """LlmClient.chat sends per-request max_tokens when given, cfg default otherwise."""
    cfg = LlmConfig(max_tokens=1024, timeout_s=60.0)
    client, calls = _make_client(lambda _p: _reply(content="ok"), cfg)

    # Default: uses cfg.max_tokens.
    client.chat([{"role": "user", "content": "hi"}])
    assert calls[0]["body"]["max_tokens"] == 1024

    # Override: uses the supplied value.
    client.chat([{"role": "user", "content": "hi"}], max_tokens=128)
    assert calls[1]["body"]["max_tokens"] == 128


def test_chat_per_request_timeout_override():
    """LlmClient.chat uses per-request timeout_s when given, cfg default otherwise."""
    cfg = LlmConfig(max_tokens=1024, timeout_s=60.0)
    recorded_timeouts: list[float] = []

    def opener(payload, timeout):
        recorded_timeouts.append(timeout)
        return _reply(content="ok")

    client = LlmClient(cfg, opener=opener)

    # Default timeout.
    client.chat([{"role": "user", "content": "hi"}])
    assert recorded_timeouts[0] == 60.0

    # Per-request override.
    client.chat([{"role": "user", "content": "hi"}], timeout_s=5.0)
    assert recorded_timeouts[1] == 5.0


def test_summarizer_uses_compaction_max_tokens_and_timeout():
    """The summarizer call uses compaction_max_tokens and compaction_timeout_s."""
    cfg = _compact_cfg(
        budget_chars=200,
        trigger_ratio=0.8,
        keep_turns=1,
        compaction_max_tokens=42,
        compaction_timeout_s=7.0,
    )

    recorded: list[dict] = []

    def opener(payload, timeout):
        msgs = payload["body"]["messages"]
        if msgs[0]["content"] == "COMPACT":
            recorded.append({"max_tokens": payload["body"]["max_tokens"], "timeout": timeout})
            return _reply(content="summary")
        return _reply(content="ok")

    client  = LlmClient(cfg, opener=opener)
    session = ChatSession(cfg, client)

    for _ in range(4):
        session.send("x" * 40)

    assert len(recorded) >= 1
    assert recorded[0]["max_tokens"] == 42
    assert recorded[0]["timeout"] == 7.0


def test_notify_called_on_compaction():
    """notify('compacted', turns=N) is called after a successful compaction."""
    cfg = _compact_cfg(budget_chars=200, trigger_ratio=0.8, keep_turns=1)

    notifications: list[tuple] = []

    def notify(kind, **fields):
        notifications.append((kind, fields))

    def responder(payload):
        msgs = payload["body"]["messages"]
        if msgs[0]["content"] == "COMPACT":
            return _reply(content="summary")
        return _reply(content="ok")

    client  = LlmClient(cfg, opener=lambda p, t: responder(p))
    session = ChatSession(cfg, client, notify=notify)

    for _ in range(4):
        session.send("x" * 40)

    assert any(kind == "compacted" for kind, _ in notifications)
    compacted = [(k, f) for k, f in notifications if k == "compacted"]
    assert compacted[0][1]["turns"] > 0


def test_notify_not_called_when_none():
    """ChatSession works without a notify callback (default None)."""
    cfg = _compact_cfg(budget_chars=200, trigger_ratio=0.8, keep_turns=1)

    def responder(payload):
        msgs = payload["body"]["messages"]
        if msgs[0]["content"] == "COMPACT":
            return _reply(content="summary")
        return _reply(content="ok")

    client  = LlmClient(cfg, opener=lambda p, t: responder(p))
    session = ChatSession(cfg, client)  # no notify

    # Must not raise.
    for _ in range(4):
        session.send("x" * 40)


def test_system_prompt_is_first_message_after_compaction():
    """After compaction the chat system prompt is always the first message."""
    cfg = _compact_cfg(budget_chars=200, trigger_ratio=0.8, keep_turns=1)

    def responder(payload):
        msgs = payload["body"]["messages"]
        if msgs[0]["content"] == "COMPACT":
            return _reply(content="summary text")
        return _reply(content="ok")

    client, calls = _make_client(responder, cfg)
    session       = ChatSession(cfg, client)

    for _ in range(4):
        session.send("x" * 40)

    # Find the last chat call (system prompt == "SYS", not "COMPACT").
    chat_calls = [c for c in calls if c["body"]["messages"][0]["content"] == "SYS"]
    assert chat_calls, "no chat calls found after compaction"
    last_msgs = chat_calls[-1]["body"]["messages"]
    assert last_msgs[0]["role"] == "system"
    assert last_msgs[0]["content"] == "SYS"


def test_chat_session_truncated_empty_reply_notifies_and_is_not_appended():
    """A length-limited empty reply is reported and never stored as an empty turn."""
    notifications: list[tuple] = []

    def notify(kind, **fields):
        notifications.append((kind, fields))

    client, _ = _make_client(lambda _p: _reply(content="", finish="length"))
    session = ChatSession(LlmConfig(), client, notify=notify)

    assert session.send("hi") == ""
    assert session.history == [{"role": "user", "content": "hi"}]
    assert notifications == [("truncated", {"message": "response truncated"})]


def test_chat_session_truncated_partial_reply_is_appended_and_notified():
    """A truncated but non-empty reply is kept and reported as truncated."""
    notifications: list[tuple] = []

    def notify(kind, **fields):
        notifications.append((kind, fields))

    client, _ = _make_client(lambda _p: _reply(content="partial answer", finish="length"))
    session = ChatSession(LlmConfig(), client, notify=notify)

    assert session.send("hi") == "partial answer"
    assert session.history[-1] == {"role": "assistant", "content": "partial answer"}
    assert any(kind == "truncated" for kind, _ in notifications)


def test_on_io_receives_request_then_response():
    """chat(..., on_io=cb) reports the request (url+body) then the response payload."""
    client, _ = _make_client(lambda _p: _reply(content="ok"))
    seen: list[dict] = []

    response = client.chat([{"role": "user", "content": "hi"}], on_io=seen.append)

    assert response.content == "ok"
    assert [event["phase"] for event in seen] == ["request", "response"]
    assert seen[0]["url"] == LlmConfig().endpoint
    assert seen[0]["body"]["messages"] == [{"role": "user", "content": "hi"}]
    assert seen[0]["body"]["max_tokens"] == LlmConfig().max_tokens
    assert seen[1]["payload"]["choices"][0]["message"]["content"] == "ok"
    assert seen[1]["error"] is None


def test_on_io_response_phase_reports_duration_ms():
    """The response phase carries a non-negative float duration_ms."""
    client, _ = _make_client(lambda _p: _reply(content="ok"))
    seen: list[dict] = []

    client.chat([{"role": "user", "content": "hi"}], on_io=seen.append)

    assert "duration_ms" not in seen[0]  # request phase carries no timing
    assert isinstance(seen[1]["duration_ms"], float)
    assert seen[1]["duration_ms"] >= 0.0


def test_on_io_response_phase_reports_opener_error():
    """An opener raising still yields a response phase with the error set."""
    import urllib.error

    def opener(_payload, _timeout):
        raise urllib.error.URLError("down")

    client = LlmClient(LlmConfig(), opener=opener)
    seen: list[dict] = []

    with pytest.raises(LlmError):
        client.chat([{"role": "user", "content": "hi"}], on_io=seen.append)

    assert [event["phase"] for event in seen] == ["request", "response"]
    assert seen[1]["payload"] is None
    assert "down" in seen[1]["error"]
    assert isinstance(seen[1]["duration_ms"], float)
    assert seen[1]["duration_ms"] >= 0.0


def test_raising_on_io_does_not_break_chat():
    """A throwing on_io callback is swallowed and the reply still parses."""

    def bad(_event):
        raise RuntimeError("boom")

    client, _ = _make_client(lambda _p: _reply(content="ok"))

    assert client.chat([{"role": "user", "content": "hi"}], on_io=bad).content == "ok"


# ---------------------------------------------------------------------------
# ConversationEngine tests (Todo 2)
# ---------------------------------------------------------------------------

def _fake_engine_client(record=None, content="ok", finish="stop"):
    """Old-style fake: accepts no slot_id/on_io keywords, returns LlmResponse."""

    class _Fake:
        def chat(self, messages, max_tokens=None, timeout_s=None):
            if record is not None:
                record.append(
                    {
                        "messages": messages,
                        "max_tokens": max_tokens,
                        "timeout_s": timeout_s,
                    }
                )
            return LlmResponse(content=content, finish_reason=finish)

    return _Fake()


def test_engine_custom_system_prompt_is_first():
    client, calls = _make_client(lambda _p: _reply(content="ok"))
    engine = ConversationEngine(LlmConfig(), client, system_prompt="CUSTOM")

    assert engine.send("hi") == "ok"
    assert calls[0]["body"]["messages"][0] == {
        "role": "system",
        "content": "CUSTOM",
    }


def test_engine_custom_compaction_prompt_is_used():
    cfg = _compact_cfg(budget_chars=200, trigger_ratio=0.8, keep_turns=1)
    seen_prompts: list[str] = []

    def responder(payload):
        msgs = payload["body"]["messages"]
        if msgs[0]["content"] == "BRIEF-COMPACT":
            seen_prompts.append(msgs[0]["content"])
            return _reply(content="brief summary")
        assert msgs[0]["content"] == "SYS2"
        return _reply(content="ok")

    client, _ = _make_client(responder, cfg)
    engine = ConversationEngine(
        cfg,
        client,
        system_prompt="SYS2",
        compaction_system_prompt="BRIEF-COMPACT",
    )

    for _ in range(4):
        engine.send("x" * 40)

    assert seen_prompts, "summarizer never ran with the custom prompt"
    assert engine.summary == "brief summary"


def test_engine_slot_id_sent_only_when_set():
    client, calls = _make_client(lambda _p: _reply(content="ok"))

    engine = ConversationEngine(LlmConfig(), client, system_prompt="SYS")
    engine.send("hi")
    assert "id_slot" not in calls[0]["body"]

    engine = ConversationEngine(LlmConfig(), client, system_prompt="SYS", slot_id=3)
    engine.send("hi")
    assert calls[1]["body"]["id_slot"] == 3


def test_engine_omits_slot_and_on_io_for_old_fake_when_none():
    record: list[dict] = []
    fake = _fake_engine_client(record)
    engine = ConversationEngine(LlmConfig(), fake, system_prompt="SYS")

    assert engine.send("hi") == "ok"
    assert record[0]["messages"][0] == {"role": "system", "content": "SYS"}
    assert record[0]["max_tokens"] is None
    assert record[0]["timeout_s"] is None


def test_engine_forwards_per_call_on_io():
    client, _ = _make_client(lambda _p: _reply(content="ok"))
    engine = ConversationEngine(LlmConfig(), client, system_prompt="SYS")
    seen: list[dict] = []

    assert engine.send("hi", on_io=seen.append) == "ok"
    assert [event["phase"] for event in seen] == ["request", "response"]


def test_engine_forwards_slot_and_on_io_when_set():
    record: list[dict] = []

    class _NewFake:
        def chat(self, messages, max_tokens=None, timeout_s=None, **kwargs):
            record.append(kwargs)
            return LlmResponse(content="ok", finish_reason="stop")

    def hook(_event):
        pass

    engine = ConversationEngine(
        LlmConfig(), _NewFake(), system_prompt="SYS", slot_id=7
    )
    assert engine.send("hi", on_io=hook) == "ok"
    assert record[0]["slot_id"] == 7
    assert record[0]["on_io"] is hook


def test_engine_forwards_max_tokens_and_timeout():
    record: list[dict] = []
    fake = _fake_engine_client(record)
    engine = ConversationEngine(
        LlmConfig(), fake, system_prompt="SYS", max_tokens=11, timeout_s=2.5
    )

    engine.send("hi")
    assert record[0]["max_tokens"] == 11
    assert record[0]["timeout_s"] == 2.5


def test_last_reply_returns_latest_assistant():
    fake = _fake_engine_client(content="first")
    engine = ConversationEngine(LlmConfig(), fake, system_prompt="SYS")

    assert engine.last_reply() == ""
    engine.send("hi")
    assert engine.last_reply() == "first"
    engine.history.append({"role": "user", "content": "again"})
    engine.history.append({"role": "assistant", "content": "second"})
    assert engine.last_reply() == "second"


def test_finalize_none_restores_history_and_summary_copy():
    notifications: list[tuple] = []

    def notify(kind, **fields):
        notifications.append((kind, fields))

    client, _ = _make_client(lambda _p: _reply(content="raw reply"))
    engine = ConversationEngine(LlmConfig(), client, system_prompt="SYS", notify=notify)
    engine.history = [
        {"role": "user", "content": "u1"},
        {"role": "assistant", "content": "a1"},
    ]
    engine.summary = "prior summary"

    assert engine.send("new?", finalize=lambda reply, finish: None) == ""
    assert engine.history == [
        {"role": "user", "content": "u1"},
        {"role": "assistant", "content": "a1"},
    ]
    assert engine.summary == "prior summary"
    assert notifications == []


def test_finalize_transformed_string_is_appended():
    client, _ = _make_client(lambda _p: _reply(content="raw"))
    engine = ConversationEngine(LlmConfig(), client, system_prompt="SYS")

    seen_args: list[tuple] = []

    def finalize(reply, finish):
        seen_args.append((reply, finish))
        return "EDITED: " + reply

    assert engine.send("hi", finalize=finalize) == "EDITED: raw"
    assert seen_args == [("raw", "stop")]
    assert engine.history[-1] == {"role": "assistant", "content": "EDITED: raw"}
    assert engine.last_reply() == "EDITED: raw"


def test_finalize_empty_string_is_appended_not_restored():
    client, _ = _make_client(lambda _p: _reply(content="raw"))
    engine = ConversationEngine(LlmConfig(), client, system_prompt="SYS")
    engine.history = [
        {"role": "user", "content": "u1"},
        {"role": "assistant", "content": "a1"},
    ]

    assert engine.send("new?", finalize=lambda reply, finish: "") == ""
    assert engine.history == [
        {"role": "user", "content": "u1"},
        {"role": "assistant", "content": "a1"},
        {"role": "user", "content": "new?"},
        {"role": "assistant", "content": ""},
    ]
    assert engine.history[-1] == {"role": "assistant", "content": ""}


def test_finalize_reject_after_reset_mid_flight_stays_empty():
    client, _ = _make_client(lambda _p: _reply(content="late"))

    holder: list[ConversationEngine] = []
    engine = ConversationEngine(LlmConfig(), client, system_prompt="SYS")
    holder.append(engine)
    engine.summary = "prior summary"

    def finalize(reply, finish):
        holder[0].reset()
        return None

    assert engine.send("hi", finalize=finalize) == ""
    assert engine.history == []
    assert engine.summary == ""


def test_finalize_snapshot_is_a_copy():
    client, _ = _make_client(lambda _p: _reply(content="raw"))
    engine = ConversationEngine(LlmConfig(), client, system_prompt="SYS")
    engine.history = [
        {"role": "user", "content": "u1"},
        {"role": "assistant", "content": "a1"},
    ]

    def finalize(reply, finish):
        engine.history.append({"role": "assistant", "content": "POISON"})
        engine.history.append({"role": "user", "content": "POISON"})
        return None

    assert engine.send("new?", finalize=finalize) == ""
    assert engine.history == [
        {"role": "user", "content": "u1"},
        {"role": "assistant", "content": "a1"},
    ]


def test_summarizer_never_pinned_to_engine_slot():
    cfg = _compact_cfg(budget_chars=200, trigger_ratio=0.8, keep_turns=1)
    client, calls = _make_client(
        lambda p: (
            _reply(content="s")
            if p["body"]["messages"][0]["content"] == "CUSTOM-COMPACT"
            else _reply(content="ok")
        ),
        cfg,
    )
    engine = ConversationEngine(
        cfg,
        client,
        system_prompt="SYS",
        compaction_system_prompt="CUSTOM-COMPACT",
        slot_id=5,
    )

    for _ in range(4):
        engine.send("x" * 40)

    chat_calls = [c for c in calls if c["body"]["messages"][0]["content"] == "SYS"]
    compact_calls = [
        c for c in calls if c["body"]["messages"][0]["content"] == "CUSTOM-COMPACT"
    ]
    assert chat_calls and all(c["body"]["id_slot"] == 5 for c in chat_calls)
    assert compact_calls and all("id_slot" not in c["body"] for c in compact_calls)


def test_chat_session_is_engine_with_chat_prompt():
    client, calls = _make_client(lambda _p: _reply(content="hello"))
    session = ChatSession(LlmConfig(), client)

    assert isinstance(session, ConversationEngine)
    assert session.system == LlmConfig().system_prompt_chat
    assert session.send("hi") == "hello"
    assert calls[0]["body"]["messages"][0]["content"] == LlmConfig().system_prompt_chat
    assert session.last_reply() == "hello"


def test_brief_engine_sends_static_system_and_single_user_turn():
    client, calls = _make_client(lambda _p: _reply(content="BRIEF v1"))
    engine = BriefEngine(LlmConfig(), client, system_prompt="SYS", max_tokens=99)

    assert engine.send("CURRENT LEDGER:\n(none yet)") == "BRIEF v1"

    assert len(calls) == 1
    messages = calls[0]["body"]["messages"]
    assert messages[0] == {"role": "system", "content": "SYS"}
    assert len(messages) == 2
    assert messages[1]["content"].startswith("CURRENT LEDGER:")
    assert calls[0]["body"]["max_tokens"] == 99


def test_brief_engine_second_call_still_sends_two_messages():
    client, calls = _make_client(lambda _p: _reply(content="BRIEF"))
    engine = BriefEngine(LlmConfig(), client, system_prompt="SYS")

    engine.send("first")
    engine.send("second")

    assert len(calls) == 2
    for call in calls:
        assert len(call["body"]["messages"]) == 2


def test_brief_engine_is_stateless_between_calls():
    client, calls = _make_client(lambda _p: _reply(content="BRIEF"))
    engine = BriefEngine(LlmConfig(), client, system_prompt="SYS")

    engine.send("first")
    engine.send("second")

    assert calls[0]["body"]["messages"][1]["content"] == "first"
    assert calls[1]["body"]["messages"][1]["content"] == "second"


def test_brief_engine_finalize_none_returns_empty():
    client, _ = _make_client(lambda _p: _reply(content="raw"))
    engine = BriefEngine(LlmConfig(), client, system_prompt="SYS")

    assert engine.send("q", finalize=lambda reply, finish: None) == ""


def test_brief_engine_finalize_transform_is_returned():
    client, _ = _make_client(lambda _p: _reply(content="raw"))
    engine = BriefEngine(LlmConfig(), client, system_prompt="SYS")

    result = engine.send("q", finalize=lambda reply, finish: "CLEANED: " + reply)

    assert result == "CLEANED: raw"


def test_brief_engine_length_reply_is_returned_with_finish_reason():
    client, _ = _make_client(lambda _p: _reply(content="half", finish="length"))
    engine = BriefEngine(LlmConfig(), client, system_prompt="SYS")
    seen: list[str] = []

    result = engine.send("q", finalize=lambda reply, finish: seen.append(finish) or reply)

    assert result == "half"
    assert seen == ["length"]


def test_brief_engine_reset_invalidates_in_flight_reply():
    client, _ = _make_client(lambda _p: _reply(content="BRIEF"))
    engine = BriefEngine(LlmConfig(), client, system_prompt="SYS")

    def finalize(reply, finish):
        engine.reset()
        return reply

    assert engine.send("q", finalize=finalize) == ""


def test_brief_engine_passes_slot_and_io_hook():
    client, calls = _make_client(lambda _p: _reply(content="BRIEF"))
    engine = BriefEngine(LlmConfig(), client, system_prompt="SYS", slot_id=5)
    seen: list[dict] = []

    engine.send("q", on_io=lambda part: seen.append(part))

    assert calls[0]["body"]["id_slot"] == 5
    assert any(part["phase"] == "request" for part in seen)