"""Tests for DialogueRunner — usage capture, message finalization, and billing context."""

import json
from collections import defaultdict
from types import SimpleNamespace
from typing import Any, AsyncIterator, Callable, Dict, List, Optional
from unittest.mock import AsyncMock, MagicMock, patch
from uuid import uuid4

import pytest

from swarm.agent import Agent
from swarm.chunk import (
    AssistantMessageChunk,
    ChunkType,
    DelimiterChunk,
    TokenUsage,
    ToolResponseChunk,
)
from swarm.context import RunContext
from swarm.message import SwarmAssistantMessage, SwarmToolMessage, SwarmUserMessage
from swarm.providers.protocols import ModelRouterProtocol
from swarm.runners.dialogue import DialogueRunner, _LoopState
from swarm.tool import ResultMetadata, ToolExecutionResponse
from swarm.types import ContentType


def _make_state(
    *,
    agent_model: str | list[str] = "gpt-4o",
    current_call_usage: TokenUsage | None = None,
) -> _LoopState:
    """Build a minimal _LoopState for unit tests."""
    agent = SimpleNamespace(
        name="test-agent",
        id="agent-1",
        model=agent_model,
        instructions="hello",
        tools=[],
        model_settings=None,
    )
    state = _LoopState(
        message_history=[],
        current_message_uuid="msg-uuid-1",
        initial_len=0,
        active_agent=agent,
        system_prompt="system",
        context_variables={},
    )
    state.current_call_usage = current_call_usage
    return state


def _make_message_accumulator(
    msg_id: str = "msg-uuid-1",
) -> Dict[str, Any]:
    return {
        "id": msg_id,
        "content": "Hello world",
        "sender": "agent-1",
        "role": "assistant",
        "function_call": None,
        "tool_calls": defaultdict(
            lambda: {
                "function": {"arguments": "", "name": ""},
                "id": "",
                "type": "",
            }
        ),
    }


class TestFinalizeAssistantMessage:
    def test_attaches_usage_when_present(self) -> None:
        usage = TokenUsage(model="gpt-4o", input_tokens=50, output_tokens=20)
        state = _make_state(current_call_usage=usage)
        runner = DialogueRunner()

        message = _make_message_accumulator()
        result = runner._finalize_assistant_message(message, state)

        assert isinstance(result, SwarmAssistantMessage)
        assert result.usage is not None
        assert result.usage.model == "gpt-4o"
        assert result.usage.input_tokens == 50
        assert result.usage.output_tokens == 20

    def test_no_usage_when_current_call_usage_is_none(self) -> None:
        state = _make_state(current_call_usage=None)
        runner = DialogueRunner()

        message = _make_message_accumulator()
        result = runner._finalize_assistant_message(message, state)

        assert isinstance(result, SwarmAssistantMessage)
        assert result.usage is None

    def test_message_added_to_history(self) -> None:
        state = _make_state()
        runner = DialogueRunner()

        message = _make_message_accumulator()
        result = runner._finalize_assistant_message(message, state)

        assert state.message_history[-1] is result


class TestProcessStreamChunkUsage:
    """Test that _process_stream_chunk correctly captures usage."""

    def test_stream_chunk_id_uses_runner_uuid_not_provider_completion_id(self) -> None:
        state = _make_state()
        runner = DialogueRunner()
        message = _make_message_accumulator(msg_id=state.current_message_uuid)

        chunk = SimpleNamespace(
            choices=[
                SimpleNamespace(
                    delta=SimpleNamespace(
                        json=lambda: '{"id":"chatcmpl_abc123","content":"hello"}'
                    )
                )
            ],
            usage=None,
        )

        chunks = runner._process_stream_chunk(chunk, message, state)

        assert len(chunks) == 1
        assistant_chunk = chunks[0]
        assert assistant_chunk.id == state.current_message_uuid
        assert assistant_chunk.id != "chatcmpl_abc123"
        assert message["id"] == state.current_message_uuid

    def test_usage_sets_current_call_usage_and_appends_accumulated(self) -> None:
        state = _make_state(agent_model="gpt-4o")
        runner = DialogueRunner()
        message = _make_message_accumulator()

        # Simulate a final chunk with usage
        chunk = SimpleNamespace(
            choices=[],
            usage=SimpleNamespace(prompt_tokens=100, completion_tokens=50),
            model="gpt-4o-2024-01-01",
        )

        runner._process_stream_chunk(chunk, message, state)

        assert state.current_call_usage is not None
        # Should prefer chunk.model
        assert state.current_call_usage.model == "gpt-4o-2024-01-01"
        assert state.current_call_usage.input_tokens == 100
        assert state.current_call_usage.output_tokens == 50

        # Should also be appended to accumulated_usage
        accumulated = state.execution_metadata.get("accumulated_usage", [])
        assert len(accumulated) == 1
        assert accumulated[0] is state.current_call_usage

    def test_model_resolution_list_agent_model(self) -> None:
        """When agent.model is a list and chunk has no model, use first element."""
        state = _make_state(agent_model=["gpt-4o", "gpt-4o-mini"])
        runner = DialogueRunner()
        message = _make_message_accumulator()

        # Chunk without .model attribute
        chunk = SimpleNamespace(
            choices=[],
            usage=SimpleNamespace(prompt_tokens=10, completion_tokens=5),
        )

        runner._process_stream_chunk(chunk, message, state)

        assert state.current_call_usage is not None
        assert state.current_call_usage.model == "gpt-4o"

    def test_model_resolution_string_agent_model(self) -> None:
        """When agent.model is a string and chunk has no model, use it directly."""
        state = _make_state(agent_model="claude-3-opus")
        runner = DialogueRunner()
        message = _make_message_accumulator()

        chunk = SimpleNamespace(
            choices=[],
            usage=SimpleNamespace(prompt_tokens=10, completion_tokens=5),
        )

        runner._process_stream_chunk(chunk, message, state)

        assert state.current_call_usage is not None
        assert state.current_call_usage.model == "claude-3-opus"

    def test_model_resolution_empty_list_falls_back_to_unknown(self) -> None:
        """When agent.model is an empty list and chunk has no model, use 'unknown'."""
        state = _make_state(agent_model=[])
        runner = DialogueRunner()
        message = _make_message_accumulator()

        chunk = SimpleNamespace(
            choices=[],
            usage=SimpleNamespace(prompt_tokens=10, completion_tokens=5),
        )

        runner._process_stream_chunk(chunk, message, state)

        assert state.current_call_usage is not None
        assert state.current_call_usage.model == "unknown"

    def test_current_call_usage_reset_between_iterations(self) -> None:
        """Verify the reset documented in the plan (tested via state initialization)."""
        state = _make_state(
            current_call_usage=TokenUsage(
                model="old-model", input_tokens=99, output_tokens=99
            )
        )
        # The reset happens in run_streamed before each completion loop
        # We verify that setting to None works
        state.current_call_usage = None
        assert state.current_call_usage is None


class TestPersistMessagesBillingContext:
    """Test that _persist_messages passes billing context to message store."""

    @pytest.mark.asyncio
    async def test_passes_user_id_audience_node_key(self) -> None:
        runner = DialogueRunner()
        message_store = MagicMock()
        message_store.persist_messages = AsyncMock()

        msg = SwarmAssistantMessage(
            id="msg-1",
            content=[{"type": "text", "text": "test"}],
        )

        context = SimpleNamespace(
            conversation_id="conv-1",
            designer_fullname="Designer A",
            user_id="user-123",
            audience="consumer",
            node_key="room_capture",
        )
        agent = SimpleNamespace(model="gpt-4o", name="test")

        await runner._persist_messages(
            messages=[msg],
            context=context,
            context_variables={},
            agent=agent,
            message_store=message_store,
        )

        message_store.persist_messages.assert_awaited_once()
        call_kwargs = message_store.persist_messages.await_args.kwargs
        assert call_kwargs["user_id"] == "user-123"
        assert call_kwargs["audience"] == "consumer"
        assert call_kwargs["node_key"] == "room_capture"

    @pytest.mark.asyncio
    async def test_passes_none_when_context_lacks_billing_attrs(self) -> None:
        runner = DialogueRunner()
        message_store = MagicMock()
        message_store.persist_messages = AsyncMock()

        msg = SwarmAssistantMessage(
            id="msg-1",
            content=[{"type": "text", "text": "test"}],
        )

        # Context without billing attributes
        context = SimpleNamespace(
            conversation_id="conv-1",
            designer_fullname="Designer A",
        )
        agent = SimpleNamespace(model="gpt-4o", name="test")

        await runner._persist_messages(
            messages=[msg],
            context=context,
            context_variables={},
            agent=agent,
            message_store=message_store,
        )

        call_kwargs = message_store.persist_messages.await_args.kwargs
        assert call_kwargs["user_id"] is None
        assert call_kwargs["audience"] is None
        assert call_kwargs["node_key"] is None


# ── Helpers for _handle_display_tool tests ─────────────────────────────────


def _make_tool_msg(
    *,
    content_type: ContentType = ContentType.INTERACTIVE,
    data_text: Optional[str] = "Pick a room type",
    data_title: Optional[str] = None,
    top_title: Optional[str] = None,
    tool_name: str = "room_type_selection",
) -> SwarmToolMessage:
    """Build a SwarmToolMessage with component JSON content."""
    data: Dict[str, Any] = {}
    if data_text is not None:
        data["text"] = data_text
    if data_title is not None:
        data["title"] = data_title

    content: Dict[str, Any] = {"data": data}
    if top_title is not None:
        content["title"] = top_title

    return SwarmToolMessage(
        id=str(uuid4()),
        content=[{"type": "text", "text": json.dumps(content)}],
        tool_call_id=f"call-{uuid4()}",
        tool_name=tool_name,
        content_type=content_type,
    )


def _make_tool_response(
    tool_msg: SwarmToolMessage,
) -> ToolExecutionResponse:
    return ToolExecutionResponse(
        messages=[tool_msg],
        context_variables={},
        metadata=ResultMetadata(),
    )


def _make_fake_message_store() -> MagicMock:
    store = MagicMock()
    store.persist_messages = AsyncMock()
    return store


def _make_fake_stream_manager() -> MagicMock:
    mgr = MagicMock()
    mgr.publish_status = AsyncMock()
    return mgr


def _make_context() -> SimpleNamespace:
    return SimpleNamespace(
        conversation_id="conv-1",
        designer_fullname="Designer A",
        user_id="user-123",
        audience="consumer",
        node_key="room_capture",
    )


async def _collect_chunks(
    runner: DialogueRunner,
    tool_msg: SwarmToolMessage,
    tool_response: ToolExecutionResponse,
    state: _LoopState,
    context: SimpleNamespace,
    message_store: MagicMock,
    stream_manager: MagicMock,
) -> List[ChunkType]:
    chunks: List[ChunkType] = []
    async for chunk in runner._handle_display_tool(
        tool_msg, tool_response, state, context, message_store, stream_manager
    ):
        chunks.append(chunk)
    return chunks


class TestHandleDisplayToolInteractive:
    """INTERACTIVE display tools: no synthetic text chunks streamed."""

    @pytest.mark.asyncio
    async def test_interactive_no_text_chunks_streamed(self) -> None:
        runner = DialogueRunner()
        state = _make_state()
        tool_msg = _make_tool_msg(content_type=ContentType.INTERACTIVE)
        tool_response = _make_tool_response(tool_msg)
        message_store = _make_fake_message_store()
        stream_manager = _make_fake_stream_manager()

        chunks = await _collect_chunks(
            runner,
            tool_msg,
            tool_response,
            state,
            _make_context(),
            message_store,
            stream_manager,
        )

        # No AssistantMessageChunk or word-level DelimiterChunks
        assistant_chunks = [c for c in chunks if isinstance(c, AssistantMessageChunk)]
        assert assistant_chunks == []

        # ToolResponseChunk is still present
        tool_chunks = [c for c in chunks if isinstance(c, ToolResponseChunk)]
        assert len(tool_chunks) == 1

    @pytest.mark.asyncio
    async def test_interactive_text_still_persisted(self) -> None:
        runner = DialogueRunner()
        state = _make_state()
        tool_msg = _make_tool_msg(
            content_type=ContentType.INTERACTIVE, data_text="Pick a room type"
        )
        tool_response = _make_tool_response(tool_msg)
        message_store = _make_fake_message_store()
        stream_manager = _make_fake_stream_manager()

        await _collect_chunks(
            runner,
            tool_msg,
            tool_response,
            state,
            _make_context(),
            message_store,
            stream_manager,
        )

        # Text message persisted with is_client_message_override=False
        message_store.persist_messages.assert_awaited()
        call_kwargs = message_store.persist_messages.await_args.kwargs
        assert call_kwargs["is_client_message_override"] is False

        # The text message is in history
        text_msgs = [
            m
            for m in state.message_history
            if isinstance(m, SwarmAssistantMessage) and m.content
        ]
        assert len(text_msgs) == 1
        assert text_msgs[0].content[0]["text"] == "Pick a room type"

    @pytest.mark.asyncio
    async def test_interactive_fallback_chain_uses_title(self) -> None:
        """When data.text is missing, falls back to data.title."""
        runner = DialogueRunner()
        state = _make_state()
        tool_msg = _make_tool_msg(
            content_type=ContentType.INTERACTIVE,
            data_text=None,
            data_title="Select your style",
        )
        tool_response = _make_tool_response(tool_msg)
        message_store = _make_fake_message_store()
        stream_manager = _make_fake_stream_manager()

        await _collect_chunks(
            runner,
            tool_msg,
            tool_response,
            state,
            _make_context(),
            message_store,
            stream_manager,
        )

        text_msgs = [
            m
            for m in state.message_history
            if isinstance(m, SwarmAssistantMessage) and m.content
        ]
        assert len(text_msgs) == 1
        assert text_msgs[0].content[0]["text"] == "Select your style"

    @pytest.mark.asyncio
    async def test_interactive_empty_text_still_persists_and_warns(
        self, caplog: pytest.LogCaptureFixture
    ) -> None:
        """When all fields are empty, persist with empty text and log warning."""
        runner = DialogueRunner()
        state = _make_state()
        tool_msg = _make_tool_msg(
            content_type=ContentType.INTERACTIVE,
            data_text=None,
            data_title=None,
            top_title=None,
        )
        tool_response = _make_tool_response(tool_msg)
        message_store = _make_fake_message_store()
        stream_manager = _make_fake_stream_manager()

        with caplog.at_level("WARNING", logger="swarm.util"):
            await _collect_chunks(
                runner,
                tool_msg,
                tool_response,
                state,
                _make_context(),
                message_store,
                stream_manager,
            )

        # Still persisted (gate is `if is_interactive or display_text`)
        message_store.persist_messages.assert_awaited()
        text_msgs = [
            m for m in state.message_history if isinstance(m, SwarmAssistantMessage)
        ]
        assert len(text_msgs) == 1
        assert text_msgs[0].content[0]["text"] == ""

        # Warning logged when all fields are empty
        assert any("interactive_display.no_text" in r.message for r in caplog.records)


class TestHandleDisplayToolRich:
    """RICH display tools: synthetic text chunks still emitted (unchanged)."""

    @pytest.mark.asyncio
    async def test_rich_streams_text_chunks(self) -> None:
        runner = DialogueRunner()
        state = _make_state()
        tool_msg = _make_tool_msg(
            content_type=ContentType.RICH,
            data_text="Here is your plan",
        )
        # Fake actions so it classifies as display
        tool_response = _make_tool_response(tool_msg)
        message_store = _make_fake_message_store()
        stream_manager = _make_fake_stream_manager()

        chunks = await _collect_chunks(
            runner,
            tool_msg,
            tool_response,
            state,
            _make_context(),
            message_store,
            stream_manager,
        )

        # Should have DelimiterChunk(start), AssistantMessageChunks, DelimiterChunk(end)
        delimiters = [c for c in chunks if isinstance(c, DelimiterChunk)]
        assistant_chunks = [c for c in chunks if isinstance(c, AssistantMessageChunk)]

        assert any(d.delim == "start" for d in delimiters)
        assert any(d.delim == "end" for d in delimiters)
        assert len(assistant_chunks) > 0

        # Reconstructed text matches
        full_text = "".join(c.content for c in assistant_chunks if c.content)
        assert full_text == "Here is your plan"

    @pytest.mark.asyncio
    async def test_rich_no_override_on_persist(self) -> None:
        runner = DialogueRunner()
        state = _make_state()
        tool_msg = _make_tool_msg(
            content_type=ContentType.RICH,
            data_text="Here is your plan",
        )
        tool_response = _make_tool_response(tool_msg)
        message_store = _make_fake_message_store()
        stream_manager = _make_fake_stream_manager()

        await _collect_chunks(
            runner,
            tool_msg,
            tool_response,
            state,
            _make_context(),
            message_store,
            stream_manager,
        )

        call_kwargs = message_store.persist_messages.await_args.kwargs
        assert call_kwargs["is_client_message_override"] is None


# ── run_streamed regression test infrastructure ────────────────────────────


class _FakeContext:
    """Minimal ContextProtocol implementation for run_streamed tests."""

    def __init__(self) -> None:
        self._context_variables: Dict[str, Any] = {}
        self.workflow_context: Dict[str, Any] = {}

    @property
    def conversation_id(self) -> str:
        return "conv-1"

    @property
    def current_context_variables(self) -> Dict[str, Any]:
        return self._context_variables

    def get_history_for_completion(self) -> List[Any]:
        return []

    @property
    def designer_fullname(self) -> str:
        return "Designer A"


class _FakeStreamManager:
    def __init__(self) -> None:
        self.statuses: List[str] = []

    async def publish_status(self, status: str) -> None:
        self.statuses.append(status)

    async def publish_error(self, message: str, code: int = 500) -> None:
        pass

    async def publish_message(self, content: Any, **kwargs: Any) -> None:
        pass


class _FakeMessageStore:
    def __init__(self) -> None:
        self.calls: List[Dict[str, Any]] = []

    async def persist_messages(
        self,
        *,
        messages: Any,
        conversation_id: str,
        context_variables: Dict[str, Any],
        designer_model: str,
        designer_fullname: str,
        user_id: Optional[str] = None,
        audience: Optional[str] = None,
        node_key: Optional[str] = None,
        is_client_message_override: Optional[bool] = None,
    ) -> None:
        self.calls.append(
            {
                "messages": messages,
                "conversation_id": conversation_id,
                "is_client_message_override": is_client_message_override,
            }
        )


def _make_stream_chunk(
    *, content: Optional[str] = None, tool_call_name: Optional[str] = None
) -> Any:
    """Build a fake streaming chunk from the LLM."""
    if tool_call_name:
        delta_json = json.dumps(
            {
                "tool_calls": [
                    {
                        "index": 0,
                        "id": f"call-{uuid4()}",
                        "type": "function",
                        "function": {"name": tool_call_name, "arguments": "{}"},
                    }
                ]
            }
        )
    elif content:
        delta_json = json.dumps({"content": content})
    else:
        delta_json = json.dumps({})

    return SimpleNamespace(
        choices=[SimpleNamespace(delta=SimpleNamespace(json=lambda dj=delta_json: dj))],
        usage=None,
    )


async def _fake_stream(*chunks: Any) -> AsyncIterator[Any]:
    for c in chunks:
        yield c


class _FakeStreamingRouter:
    def __init__(self, stream_fn: Callable[[Any], AsyncIterator[Any]]) -> None:
        self._stream_fn = stream_fn

    async def create_completion_async(self, request: Any) -> Any:
        _ = request
        raise AssertionError(
            "create_completion_async should not be called in DialogueRunner tests"
        )

    def stream_completion(self, request: Any) -> AsyncIterator[Any]:
        return self._stream_fn(request)


def _make_run_context(
    *,
    message_store: Optional[_FakeMessageStore] = None,
    stream_manager: Optional[_FakeStreamManager] = None,
    router: ModelRouterProtocol,
) -> RunContext[_FakeContext]:
    return RunContext(
        deps=_FakeContext(),
        stream_manager=stream_manager or _FakeStreamManager(),
        message_store=message_store or _FakeMessageStore(),
        router=router,
        retry_policy=None,
    )


class TestRunStreamedAutoContineRegression:
    """Regression: save tool → agent switch → auto-continue → new agent calls tool."""

    @pytest.mark.asyncio
    @patch("swarm.runners.dialogue.execute_tools")
    async def test_save_tool_agent_switch_auto_continue(
        self, mock_execute_tools: AsyncMock
    ) -> None:
        """Save tool triggers agent switch; new agent gets text-only turn then
        auto-continues and calls a tool."""
        agent_a = Agent(name="AgentA", model="gpt-4o", instructions="A", tools=[])
        agent_b = Agent(name="AgentB", model="gpt-4o", instructions="B", tools=[])

        # Turn 1: Agent A calls save_tool
        # Turn 2 (auto-continued): Agent B returns text only → triggers auto-continue
        # Turn 3: Agent B calls a display tool → breaks loop
        call_count = 0

        save_tool_msg = SwarmToolMessage(
            id=str(uuid4()),
            content=[{"type": "text", "text": "saved"}],
            tool_call_id="call-save",
            tool_name="save_discovery",
            content_type=ContentType.TEXT,
        )

        interactive_tool_msg = _make_tool_msg(
            content_type=ContentType.INTERACTIVE,
            data_text="Pick a room type",
        )

        def execute_side_effect(*args: Any, **kwargs: Any) -> ToolExecutionResponse:
            nonlocal call_count
            call_count += 1
            if call_count == 1:
                # First tool call: save tool with agent switch
                return ToolExecutionResponse(
                    messages=[save_tool_msg],
                    context_variables={},
                    agent=agent_b,
                    metadata=ResultMetadata(),
                )
            # Second tool call: interactive display tool
            return ToolExecutionResponse(
                messages=[interactive_tool_msg],
                context_variables={},
                metadata=ResultMetadata(),
            )

        mock_execute_tools.side_effect = execute_side_effect

        stream_call_count = 0

        def stream_side_effect(request: Any) -> AsyncIterator[Any]:
            nonlocal stream_call_count
            stream_call_count += 1
            if stream_call_count == 1:
                # Agent A: tool call
                return _fake_stream(_make_stream_chunk(tool_call_name="save_discovery"))
            if stream_call_count == 2:
                # Agent B: text only (triggers auto-continue)
                return _fake_stream(_make_stream_chunk(content="Hello!"))
            # Agent B (after auto-continue): tool call
            return _fake_stream(
                _make_stream_chunk(tool_call_name="room_type_selection")
            )

        fake_router = _FakeStreamingRouter(stream_side_effect)

        message_store = _FakeMessageStore()
        run_context = _make_run_context(message_store=message_store, router=fake_router)

        runner = DialogueRunner()
        chunks: List[ChunkType] = []
        async for chunk in runner.run_streamed(agent_a, run_context, max_turns=25):
            chunks.append(chunk)

        # Should have: 3 LLM calls (AgentA tool, AgentB text, AgentB tool)
        assert stream_call_count == 3

        # Should have: 2 tool executions
        assert call_count == 2

        # Auto-continue: synthetic "continue" user message should be persisted
        persisted_msg_types = []
        for call in message_store.calls:
            for msg in call["messages"]:
                if isinstance(msg, SwarmUserMessage):
                    persisted_msg_types.append("user")
                elif isinstance(msg, SwarmAssistantMessage):
                    persisted_msg_types.append("assistant")
                elif isinstance(msg, SwarmToolMessage):
                    persisted_msg_types.append("tool")

        assert "user" in persisted_msg_types  # synthetic "continue" message

        # Loop should end with complete delimiter
        complete_chunks = [
            c for c in chunks if isinstance(c, DelimiterChunk) and c.delim == "complete"
        ]
        assert len(complete_chunks) == 1


class TestRunStreamedInteractiveOrderingInvariant:
    """Regression: interactive tool → synthetic text message ordering."""

    @pytest.mark.asyncio
    @patch("swarm.runners.dialogue.execute_tools")
    async def test_interactive_persist_order(
        self, mock_execute_tools: AsyncMock
    ) -> None:
        """For interactive display tools, the tool response messages are in history
        BEFORE the synthetic assistant text message (appended in _handle_display_tool)."""
        agent = Agent(name="Agent", model="gpt-4o", instructions="sys", tools=[])

        interactive_tool_msg = _make_tool_msg(
            content_type=ContentType.INTERACTIVE,
            data_text="Pick a room type",
        )

        mock_execute_tools.return_value = ToolExecutionResponse(
            messages=[interactive_tool_msg],
            context_variables={},
            metadata=ResultMetadata(),
        )

        fake_router = _FakeStreamingRouter(
            lambda request: _fake_stream(
                _make_stream_chunk(tool_call_name="room_type_selection")
            )
        )

        message_store = _FakeMessageStore()
        run_context = _make_run_context(message_store=message_store, router=fake_router)

        runner = DialogueRunner()
        chunks: List[ChunkType] = []
        async for chunk in runner.run_streamed(agent, run_context, max_turns=5):
            chunks.append(chunk)

        # Extract message history from the run_context's deps — but we need to
        # check the persisted order via message_store calls instead
        persist_order: List[str] = []
        for call in message_store.calls:
            for msg in call["messages"]:
                if isinstance(msg, SwarmToolMessage):
                    persist_order.append("tool")
                elif isinstance(msg, SwarmAssistantMessage):
                    if msg.tool_calls:
                        persist_order.append("assistant_tool_call")
                    else:
                        persist_order.append("assistant_text")

        # Ordering: assistant_tool_call → tool → assistant_text (synthetic)
        assert "assistant_tool_call" in persist_order
        assert "tool" in persist_order
        assert "assistant_text" in persist_order

        tool_idx = persist_order.index("tool")
        text_idx = persist_order.index("assistant_text")
        assert tool_idx < text_idx, (
            f"Tool response (idx={tool_idx}) should be persisted before "
            f"synthetic text (idx={text_idx})"
        )

        # Interactive text should have is_client_message_override=False
        text_persist_calls = [
            c for c in message_store.calls if c["is_client_message_override"] is False
        ]
        assert len(text_persist_calls) == 1
