"""Unit tests for navi.core.agent.Agent.

Uses InMemorySessionStore, FakeLLMBackend, and FakeTool so tests run
without a real database or LLM server.
"""

import asyncio
import copy
from datetime import datetime, timezone

import pytest
import pytest_asyncio

from navi.core.agent import Agent, _is_casual_message
from navi.core.events import (
    CompressionStarted,
    ContextCompressed,
    ModelInfo,
    StreamEnd,
    SubagentComplete,
)
from navi.core.registry import BackendRegistry, ProfileRegistry
from navi.core.session import InMemorySessionStore, Session
from navi.exceptions import MaxIterationsReached, NothingToCompactError, SessionNotFound
from navi.llm.base import LLMChunk, Message, ToolCallRequest
from tests.conftest_factory import FakeLLMBackend, make_profile, make_registry_with_tools


@pytest.fixture
def agent():
    sessions = InMemorySessionStore()
    profiles = ProfileRegistry()
    profile = make_profile("test")
    profile.planning_phase1_enabled = False
    profile.planning_phase2_enabled = False
    profile.planning_phase3_enabled = False
    profiles.register(profile)
    tools = make_registry_with_tools()
    backends = BackendRegistry()
    backends.register("ollama", FakeLLMBackend(responses=["hello"]))
    return Agent(
        session_store=sessions,
        profile_registry=profiles,
        tool_registry=tools,
        backend_registry=backends,
    )


@pytest_asyncio.fixture
async def session(agent):
    return await agent._sessions.create(profile_id="test")


# ─── run() tests ───────────────────────────────────────────────────────────


class TestAgentRun:
    @pytest.mark.asyncio
    async def test_run_single_iteration(self, agent, session):
        backend = FakeLLMBackend(responses=["hello"])
        agent._backends.register("ollama", backend)

        result = await agent.run(session.id, "hi")
        assert result == "hello"
        saved = await agent._sessions.get(session.id)
        # user display + user context + assistant
        assert len(saved.messages) == 3
        assert saved.messages[0].role == "user"
        assert saved.messages[1].role == "user"
        assert saved.messages[2].role == "assistant"
        assert saved.messages[2].content == "hello"

        # Flags: display-only user message, context-only user message
        assert saved.messages[0].is_display is True
        assert saved.messages[0].is_context is False
        assert saved.messages[1].is_display is False
        assert saved.messages[1].is_context is True
        # Assistant message is both display and context
        assert saved.messages[2].is_display is True
        assert saved.messages[2].is_context is True

    @pytest.mark.asyncio
    async def test_run_session_not_found(self, agent):
        with pytest.raises(SessionNotFound):
            await agent.run("nonexistent-id", "hi")

    @pytest.mark.asyncio
    async def test_run_tool_calls_then_stop(self, agent, session):
        """Tool-calling turn followed by a final stop turn."""
        backend = FakeLLMBackend(
            responses=["", "done"],
            tool_calls=[
                [ToolCallRequest(id="1", name="test_tool", arguments={})],
                None,
            ],
        )
        agent._backends.register("ollama", backend)

        result = await agent.run(session.id, "do something")
        assert result == "done"
        saved = await agent._sessions.get(session.id)
        # user display + user context + assistant(tool) + tool_result + assistant(final)
        assert len(saved.messages) == 5
        assert saved.messages[3].role == "tool"
        assert saved.messages[4].content == "done"

    @pytest.mark.asyncio
    async def test_run_token_accumulation(self, agent, session):
        """_turn_tokens accumulates completion tokens across tool-calling iterations."""
        backend = FakeLLMBackend(
            responses=["", "done"],
            tool_calls=[
                [ToolCallRequest(id="1", name="test_tool", arguments={})],
                None,
            ],
            prompt_tokens=10,
            completion_tokens=5,
        )
        agent._backends.register("ollama", backend)

        await agent.run(session.id, "do something")
        saved = await agent._sessions.get(session.id)
        final_msg = saved.messages[-1]
        # Two iterations × 5 completion tokens = 10 tokens
        assert final_msg.token_count == 10

    @pytest.mark.asyncio
    async def test_run_max_iterations(self, agent, session):
        """After max_iterations tool turns, MaxIterationsReached is raised."""
        profile = agent._profiles.get("test")
        profile.max_iterations = 2

        backend = FakeLLMBackend(
            responses=["", ""],
            tool_calls=[
                [ToolCallRequest(id="1", name="test_tool", arguments={})],
                [ToolCallRequest(id="2", name="test_tool", arguments={})],
            ],
        )
        agent._backends.register("ollama", backend)

        with pytest.raises(MaxIterationsReached):
            await agent.run(session.id, "loop forever")


# ─── run_stream() tests ──────────────────────────────────────────────────────


class TestAgentRunStream:
    @pytest.mark.asyncio
    async def test_run_stream_single_iteration(self, agent, session):
        backend = FakeLLMBackend(responses=["streamed hello"])
        agent._backends.register("ollama", backend)

        events = []
        async for ev in agent.run_stream(session.id, "hi"):
            events.append(type(ev).__name__)

        assert events[-1] == "StreamEnd"
        saved = await agent._sessions.get(session.id)
        assert saved.messages[-1].content == "streamed hello"

    @pytest.mark.asyncio
    async def test_run_stream_emits_model_info(self, agent, session):
        """The agent emits a ModelInfo event carrying the resolved model."""
        from typing import AsyncGenerator

        from navi.llm.base import LLMBackend

        class ModelStampingBackend(LLMBackend):
            async def complete(self, messages, tools=None, temperature=0.7, model=None,
                               think=None, max_tokens=None, **kw):
                raise NotImplementedError

            async def stream_complete(self, messages, tools=None, temperature=0.7,
                                      model=None, think=None, **kw) -> AsyncGenerator[LLMChunk, None]:
                yield LLMChunk(delta="hello", model="resolved-model")
                yield LLMChunk(finish_reason="stop", prompt_tokens=5, completion_tokens=1)

            async def embed(self, texts, model=None):
                return [[0.1] * 768 for _ in texts]

        agent._backends.register("ollama", ModelStampingBackend())
        events = []
        async for ev in agent.run_stream(session.id, "hi"):
            events.append(ev)

        infos = [ev for ev in events if isinstance(ev, ModelInfo)]
        assert len(infos) == 1
        assert infos[0].model == "resolved-model"

    @pytest.mark.asyncio
    async def test_run_stream_emits_model_info_once_per_turn(self, agent, session):
        """ModelInfo is not re-emitted across iterations if the model stays the same."""
        from typing import AsyncGenerator

        from navi.llm.base import LLMBackend

        class ModelStampingBackend(LLMBackend):
            def __init__(self):
                self._call = 0

            async def complete(self, messages, tools=None, temperature=0.7, model=None,
                               think=None, max_tokens=None, **kw):
                raise NotImplementedError

            async def stream_complete(self, messages, tools=None, temperature=0.7,
                                      model=None, think=None, **kw) -> AsyncGenerator[LLMChunk, None]:
                self._call += 1
                if self._call == 1:
                    yield LLMChunk(
                        delta="", model="same-model",
                        finish_reason="tool_calls",
                        tool_calls=[ToolCallRequest(id="1", name="test_tool", arguments={})],
                    )
                else:
                    yield LLMChunk(delta="final answer", model="same-model")
                    yield LLMChunk(finish_reason="stop", prompt_tokens=5, completion_tokens=1)

            async def embed(self, texts, model=None):
                return [[0.1] * 768 for _ in texts]

        agent._backends.register("ollama", ModelStampingBackend())
        events = []
        async for ev in agent.run_stream(session.id, "do something"):
            events.append(ev)

        # Two iterations happened (tool call then final), but ModelInfo fires once.
        infos = [ev for ev in events if isinstance(ev, ModelInfo)]
        assert len(infos) == 1
        assert infos[0].model == "same-model"

    @pytest.mark.asyncio
    async def test_run_stream_tool_calls(self, agent, session):
        backend = FakeLLMBackend(
            responses=["", "final"],
            tool_calls=[
                [ToolCallRequest(id="1", name="test_tool", arguments={})],
                None,
            ],
        )
        agent._backends.register("ollama", backend)

        events = []
        async for ev in agent.run_stream(session.id, "do something"):
            events.append(type(ev).__name__)

        assert "ToolStarted" in events
        assert "ToolEvent" in events
        assert events[-1] == "StreamEnd"

    @pytest.mark.asyncio
    async def test_run_stream_emits_todo_updated_after_tool_turn(self, agent, session):
        """A TodoUpdated event is emitted after each tool-execution turn so the
        UI side panel can reflect todo changes from the todo tool."""
        from navi.core.events import TodoUpdated

        backend = FakeLLMBackend(
            responses=["", "final"],
            tool_calls=[
                [ToolCallRequest(id="1", name="test_tool", arguments={})],
                None,
            ],
        )
        agent._backends.register("ollama", backend)

        events = []
        async for ev in agent.run_stream(session.id, "do something"):
            events.append(ev)

        todo_events = [ev for ev in events if isinstance(ev, TodoUpdated)]
        assert len(todo_events) >= 1
        # The event carries the session id and a (possibly empty) tasks list.
        assert todo_events[0].session_id == session.id
        assert isinstance(todo_events[0].tasks, list)

    @pytest.mark.asyncio
    async def test_run_stream_stop_event(self, agent, session):
        """Cooperative stop mid-stream yields StreamStopped."""
        from navi.tools._internal.base import current_stop_event

        stop = asyncio.Event()
        token = current_stop_event.set(stop)
        try:
            async def _slow_stream(self, **kwargs):
                yield LLMChunk(delta="a")
                await asyncio.sleep(10)
                yield LLMChunk(delta="b")

            backend = FakeLLMBackend()
            # Monkey-patch stream_complete to be slow
            backend.stream_complete = _slow_stream
            agent._backends.register("ollama", backend)

            stop.set()
            events = []
            async for ev in agent.run_stream(session.id, "hi"):
                events.append(type(ev).__name__)

            assert "StreamStopped" in events
        finally:
            current_stop_event.reset(token)

    @pytest.mark.asyncio
    async def test_run_stream_token_count(self, agent, session):
        backend = FakeLLMBackend(
            responses=["final"],
            prompt_tokens=100,
            completion_tokens=50,
        )
        agent._backends.register("ollama", backend)

        events = []
        async for ev in agent.run_stream(session.id, "hi"):
            if isinstance(ev, StreamEnd):
                events.append(ev)

        assert events[0].token_count == 50
        saved = await agent._sessions.get(session.id)
        assert saved.messages[-1].token_count == 50


# ─── compact_stream() tests ──────────────────────────────────────────────────


class TestAgentCompactStream:
    """Forced /compact: bypasses the token threshold and runs the real compressor."""

    @pytest_asyncio.fixture
    async def session_with_history(self, agent, session):
        """A session with enough turns that there is something to summarize."""
        # context_keep_recent defaults to 8 turns; add 10 so the oldest turns
        # fall into to_summarize (>= 2 messages) and compress_context does work.
        for i in range(10):
            session.context.append(Message(role="user", content=f"user {i}"))
            session.context.append(Message(role="assistant", content=f"assistant {i}"))
            session.messages.append(Message(role="user", content=f"user {i}"))
            session.messages.append(Message(role="assistant", content=f"assistant {i}"))
        await agent._sessions.save(session)
        return session

    @pytest.mark.asyncio
    async def test_compact_emits_started_then_compressed(self, agent, session_with_history):
        # The agent fixture's FakeLLMBackend returns "hello" — a valid summary.
        events = [ev async for ev in agent.compact_stream(session_with_history.id)]

        assert isinstance(events[0], CompressionStarted)
        assert isinstance(events[-1], ContextCompressed)
        # Compression actually shrank the LLM context (10 turns -> summary + kept recent).
        assert events[-1].messages_after < events[-1].messages_before

    @pytest_asyncio.fixture
    async def session_one_long_turn(self, agent, session):
        """The navi_code shape: a single long autonomous turn — one user message
        followed by many tool/assistant iterations. This is ONE turn, so the
        turn-based partition (len(turns) <= keep_recent) finds nothing to
        summarize. Forced compact must still compress via the intra-turn split
        (keep_recent_messages), otherwise the user always sees
        "Nothing to compact yet — the context is still small" regardless of how
        long the turn grew. See keep_recent_messages in compact_stream."""
        session.context.append(Message(role="user", content="do the task"))
        session.messages.append(Message(role="user", content="do the task"))
        for i in range(20):
            session.context.append(Message(role="assistant", content=f"step {i}"))
            session.messages.append(Message(role="assistant", content=f"step {i}"))
        await agent._sessions.save(session)
        return session

    @pytest.mark.asyncio
    async def test_compact_compresses_single_long_turn(self, agent, session_one_long_turn):
        """Regression: a single long turn (navi_code shape) must compress, not
        report 'nothing to compact'. The intra-turn fallback in compact_stream
        (keep_recent_messages=max(12, context_keep_recent*2)) is what makes this
        work — without it forced compact always returns NothingToCompactError
        once the conversation is a single turn."""
        events = [ev async for ev in agent.compact_stream(session_one_long_turn.id)]

        assert isinstance(events[0], CompressionStarted)
        assert isinstance(events[-1], ContextCompressed)
        assert events[-1].messages_after < events[-1].messages_before

    @pytest.mark.asyncio
    async def test_compact_bypasses_threshold_even_when_context_small(self, agent, session):
        """Forced compact is NOT gated by should_compress — but still needs >= 2
        summarizable messages. A near-empty session raises NothingToCompactError
        (not silently does nothing), which the orchestrator surfaces as feedback."""
        session.context.append(Message(role="user", content="hi"))
        session.context.append(Message(role="assistant", content="hello"))
        await agent._sessions.save(session)

        with pytest.raises(NothingToCompactError):
            async for _ in agent.compact_stream(session.id):
                pass

    @pytest.mark.asyncio
    async def test_compact_raises_when_compression_disabled(self, agent, session_with_history, monkeypatch):
        import navi.core.agent as agent_mod
        import navi.config as config

        # Settings is frozen — swap the whole object the agent module sees.
        disabled = config.Settings(
            database_url=config.settings.database_url,
            context_compression_enabled=False,
        )
        monkeypatch.setattr(agent_mod, "settings", disabled)
        with pytest.raises(NothingToCompactError):
            async for _ in agent.compact_stream(session_with_history.id):
                pass

    @pytest.mark.asyncio
    async def test_compact_unknown_session_raises(self, agent):
        with pytest.raises(SessionNotFound):
            async for _ in agent.compact_stream("no-such-session"):
                pass

    @pytest.mark.asyncio
    async def test_midturn_no_compression_started_when_nothing_to_compress(self, agent, session, monkeypatch):
        """C: even when the token gate fires, CompressionStarted is NOT emitted
        if would_compress says the partition cannot shrink the context — no
        misleading "compression" status with no work done. The gate is forced
        True here so we exercise the would_compress guard directly."""
        import navi.core.agent as agent_mod
        from navi.core.events import CompressionStarted

        # Small context that partition cannot shrink (would_compress False).
        session.context.append(Message(role="user", content="hi"))
        session.context.append(Message(role="assistant", content="hello"))
        await agent._sessions.save(session)

        # Force the token gate open; without this the small context would not
        # reach the would_compress check at all.
        monkeypatch.setattr(agent_mod, "should_compress", lambda *a, **k: True)

        events = [
            ev
            async for ev in agent._compression_events_midturn(
                session,
                llm=agent._get_backend("ollama"),
                profile=agent._profiles.get("test"),
                session_id=session.id,
                iteration=1,
                ctx_injections=[],
                mem=None,
            )
        ]
        assert not any(isinstance(ev, CompressionStarted) for ev in events)


# ─── run_ephemeral() tests ───────────────────────────────────────────────────


class TestAgentRunEphemeral:
    @pytest.mark.asyncio
    async def test_run_ephemeral_complete(self, agent):
        backend = FakeLLMBackend(responses=["subagent result"])
        agent._backends.register("ollama", backend)

        result, ok = await agent.run_ephemeral("task", profile_id="test")
        assert "subagent result" in result
        assert "[Sub-agent stopped: completed]" in result
        assert ok is True

    @pytest.mark.asyncio
    async def test_run_ephemeral_max_iterations(self, agent):
        backend = FakeLLMBackend(
            responses=[""],
            tool_calls=[
                [ToolCallRequest(id="1", name="test_tool", arguments={})],
            ],
        )
        agent._backends.register("ollama", backend)

        result, ok = await agent.run_ephemeral(
            "task", profile_id="test", max_iterations=1
        )
        assert ok is False
        assert "iteration limit" in result.lower()

    @pytest.mark.skip(reason="run_ephemeral uses 'import time as _time' inside the function; CPython LOAD_GLOBAL caching makes module-level mock replacement unreliable in pytest-asyncio.")
    @pytest.mark.asyncio
    async def test_run_ephemeral_timeout(self, agent):
        pass

    @pytest.mark.asyncio
    async def test_run_ephemeral_planning_tokens_accumulated(self, agent):
        """Planning phase AIHelperTokensUsed contributes to SubagentComplete."""
        from navi.core.events import AIHelperTokensUsed
        from navi.tools._internal.base import current_event_sink

        backend = FakeLLMBackend(responses=["final"])
        agent._backends.register("ollama", backend)

        # Force planning by setting subagent_planning_enabled on profile
        profile = agent._profiles.get("test")
        profile.subagent_planning_enabled = True

        sink = asyncio.Queue()
        token = current_event_sink.set(sink)
        try:
            # Mock planning to emit AIHelperTokensUsed
            original_planning_run = agent._planning.run

            async def _mock_planning(*args, **kwargs):
                yield AIHelperTokensUsed(prompt_tokens=5, completion_tokens=10)
                yield AIHelperTokensUsed(prompt_tokens=3, completion_tokens=7)

            agent._planning.run = _mock_planning

            result, ok = await agent.run_ephemeral("task", profile_id="test")
            assert ok is True

            # Drain sink for SubagentComplete
            subagent_complete = None
            while not sink.empty():
                item = await sink.get()
                if isinstance(item, SubagentComplete):
                    subagent_complete = item

            # Planning completion tokens: 10 + 7 = 17
            # Final LLM call: 0 (no tokens in FakeLLMBackend default)
            assert subagent_complete is not None
            assert subagent_complete.token_count == 17
        finally:
            current_event_sink.reset(token)
            agent._planning.run = original_planning_run

    @pytest.mark.asyncio
    async def test_run_ephemeral_thinking_stall(self, agent):
        """Subagent that produces only thinking for too long is aborted."""
        async def _thinking_only(self, **kwargs):
            for _ in range(200):
                yield LLMChunk(thinking="thinking " * 100)
            yield LLMChunk(delta="done", finish_reason="stop")

        backend = FakeLLMBackend()
        backend.stream_complete = _thinking_only
        agent._backends.register("ollama", backend)

        result, ok = await agent.run_ephemeral("task", profile_id="test")
        assert ok is False
        assert "thinking" in result.lower() or "stall" in result.lower()


# ─── _is_casual_message heuristic tests ──────────────────────────────────────


class TestIsCasualMessage:
    def test_greetings_are_casual(self):
        assert _is_casual_message("привет")
        assert _is_casual_message("hi")
        assert _is_casual_message("Hello")
        assert _is_casual_message("здравствуйте")

    def test_social_phrases_are_casual(self):
        assert _is_casual_message("как дела?")
        assert _is_casual_message("how are you")
        assert _is_casual_message("спасибо")
        assert _is_casual_message("thanks")
        assert _is_casual_message("пока")

    def test_very_short_messages_are_casual(self):
        assert _is_casual_message("ok")
        assert _is_casual_message("да")
        assert _is_casual_message("yo")

    def test_tool_or_command_markers_are_not_casual(self):
        assert not _is_casual_message("/help")
        assert not _is_casual_message("!restart")
        assert not _is_casual_message("ping @user")

    def test_urls_and_paths_are_not_casual(self):
        assert not _is_casual_message("https://example.com")
        assert not _is_casual_message("read /home/user/file.py")
        assert not _is_casual_message("~/notes.md")

    def test_long_messages_are_not_casual(self):
        assert not _is_casual_message("привет, расскажи подробно как настроить сервер и что для этого нужно сделать")

    def test_actual_task_requests_are_not_casual(self):
        assert not _is_casual_message("напиши скрипт для бэкапа")
        assert not _is_casual_message("проверь почему не работает ssh")
        assert not _is_casual_message("what is the weather in Berlin tomorrow")


# ─── Persistence-on-crash tests (B1: incremental flush) ─────────────────────


class _SnapshotSessionStore(InMemorySessionStore):
    """Session store that mimics a real DB boundary.

    ``save()`` deep-copies the session into a snapshot; ``get()`` returns the last
    saved snapshot (a fresh copy), not the live object the agent is mutating. So
    mutations that were never ``save()``d are invisible to ``get()`` — exactly like
    ``PgSessionStore`` where messages with ``sequence_number < 0`` only persist on
    ``save()``. This lets us assert that a crash mid-turn does not lose work that
    was incrementally flushed (B1).
    """

    def __init__(self) -> None:
        super().__init__()
        self._snapshots: dict[str, Session] = {}

    async def save(self, session: Session) -> None:
        session.last_active = datetime.now(timezone.utc)
        # Mirror PgSessionStore: assign sequence numbers to new messages on save.
        for m in session.messages:
            if m.sequence_number < 0:
                m.sequence_number = session.db_next_sequence
                session.db_next_sequence += 1
        self._snapshots[session.id] = copy.deepcopy(session)
        self._sessions[session.id] = session

    async def get(self, session_id: str) -> Session | None:
        snap = self._snapshots.get(session_id)
        if snap is not None:
            return copy.deepcopy(snap)
        return self._sessions.get(session_id)


class _CrashBackend(FakeLLMBackend):
    """FakeLLMBackend whose ``stream_complete`` raises on the Nth call.

    Simulates a server crash / CancelledError arriving mid-turn, so we can assert
    that tool work completed before the crash was already persisted (B1).
    """

    def __init__(
        self,
        responses: list[str],
        tool_calls: list[list[ToolCallRequest] | None],
        raise_on_call: int,
        exc: BaseException = RuntimeError("simulated crash"),
    ) -> None:
        super().__init__(responses=responses, tool_calls=tool_calls)
        self._raise_on_call = raise_on_call
        self._exc = exc
        self._sc_idx = 0

    async def stream_complete(self, messages, tools=None, temperature=0.7, model=None,
                              think=None, **kwargs):
        idx = self._sc_idx
        self._sc_idx += 1
        if idx == self._raise_on_call:
            raise self._exc
        content = self._responses[idx] if idx < len(self._responses) else ""
        tcalls = self._tool_calls[idx] if idx < len(self._tool_calls) else None
        if content:
            yield LLMChunk(delta=content)
        yield LLMChunk(
            finish_reason="tool_calls" if tcalls else "stop",
            tool_calls=tcalls,
        )


def _build_persistence_agent(store: InMemorySessionStore, backend) -> Agent:
    profiles = ProfileRegistry()
    profile = make_profile("test")
    # Disable planning so it neither appends plan messages nor consumes LLM calls —
    # the test controls backend calls precisely via stream_complete.
    profile.planning_phase1_enabled = False
    profile.planning_phase2_enabled = False
    profile.planning_phase3_enabled = False
    profiles.register(profile)
    tools = make_registry_with_tools()
    backends = BackendRegistry()
    backends.register("ollama", backend)
    return Agent(
        session_store=store,
        profile_registry=profiles,
        tool_registry=tools,
        backend_registry=backends,
    )


class TestAgentPersistenceOnCrash:
    @pytest.mark.asyncio
    async def test_tool_results_persisted_on_crash_midturn(self):
        """A crash on a later LLM call must not lose tool results already completed
        in earlier iterations — B1 persists each tool result as it completes."""
        store = _SnapshotSessionStore()
        backend = _CrashBackend(
            responses=["", "unused"],
            tool_calls=[
                [ToolCallRequest(id="1", name="test_tool", arguments={})],
                None,
            ],
            raise_on_call=1,
        )
        agent = _build_persistence_agent(store, backend)
        sess = await agent._sessions.create(profile_id="test")

        with pytest.raises(RuntimeError, match="simulated crash"):
            async for _ in agent.run_stream(sess.id, "do work"):
                pass

        saved = await agent._sessions.get(sess.id)
        roles = [m.role for m in saved.messages]
        # Iteration 0's assistant tool-call decision + tool result were flushed by
        # B1 before the crash on iteration 1. Without B1 only the user message
        # (saved at turn start) would survive.
        assert "tool" in roles, f"tool result missing from saved state: {roles}"
        assert any(
            m.role == "assistant" and m.tool_calls for m in saved.messages
        ), f"assistant tool-call msg missing: {roles}"

    @pytest.mark.asyncio
    async def test_cancel_flushes_completed_tool_work(self):
        """CancelledError mid-turn (server restart/shutdown) must not lose tool
        results already completed — B1 flushed them incrementally before the cancel."""
        store = _SnapshotSessionStore()
        backend = _CrashBackend(
            responses=["", "unused"],
            tool_calls=[
                [ToolCallRequest(id="1", name="test_tool", arguments={})],
                None,
            ],
            raise_on_call=1,
            exc=asyncio.CancelledError(),
        )
        agent = _build_persistence_agent(store, backend)
        sess = await agent._sessions.create(profile_id="test")

        with pytest.raises(asyncio.CancelledError):
            async for _ in agent.run_stream(sess.id, "do work"):
                pass

        saved = await agent._sessions.get(sess.id)
        roles = [m.role for m in saved.messages]
        assert "tool" in roles, f"tool result missing after cancel: {roles}"