agent-lab/tests/test_debug_runtime.py

import asyncio
import importlib
import json
import logging
from collections.abc import AsyncIterator
from typing import Any

import pytest

from agent_lab.application.contracts import AgentParams, DebugRunRequest, EventAgentParams
from agent_lab.application.event_agent import EventAgentRequest
from agent_lab.application.runtime import DebugRuntime
from agent_lab.application.tools import ToolDefinition, ToolExecutionContext, ToolRegistry
from agent_lab.domain.events import ToolCallEvent
from agent_lab.domain.messages import ChatMessage, StreamItem, TokenUsage
from agent_lab.infrastructure.sqlite_store import SQLiteSessionStore


def _runtime_queues_class():
    module = importlib.import_module("agent_lab.application.queues")
    return module.RuntimeQueues


async def _collect_outputs(stream: AsyncIterator[dict[str, Any]]) -> list[dict[str, Any]]:
    return [message async for message in stream]


def _without_audit(outputs: list[dict[str, Any]]) -> list[dict[str, Any]]:
    return [message for message in outputs if message["type"] != "audit"]


def _message_types(outputs: list[dict[str, Any]]) -> list[str]:
    return [message["type"] for message in _without_audit(outputs)]


def _event_tool_call_from_tools(
    tools: list[dict],
    messages: list[ChatMessage],
) -> StreamItem:
    tool_name = tools[0]["function"]["name"]
    content = ""
    for role in ("assistant", "user"):
        content = next(
            (
                message.content
                for message in reversed(messages)
                if message.role == role and message.content.strip()
            ),
            "",
        )
        if content:
            break
    arguments = {
        "message": content,
        "query": content,
        "title": content,
    }
    return StreamItem.provider_tool_call(
        ToolCallEvent(
            id="event_agent_call_1",
            name=tool_name,
            arguments=arguments,
            raw_arguments=json.dumps(arguments),
        )
    )


async def _next_non_audit(
    stream: AsyncIterator[dict[str, Any]],
) -> dict[str, Any]:
    while True:
        message = await anext(stream)
        if message["type"] != "audit":
            return message


async def _next_non_audit_from_queue(queues: Any) -> dict[str, Any]:
    while True:
        message = await queues.output.get()
        if message["type"] != "audit":
            return message


class RecordingQueue(asyncio.Queue):
    def __init__(self, name: str, log: list[tuple[str, str, str]]) -> None:
        super().__init__()
        self.name = name
        self.log = log

    async def put(self, item: Any) -> None:
        self.log.append((self.name, "put", self._describe(item)))
        await super().put(item)

    async def get(self) -> Any:
        item = await super().get()
        self.log.append((self.name, "get", self._describe(item)))
        return item

    def _describe(self, item: Any) -> str:
        if isinstance(item, ChatMessage):
            if item.role == "tool":
                return f"tool:{item.tool_call_id}"
            return item.role
        if isinstance(item, ToolCallEvent):
            return f"event:{item.name}:{item.id}"
        if isinstance(item, EventAgentRequest):
            events = ",".join(f"{event.name}:{event.id}" for event in item.events)
            return f"event_request:{events}"
        if isinstance(item, dict):
            return f"output:{item.get('type')}"
        return type(item).__name__


class FakeChatClient:
    def __init__(self) -> None:
        self.calls = 0

    async def stream_chat(
        self,
        messages: list[ChatMessage],
        tools: list[dict],
        params: AgentParams,
        tool_choice: dict[str, Any] | None = None,
    ) -> AsyncIterator[StreamItem]:
        if tools:
            yield StreamItem.raw_response_chunk(
                {
                    "choices": [
                        {
                            "delta": {
                                "tool_calls": [
                                    {
                                        "index": 0,
                                        "function": {"name": tools[0]["function"]["name"]},
                                    }
                                ]
                            },
                            "finish_reason": None,
                        }
                    ]
                }
            )
            yield _event_tool_call_from_tools(tools, messages)
            return
        self.calls += 1
        if self.calls == 1:
            yield StreamItem.raw_response_chunk(
                {
                    "choices": [
                        {
                            "delta": {"content": "<agent_events>handoff_note</agent_events>"},
                            "finish_reason": None,
                        }
                    ]
                }
            )
            yield StreamItem.text_event(
                ToolCallEvent(
                    id="call_1",
                    name="handoff_note",
                    arguments={},
                    raw_arguments="{}",
                )
            )
            return

        assert not any(message.role == "tool" for message in messages)
        assert any(message.name == "event_agent" for message in messages)
        yield StreamItem.message_delta("final answer")


class WrongSourceEventChatClient:
    def __init__(self, item: StreamItem) -> None:
        self.item = item
        self.calls = 0

    async def stream_chat(
        self,
        messages: list[ChatMessage],
        tools: list[dict],
        params: AgentParams,
        tool_choice: dict[str, Any] | None = None,
    ) -> AsyncIterator[StreamItem]:
        if tools:
            yield _event_tool_call_from_tools(tools, messages)
            return
        self.calls += 1
        if self.calls == 1:
            yield self.item
            return
        yield StreamItem.message_delta("unexpected continuation")


class IncrementingClock:
    def __init__(self, current: float = 100.0, step: float = 0.01) -> None:
        self.current = current
        self.step = step

    def __call__(self) -> float:
        self.current += self.step
        return self.current


class StrictHistoryChatClient:
    def __init__(self) -> None:
        self.calls = 0
        self.second_call_messages: list[ChatMessage] = []

    async def stream_chat(
        self,
        messages: list[ChatMessage],
        tools: list[dict],
        params: AgentParams,
        tool_choice: dict[str, Any] | None = None,
    ) -> AsyncIterator[StreamItem]:
        if tools:
            yield StreamItem.raw_response_chunk(
                {
                    "choices": [
                        {
                            "delta": {
                                "tool_calls": [
                                    {
                                        "index": 0,
                                        "function": {"name": tools[0]["function"]["name"]},
                                    }
                                ]
                            },
                            "finish_reason": None,
                        }
                    ]
                }
            )
            yield _event_tool_call_from_tools(tools, messages)
            return
        self.calls += 1
        if self.calls == 1:
            yield StreamItem.raw_response_chunk(
                {
                    "choices": [
                        {
                            "delta": {"content": "<agent_events>handoff_note</agent_events>"},
                            "finish_reason": None,
                        }
                    ]
                }
            )
            yield StreamItem.text_event(
                ToolCallEvent(
                    id="call_1",
                    name="handoff_note",
                    arguments={},
                    raw_arguments="{}",
                )
            )
            return

        self.second_call_messages = list(messages)
        yield StreamItem.message_delta("final answer")


class MultiEventChatClient:
    def __init__(self) -> None:
        self.calls = 0
        self.second_call_messages: list[ChatMessage] = []

    async def stream_chat(
        self,
        messages: list[ChatMessage],
        tools: list[dict],
        params: AgentParams,
        tool_choice: dict[str, Any] | None = None,
    ) -> AsyncIterator[StreamItem]:
        if tools:
            yield _event_tool_call_from_tools(tools, messages)
            return
        self.calls += 1
        if self.calls == 1:
            yield StreamItem.message_delta("Checking events.")
            yield StreamItem.text_event(
                ToolCallEvent(
                    id="call_1",
                    name="handoff_note",
                    arguments={"message": "ignored chat argument"},
                    raw_arguments='{"message":"ignored chat argument"}',
                )
            )
            yield StreamItem.text_event(
                ToolCallEvent(
                    id="call_2",
                    name="audit_note",
                    arguments={"message": "ignored chat argument"},
                    raw_arguments='{"message":"ignored chat argument"}',
                )
            )
            return

        self.second_call_messages = list(messages)
        yield StreamItem.message_delta("Final answer.")


class ToolCapturingChatClient:
    def __init__(self) -> None:
        self.tools: list[dict[str, Any]] = []
        self.messages: list[ChatMessage] = []

    async def stream_chat(
        self,
        messages: list[ChatMessage],
        tools: list[dict],
        params: AgentParams,
        tool_choice: dict[str, Any] | None = None,
    ) -> AsyncIterator[StreamItem]:
        self.messages = list(messages)
        self.tools = list(tools)
        yield StreamItem.message_delta("final answer")


class EventLoopLimitChatClient:
    def __init__(self) -> None:
        self.calls = 0
        self.tools_by_call: list[list[dict[str, Any]]] = []
        self.messages_by_call: list[list[ChatMessage]] = []

    async def stream_chat(
        self,
        messages: list[ChatMessage],
        tools: list[dict],
        params: AgentParams,
        tool_choice: dict[str, Any] | None = None,
    ) -> AsyncIterator[StreamItem]:
        if tools:
            yield _event_tool_call_from_tools(tools, messages)
            return
        self.calls += 1
        self.tools_by_call.append(list(tools))
        self.messages_by_call.append(list(messages))
        if self.calls == 1:
            yield StreamItem.text_event(
                ToolCallEvent(
                    id="call_1",
                    name="handoff_note",
                    arguments={},
                    raw_arguments="{}",
                )
            )
            return

        yield StreamItem.message_delta("final after event limit")


class RoundStatsChatClient:
    async def stream_chat(
        self,
        messages: list[ChatMessage],
        tools: list[dict],
        params: AgentParams,
        tool_choice: dict[str, Any] | None = None,
    ) -> AsyncIterator[StreamItem]:
        yield StreamItem.message_delta("hello")
        yield StreamItem.usage_item(
            TokenUsage(
                prompt_tokens=10,
                completion_tokens=20,
                total_tokens=30,
                cached_tokens=5,
            )
        )


class EventRoundStatsChatClient:
    def __init__(self) -> None:
        self.calls = 0

    async def stream_chat(
        self,
        messages: list[ChatMessage],
        tools: list[dict],
        params: AgentParams,
        tool_choice: dict[str, Any] | None = None,
    ) -> AsyncIterator[StreamItem]:
        if tools:
            yield StreamItem.raw_response_chunk(
                {
                    "choices": [
                        {
                            "delta": {
                                "tool_calls": [
                                    {
                                        "index": 0,
                                        "function": {
                                            "name": tools[0]["function"]["name"],
                                        },
                                    }
                                ]
                            },
                            "finish_reason": None,
                        }
                    ]
                }
            )
            yield _event_tool_call_from_tools(tools, messages)
            return
        self.calls += 1
        if self.calls == 1:
            yield StreamItem.raw_response_chunk(
                {
                    "choices": [
                        {
                            "delta": {
                                "content": "<agent_events>handoff_note</agent_events>",
                            },
                            "finish_reason": None,
                        }
                    ]
                }
            )
            yield StreamItem.text_event(
                ToolCallEvent(
                    id="call_1",
                    name="handoff_note",
                    arguments={},
                    raw_arguments="{}",
                )
            )
            yield StreamItem.usage_item(
                TokenUsage(prompt_tokens=3, completion_tokens=0, total_tokens=3)
            )
            return

        yield StreamItem.raw_response_chunk(
            {
                "choices": [
                    {
                        "delta": {"content": "final answer"},
                        "finish_reason": None,
                    }
                ]
            }
        )
        yield StreamItem.message_delta("final answer")
        yield StreamItem.usage_item(
            TokenUsage(prompt_tokens=4, completion_tokens=6, total_tokens=10)
        )


class ContextBoundaryChatClient:
    def __init__(self) -> None:
        self.calls = 0
        self.event_agent_messages: list[list[ChatMessage]] = []

    async def stream_chat(
        self,
        messages: list[ChatMessage],
        tools: list[dict],
        params: AgentParams,
        tool_choice: dict[str, Any] | None = None,
    ) -> AsyncIterator[StreamItem]:
        if tools:
            self.event_agent_messages.append(list(messages))
            yield _event_tool_call_from_tools(tools, messages)
            return
        self.calls += 1
        if self.calls == 1:
            yield StreamItem.message_delta("I will check.")
            yield StreamItem.text_event(
                ToolCallEvent(
                    id="call_1",
                    name="handoff_note",
                    arguments={},
                    raw_arguments="{}",
                )
            )
            return
        yield StreamItem.message_delta("final answer")


class TwoTurnSessionChatClient:
    def __init__(self) -> None:
        self.calls = 0
        self.messages_by_call: list[list[ChatMessage]] = []

    async def stream_chat(
        self,
        messages: list[ChatMessage],
        tools: list[dict],
        params: AgentParams,
        tool_choice: dict[str, Any] | None = None,
    ) -> AsyncIterator[StreamItem]:
        if tools:
            yield _event_tool_call_from_tools(tools, messages)
            return
        self.calls += 1
        self.messages_by_call.append(list(messages))
        if self.calls in {1, 3}:
            yield StreamItem.text_event(
                ToolCallEvent(
                    id=f"call_{self.calls}",
                    name="handoff_note",
                    arguments={},
                    raw_arguments="{}",
                )
            )
            return

        yield StreamItem.message_delta(f"final answer {self.calls}")


class SlowAfterEventChatClient:
    def __init__(self) -> None:
        self.calls = 0
        self.event_seen = asyncio.Event()
        self.release_stream = asyncio.Event()

    async def stream_chat(
        self,
        messages: list[ChatMessage],
        tools: list[dict],
        params: AgentParams,
        tool_choice: dict[str, Any] | None = None,
    ) -> AsyncIterator[StreamItem]:
        if tools:
            yield _event_tool_call_from_tools(tools, messages)
            return
        self.calls += 1
        if self.calls == 1:
            yield StreamItem.message_delta("Need event.")
            yield StreamItem.text_event(
                ToolCallEvent(
                    id="call_1",
                    name="handoff_note",
                    arguments={},
                    raw_arguments="{}",
                )
            )
            self.event_seen.set()
            await self.release_stream.wait()
            return

        yield StreamItem.message_delta("final answer")


def test_runtime_queues_exposes_input_output_and_events_queues():
    RuntimeQueues = _runtime_queues_class()

    queues = RuntimeQueues()

    assert isinstance(queues.input, asyncio.Queue)
    assert isinstance(queues.output, asyncio.Queue)
    assert isinstance(queues.events, asyncio.Queue)
    assert queues.input is not queues.output
    assert queues.input is not queues.events
    assert queues.output is not queues.events


@pytest.mark.asyncio
async def test_runtime_routes_chat_events_through_event_agent_then_continues_chat():
    request = DebugRunRequest(
        user_message="debug this",
        system_prompts=["You are a debugger."],
        pre_messages=[],
        chat_agent=AgentParams(model="fake-model", temperature=0.1, max_tokens=200),
        event_agent=EventAgentParams(enabled_tools=["handoff_note"], max_event_loops=3),
    )
    client = FakeChatClient()
    runtime = DebugRuntime(client)

    outputs = [message async for message in runtime.run(request)]
    business_outputs = _without_audit(outputs)

    assert client.calls == 2
    assert _message_types(outputs) == [
        "session_started",
        "event",
        "tool_result",
        "round_stats",
        "message_delta",
        "round_stats",
        "done",
    ]
    assert business_outputs[1]["event"]["name"] == "handoff_note"
    assert business_outputs[4]["content"] == "final answer"


@pytest.mark.asyncio
@pytest.mark.parametrize(
    "item",
    [
        StreamItem.provider_tool_call(
            ToolCallEvent(
                id="provider_call_1",
                name="handoff_note",
                arguments={"message": "wrong source"},
                raw_arguments='{"message":"wrong source"}',
            )
        ),
        StreamItem.event(
            ToolCallEvent(
                id="legacy_event_1",
                name="handoff_note",
                arguments={},
                raw_arguments="{}",
            )
        ),
    ],
    ids=["provider_tool_call", "legacy_event"],
)
async def test_runtime_ignores_non_text_event_sources(item: StreamItem):
    request = DebugRunRequest(
        user_message="debug this",
        system_prompts=[],
        pre_messages=[],
        chat_agent=AgentParams(model="fake-model"),
        event_agent=EventAgentParams(enabled_tools=["handoff_note"], max_event_loops=2),
    )
    client = WrongSourceEventChatClient(item)

    outputs = [message async for message in DebugRuntime(client).run(request)]

    assert client.calls == 1
    assert "event" not in _message_types(outputs)
    assert "tool_result" not in _message_types(outputs)


@pytest.mark.asyncio
async def test_runtime_emits_audit_events_and_backend_logs(caplog):
    request = DebugRunRequest(
        user_message="debug this",
        system_prompts=[],
        pre_messages=[],
        chat_agent=AgentParams(model="fake-model", temperature=0.1, max_tokens=200),
        event_agent=EventAgentParams(enabled_tools=["handoff_note"], max_event_loops=1),
    )
    runtime = DebugRuntime(FakeChatClient())

    with caplog.at_level(logging.INFO, logger="agent_lab.application.runtime"):
        outputs = [message async for message in runtime.run(request)]

    audit_events = [
        message["event"]
        for message in outputs
        if message["type"] == "audit"
    ]
    assert audit_events == [
        "session_started",
        "chat_round_started",
        "chat_agent_request",
        "chat_event_detected",
        "chat_agent_response",
        "event_agent_request",
        "event_agent_response",
        "event_agent_completed",
        "chat_round_finished",
        "chat_round_started",
        "chat_agent_request",
        "chat_message_stream_started",
        "chat_message_stream_finished",
        "chat_agent_response",
        "chat_round_finished",
        "session_finished",
    ]
    assert "chat_event_detected" in caplog.text
    assert "event_agent_completed" in caplog.text


@pytest.mark.asyncio
async def test_runtime_audits_chat_message_stream_boundaries_in_output_order():
    request = DebugRunRequest(
        user_message="debug this",
        system_prompts=[],
        pre_messages=[],
        chat_agent=AgentParams(model="fake-model", temperature=0.1, max_tokens=200),
        event_agent=EventAgentParams(enabled_tools=[], max_event_loops=1),
    )
    runtime = DebugRuntime(RoundStatsChatClient())

    outputs = [message async for message in runtime.run(request)]
    ordered_labels = [
        message["event"] if message["type"] == "audit" else message["type"]
        for message in outputs
    ]

    assert ordered_labels.index("chat_message_stream_started") < ordered_labels.index(
        "message_delta"
    )
    assert ordered_labels.index("message_delta") < ordered_labels.index(
        "chat_message_stream_finished"
    )
    assert ordered_labels.index("chat_message_stream_finished") < ordered_labels.index(
        "chat_agent_response"
    )

    stream_started = next(
        message
        for message in outputs
        if message.get("event") == "chat_message_stream_started"
    )
    stream_finished = next(
        message
        for message in outputs
        if message.get("event") == "chat_message_stream_finished"
    )
    assert stream_started["details"]["agent"] == "chat_agent"
    assert stream_started["details"]["round_index"] == 1
    assert stream_finished["details"]["delta_count"] == 1
    assert stream_finished["details"]["content_length"] == len("hello")


@pytest.mark.asyncio
async def test_runtime_audit_events_include_turn_relative_elapsed_time():
    request = DebugRunRequest(
        user_message="debug this",
        system_prompts=[],
        pre_messages=[],
        chat_agent=AgentParams(model="fake-model", temperature=0.1, max_tokens=200),
        event_agent=EventAgentParams(enabled_tools=["handoff_note"], max_event_loops=1),
    )
    runtime = DebugRuntime(FakeChatClient(), clock=IncrementingClock())

    queues = runtime.start_session(request)
    outputs: list[dict[str, Any]] = []
    while not any(message["type"] == "turn_completed" for message in outputs):
        outputs.append(await asyncio.wait_for(queues.output.get(), timeout=1))

    await runtime.aclose()

    turn_audits = [
        message
        for message in outputs
        if message["type"] == "audit"
    ]
    elapsed_values = [
        message["details"].get("turn_elapsed_ms") for message in turn_audits
    ]

    assert turn_audits
    assert all(isinstance(value, int) for value in elapsed_values)
    assert all(value >= 0 for value in elapsed_values)
    assert elapsed_values == sorted(elapsed_values)
    assert next(
        message
        for message in turn_audits
        if message["event"] == "event_agent_request"
    )["details"]["turn_elapsed_ms"] >= 0


@pytest.mark.asyncio
async def test_runtime_audit_includes_model_params_prompts_results_and_usage():
    request = DebugRunRequest(
        user_message="debug this",
        system_prompts=["You are a debugger."],
        pre_messages=[],
        chat_agent=AgentParams(model="chat-model", temperature=0.1, max_tokens=200),
        event_agent=EventAgentParams(
            model="event-model",
            temperature=0.4,
            max_tokens=80,
            enabled_tools=["handoff_note"],
            max_event_loops=1,
        ),
    )
    runtime = DebugRuntime(EventRoundStatsChatClient())

    outputs = [message async for message in runtime.run(request)]
    audits = [message for message in outputs if message["type"] == "audit"]

    chat_request = next(
        message for message in audits if message["event"] == "chat_agent_request"
    )
    assert chat_request["details"]["agent"] == "chat_agent"
    assert chat_request["details"]["params"]["model"] == "chat-model"
    assert chat_request["details"]["params"]["temperature"] == 0.1
    assert chat_request["details"]["params"]["max_tokens"] == 200
    assert chat_request["details"]["tools"] == []
    assert chat_request["details"]["messages"][0] == {
        "role": "system",
        "content": "You are a debugger.",
        "name": None,
        "tool_call_id": None,
        "tool_calls": [],
    }

    chat_response = next(
        message for message in audits if message["event"] == "chat_agent_response"
    )
    assert chat_response["details"]["event_names"] == ["handoff_note"]
    assert chat_response["details"]["raw_chunks"] == [
        {
            "choices": [
                {
                    "delta": {"content": "<agent_events>handoff_note</agent_events>"},
                    "finish_reason": None,
                }
            ]
        }
    ]
    assert chat_response["details"]["usage"] == {
        "prompt_tokens": 3,
        "completion_tokens": 0,
        "total_tokens": 3,
        "cached_tokens": 0,
    }

    event_request = next(
        message for message in audits if message["event"] == "event_agent_request"
    )
    assert event_request["details"]["agent"] == "event_agent"
    assert event_request["details"]["params"]["model"] == "event-model"
    assert event_request["details"]["params"]["temperature"] == 0.4
    assert event_request["details"]["events"][0]["name"] == "handoff_note"
    assert event_request["details"]["tools"][0]["function"]["name"] == "handoff_note"
    assert "Authorization" not in str(event_request["details"])

    event_response = next(
        message for message in audits if message["event"] == "event_agent_response"
    )
    assert event_response["details"]["replies"][0]["role"] == "tool"
    assert '"tool": "handoff_note"' in event_response["details"]["replies"][0]["content"]
    assert event_response["details"]["raw_model_chunks"][0]["event_name"] == "handoff_note"
    assert event_response["details"]["raw_model_chunks"][0]["chunks"] == []


@pytest.mark.asyncio
async def test_runtime_round_started_separates_available_events_from_round_budget():
    request = DebugRunRequest(
        user_message="debug this",
        system_prompts=[],
        pre_messages=[],
        chat_agent=AgentParams(model="chat-model", temperature=0.1, max_tokens=200),
        event_agent=EventAgentParams(
            enabled_tools=["handoff_note"],
            max_event_loops=1,
        ),
    )
    runtime = DebugRuntime(EventRoundStatsChatClient())

    outputs = [message async for message in runtime.run(request)]
    round_starts = [
        message
        for message in outputs
        if message.get("event") == "chat_round_started"
    ]

    assert round_starts[0]["details"]["events_enabled"] == ["handoff_note"]
    assert round_starts[0]["details"]["configured_events"] == ["handoff_note"]
    assert round_starts[0]["details"]["event_generation_enabled"] is True
    assert round_starts[0]["details"]["event_prompt_events"] == ["handoff_note"]
    assert round_starts[1]["details"]["events_enabled"] == []
    assert round_starts[1]["details"]["configured_events"] == ["handoff_note"]
    assert round_starts[1]["details"]["event_generation_enabled"] is False
    assert round_starts[1]["details"]["event_prompt_events"] == []


@pytest.mark.asyncio
async def test_runtime_event_agent_history_excludes_chat_agent_system_context():
    request = DebugRunRequest(
        user_message="debug this",
        system_prompts=["ChatAgent root prompt."],
        pre_messages=[
            ChatMessage(role="system", content="ChatAgent pre system."),
            ChatMessage(role="user", content="earlier user"),
            ChatMessage(role="assistant", content="earlier assistant"),
        ],
        chat_agent=AgentParams(model="chat-model", temperature=0.1, max_tokens=200),
        event_agent=EventAgentParams(
            model="event-model",
            enabled_tools=["handoff_note"],
            max_event_loops=1,
        ),
    )
    client = ContextBoundaryChatClient()
    runtime = DebugRuntime(client)

    outputs = [message async for message in runtime.run(request)]
    event_request = next(
        message for message in outputs if message.get("event") == "event_agent_request"
    )

    assert [
        (message["role"], message["content"])
        for message in event_request["details"]["history"]
    ] == [
        ("user", "earlier user"),
        ("assistant", "earlier assistant"),
        ("user", "debug this"),
        ("assistant", "I will check."),
    ]
    assert not any(
        message["role"] == "system"
        for message in event_request["details"]["history"]
    )
    assert client.event_agent_messages == []


@pytest.mark.asyncio
async def test_runtime_session_event_agent_history_excludes_internal_replies():
    request = DebugRunRequest(
        user_message="debug this",
        system_prompts=["ChatAgent root prompt."],
        pre_messages=[ChatMessage(role="assistant", content="prior answer")],
        chat_agent=AgentParams(model="chat-model", temperature=0.1, max_tokens=200),
        event_agent=EventAgentParams(
            model="event-model",
            enabled_tools=["handoff_note"],
            max_event_loops=1,
        ),
    )
    runtime = DebugRuntime(ContextBoundaryChatClient())

    queues = runtime.start_session(request)
    outputs: list[dict[str, Any]] = []
    while not any(message["type"] == "turn_completed" for message in outputs):
        outputs.append(await asyncio.wait_for(queues.output.get(), timeout=1))

    await runtime.aclose()

    event_request = next(
        message for message in outputs if message.get("event") == "event_agent_request"
    )
    assert [
        (message["role"], message["content"], message["name"])
        for message in event_request["details"]["history"]
    ] == [
        ("assistant", "prior answer", None),
        ("user", "debug this", None),
        ("assistant", "I will check.", None),
    ]
    assert not any(
        message["name"] == "event_agent"
        for message in event_request["details"]["history"]
    )


@pytest.mark.asyncio
async def test_runtime_outputs_event_as_soon_as_chat_stream_detects_it():
    RuntimeQueues = _runtime_queues_class()
    queues = RuntimeQueues()
    client = SlowAfterEventChatClient()
    runtime = DebugRuntime(client, queues=queues)
    request = DebugRunRequest(
        user_message="debug this",
        system_prompts=[],
        pre_messages=[],
        chat_agent=AgentParams(model="fake-model", temperature=0.1, max_tokens=200),
        event_agent=EventAgentParams(enabled_tools=["handoff_note"], max_event_loops=1),
    )

    runtime.start(request)

    try:
        assert await _next_non_audit_from_queue(queues) == {"type": "session_started"}
        assert await _next_non_audit_from_queue(queues) == {
            "type": "message_delta",
            "content": "Need event.",
        }
        await asyncio.wait_for(client.event_seen.wait(), timeout=1)
        event_message = await asyncio.wait_for(
            _next_non_audit_from_queue(queues),
            timeout=0.2,
        )

        assert event_message["type"] == "event"
        assert event_message["event"]["name"] == "handoff_note"
    finally:
        client.release_stream.set()
        await runtime.aclose()


@pytest.mark.asyncio
async def test_runtime_batches_round_events_before_continuing_chat_agent():
    def resolve_from_history(
        event: ToolCallEvent,
        context: ToolExecutionContext,
    ) -> dict[str, Any]:
        return {"message": context.history[-1].content, "event": event.name}

    registry = ToolRegistry(
        [
            ToolDefinition(
                name="handoff_note",
                description="Send a handoff note.",
                parameters={
                    "type": "object",
                    "properties": {"message": {"type": "string"}},
                    "required": ["message"],
                },
                handler=lambda event: {
                    "tool": event.name,
                    "message": event.arguments["message"],
                },
                argument_resolver=resolve_from_history,
            ),
            ToolDefinition(
                name="audit_note",
                description="Send an audit note.",
                parameters={
                    "type": "object",
                    "properties": {"message": {"type": "string"}},
                    "required": ["message"],
                },
                handler=lambda event: {
                    "tool": event.name,
                    "message": event.arguments["message"],
                },
                argument_resolver=resolve_from_history,
            ),
        ]
    )
    request = DebugRunRequest(
        user_message="debug this",
        system_prompts=[],
        pre_messages=[],
        chat_agent=AgentParams(model="fake-model", temperature=0.1, max_tokens=200),
        event_agent=EventAgentParams(
            enabled_tools=["handoff_note", "audit_note"],
            max_event_loops=2,
        ),
    )
    client = MultiEventChatClient()
    runtime = DebugRuntime(client, registry=registry)

    outputs = [message async for message in runtime.run(request)]

    assert _message_types(outputs) == [
        "session_started",
        "message_delta",
        "event",
        "event",
        "tool_result",
        "tool_result",
        "round_stats",
        "message_delta",
        "round_stats",
        "done",
    ]
    assert client.calls == 2
    assert [message.role for message in client.second_call_messages] == [
        "system",
        "user",
        "assistant",
        "user",
    ]
    assistant_message = client.second_call_messages[2]
    assert assistant_message.content == "Checking events."
    assert not any(message.role == "tool" for message in client.second_call_messages)
    assert client.second_call_messages[-1].content == (
        "EventAgent results:\n"
        '{"tool": "handoff_note", "message": "Checking events."}\n'
        '{"tool": "audit_note", "message": "Checking events."}'
    )
    assert client.second_call_messages[-1].name == "event_agent"


@pytest.mark.asyncio
async def test_runtime_start_returns_queues_for_downstream_output_consumer():
    request = DebugRunRequest(
        user_message="debug this",
        system_prompts=[],
        pre_messages=[],
        chat_agent=AgentParams(model="fake-model", temperature=0.1, max_tokens=200),
        event_agent=EventAgentParams(enabled_tools=[], max_event_loops=1),
    )
    runtime = DebugRuntime(RoundStatsChatClient())

    queues = runtime.start(request)

    outputs: list[dict[str, Any]] = []
    while True:
        message = await asyncio.wait_for(queues.output.get(), timeout=1)
        outputs.append(message)
        if message["type"] == "done":
            break
    assert _message_types(outputs) == [
        "session_started",
        "message_delta",
        "usage",
        "round_stats",
        "done",
    ]


@pytest.mark.asyncio
async def test_runtime_finalizes_chat_after_reaching_event_loop_limit():
    request = DebugRunRequest(
        user_message="debug this",
        system_prompts=[],
        pre_messages=[],
        chat_agent=AgentParams(model="fake-model", temperature=0.1, max_tokens=200),
        event_agent=EventAgentParams(enabled_tools=["handoff_note"], max_event_loops=1),
    )
    client = EventLoopLimitChatClient()
    runtime = DebugRuntime(client)

    outputs = [message async for message in runtime.run(request)]
    business_outputs = _without_audit(outputs)

    assert client.calls == 2
    assert client.tools_by_call[0] == []
    assert client.tools_by_call[1] == []
    assert "Available events:" in client.messages_by_call[0][0].content
    assert not any(
        "Available events:" in message.content
        for message in client.messages_by_call[1]
        if message.role == "system"
    )
    assert _message_types(outputs) == [
        "session_started",
        "event",
        "tool_result",
        "round_stats",
        "message_delta",
        "round_stats",
        "done",
    ]
    assert business_outputs[4]["content"] == "final after event limit"


@pytest.mark.asyncio
async def test_runtime_buffers_upstream_user_input_until_after_matching_tool_reply():
    RuntimeQueues = _runtime_queues_class()
    queues = RuntimeQueues()
    request = DebugRunRequest(
        user_message="debug this",
        system_prompts=[],
        pre_messages=[],
        chat_agent=AgentParams(model="fake-model", temperature=0.1, max_tokens=200),
        event_agent=EventAgentParams(enabled_tools=["handoff_note"], max_event_loops=2),
    )
    client = StrictHistoryChatClient()
    runtime = DebugRuntime(client, queues=queues)
    stream = runtime.run(request)

    assert await _next_non_audit(stream) == {"type": "session_started"}
    event_message = await _next_non_audit(stream)
    assert event_message["type"] == "event"
    await queues.input.put(ChatMessage(role="user", content="follow-up while tool runs"))
    remaining = [message async for message in stream]

    assert remaining[-1] == {"type": "done"}
    assert [message.role for message in client.second_call_messages] == [
        "system",
        "user",
        "assistant",
        "user",
        "user",
    ]
    assert not any(message.role == "tool" for message in client.second_call_messages)
    assert client.second_call_messages[3].content.startswith("EventAgent results:\n")
    assert client.second_call_messages[3].name == "event_agent"
    assert client.second_call_messages[4].content == "follow-up while tool runs"


@pytest.mark.asyncio
async def test_runtime_continues_when_event_agent_tool_handler_raises():
    def fail_tool(event: ToolCallEvent) -> dict[str, Any]:
        raise RuntimeError("boom")

    registry = ToolRegistry(
        [
            ToolDefinition(
                name="handoff_note",
                description="Broken handoff tool.",
                parameters={"type": "object"},
                handler=fail_tool,
            )
        ]
    )
    request = DebugRunRequest(
        user_message="debug this",
        system_prompts=[],
        pre_messages=[],
        chat_agent=AgentParams(model="fake-model", temperature=0.1, max_tokens=200),
        event_agent=EventAgentParams(enabled_tools=["handoff_note"], max_event_loops=2),
    )
    runtime = DebugRuntime(FakeChatClient(), registry=registry)

    outputs = await asyncio.wait_for(
        _collect_outputs(runtime.run(request)),
        timeout=1,
    )
    business_outputs = _without_audit(outputs)

    assert _message_types(outputs) == [
        "session_started",
        "event",
        "tool_result",
        "round_stats",
        "message_delta",
        "round_stats",
        "done",
    ]
    assert json.loads(business_outputs[2]["message"]["content"]) == {
        "tool": "handoff_note",
        "error": "tool handler failed: boom",
    }


@pytest.mark.asyncio
async def test_runtime_uses_event_and_input_queues_for_event_agent_handoff():
    RuntimeQueues = _runtime_queues_class()
    queue_log: list[tuple[str, str, str]] = []
    queues = RuntimeQueues(
        input=RecordingQueue("input", queue_log),
        output=RecordingQueue("output", queue_log),
        events=RecordingQueue("events", queue_log),
    )
    request = DebugRunRequest(
        user_message="debug this",
        system_prompts=["You are a debugger."],
        pre_messages=[],
        chat_agent=AgentParams(model="fake-model", temperature=0.1, max_tokens=200),
        event_agent=EventAgentParams(enabled_tools=["handoff_note"], max_event_loops=3),
    )
    runtime = DebugRuntime(FakeChatClient(), queues=queues)

    outputs = [message async for message in runtime.run(request)]

    assert _message_types(outputs) == [
        "session_started",
        "event",
        "tool_result",
        "round_stats",
        "message_delta",
        "round_stats",
        "done",
    ]
    assert queue_log.index(("input", "put", "user")) < queue_log.index(
        ("input", "get", "user")
    )
    assert queue_log.index(
        ("events", "put", "event_request:handoff_note:call_1")
    ) < queue_log.index(("events", "get", "event_request:handoff_note:call_1"))
    assert queue_log.index(
        ("events", "get", "event_request:handoff_note:call_1")
    ) < queue_log.index(("input", "put", "tool:call_1"))
    assert queue_log.index(("input", "put", "tool:call_1")) < queue_log.index(
        ("input", "get", "tool:call_1")
    )


@pytest.mark.asyncio
async def test_runtime_run_consumes_output_queue_in_stream_order():
    RuntimeQueues = _runtime_queues_class()
    queue_log: list[tuple[str, str, str]] = []
    queues = RuntimeQueues(
        input=RecordingQueue("input", queue_log),
        output=RecordingQueue("output", queue_log),
        events=RecordingQueue("events", queue_log),
    )
    request = DebugRunRequest(
        user_message="debug this",
        system_prompts=["You are a debugger."],
        pre_messages=[],
        chat_agent=AgentParams(model="fake-model", temperature=0.1, max_tokens=200),
        event_agent=EventAgentParams(enabled_tools=["handoff_note"], max_event_loops=3),
    )
    runtime = DebugRuntime(FakeChatClient(), queues=queues)

    outputs = [message async for message in runtime.run(request)]

    assert _message_types(outputs) == [
        "session_started",
        "event",
        "tool_result",
        "round_stats",
        "message_delta",
        "round_stats",
        "done",
    ]
    output_puts = [
        entry[2] for entry in queue_log if entry[0] == "output" and entry[1] == "put"
    ]
    output_gets = [
        entry[2] for entry in queue_log if entry[0] == "output" and entry[1] == "get"
    ]
    assert [message for message in output_puts if message != "output:audit"] == [
        "output:session_started",
        "output:event",
        "output:tool_result",
        "output:round_stats",
        "output:message_delta",
        "output:round_stats",
        "output:done",
    ]
    assert output_gets == output_puts


@pytest.mark.asyncio
async def test_runtime_continues_with_event_summary_without_tool_call_history():
    request = DebugRunRequest(
        user_message="debug this",
        system_prompts=["You are a debugger."],
        pre_messages=[],
        chat_agent=AgentParams(model="fake-model", temperature=0.1, max_tokens=200),
        event_agent=EventAgentParams(enabled_tools=["handoff_note"], max_event_loops=2),
    )
    client = StrictHistoryChatClient()
    runtime = DebugRuntime(client)

    outputs = [message async for message in runtime.run(request)]

    assert client.calls == 2
    assert [message.role for message in client.second_call_messages] == [
        "system",
        "system",
        "user",
        "assistant",
        "user",
    ]
    assistant_message = client.second_call_messages[3]
    assert assistant_message.content == ""
    assert "handoff_note" in client.second_call_messages[1].content
    assert not any(message.role == "tool" for message in client.second_call_messages)
    assert client.second_call_messages[4].content.startswith("EventAgent results:\n")
    assert client.second_call_messages[4].name == "event_agent"
    assert outputs[-1] == {"type": "done"}


@pytest.mark.asyncio
async def test_runtime_passes_event_catalog_system_message_without_chat_tools():
    registry = ToolRegistry(
        [
            ToolDefinition(
                name="handoff_note",
                description="Registry-owned handoff tool.",
                parameters={
                    "type": "object",
                    "properties": {
                        "message": {"type": "string"},
                        "priority": {"type": "number"},
                    },
                    "required": ["message"],
                },
                handler=lambda event: {"tool": event.name, "message": "handled"},
            )
        ]
    )
    request = DebugRunRequest(
        user_message="debug this",
        system_prompts=[],
        pre_messages=[],
        chat_agent=AgentParams(model="fake-model", temperature=0.1, max_tokens=200),
        event_agent=EventAgentParams(enabled_tools=["handoff_note"], max_event_loops=2),
    )
    client = ToolCapturingChatClient()
    runtime = DebugRuntime(client, registry=registry)

    outputs = [message async for message in runtime.run(request)]

    assert outputs[-1] == {"type": "done"}
    assert client.tools == []
    assert client.messages[0].role == "system"
    assert "Available events:" in client.messages[0].content
    assert "- handoff_note: Registry-owned handoff tool." in client.messages[0].content
    assert "message" not in client.messages[0].content
    assert "priority" not in client.messages[0].content


@pytest.mark.asyncio
async def test_runtime_emits_round_stats_with_clock_and_usage_after_model_turn():
    request = DebugRunRequest(
        user_message="debug this",
        system_prompts=[],
        pre_messages=[],
        chat_agent=AgentParams(model="fake-model", temperature=0.1, max_tokens=200),
        event_agent=EventAgentParams(enabled_tools=[], max_event_loops=1),
    )
    ticks = iter(
        [
            0.9,
            0.91,
            1.0,
            1.01,
            1.02,
            1.123,
            1.2,
            1.25,
            1.3,
            1.456,
            1.7,
            1.8,
        ]
    )
    runtime = DebugRuntime(RoundStatsChatClient(), clock=lambda: next(ticks))

    outputs = [message async for message in runtime.run(request)]
    business_outputs = _without_audit(outputs)

    assert _message_types(outputs) == [
        "session_started",
        "message_delta",
        "usage",
        "round_stats",
        "done",
    ]
    assert business_outputs[3] == {
        "type": "round_stats",
        "round_index": 1,
        "ttft_ms": 123,
        "elapsed_ms": 456,
        "prompt_tokens": 10,
        "completion_tokens": 20,
        "total_tokens": 30,
        "cached_tokens": 5,
        "had_event": False,
    }


@pytest.mark.asyncio
async def test_runtime_emits_round_stats_for_each_chat_call_in_event_handoff():
    request = DebugRunRequest(
        user_message="debug this",
        system_prompts=[],
        pre_messages=[],
        chat_agent=AgentParams(model="fake-model", temperature=0.1, max_tokens=200),
        event_agent=EventAgentParams(enabled_tools=["handoff_note"], max_event_loops=2),
    )
    ticks = iter(
        [
            1.9,
            1.91,
            2.0,
            2.01,
            2.02,
            2.03,
            2.04,
            2.05,
            2.06,
            2.07,
            2.25,
            2.26,
            3.0,
            3.01,
            3.02,
            3.05,
            3.1,
            3.15,
            3.18,
            3.2,
            3.23,
            3.24,
            3.25,
            3.26,
        ]
    )
    client = EventRoundStatsChatClient()
    runtime = DebugRuntime(client, clock=lambda: next(ticks))

    outputs = [message async for message in runtime.run(request)]
    stats = [message for message in outputs if message["type"] == "round_stats"]

    assert client.calls == 2
    assert stats == [
        {
            "type": "round_stats",
            "round_index": 1,
            "ttft_ms": None,
            "elapsed_ms": 250,
            "prompt_tokens": 3,
            "completion_tokens": 0,
            "total_tokens": 3,
            "cached_tokens": 0,
            "had_event": True,
        },
        {
            "type": "round_stats",
            "round_index": 2,
            "ttft_ms": 50,
            "elapsed_ms": 200,
            "prompt_tokens": 4,
            "completion_tokens": 6,
            "total_tokens": 10,
            "cached_tokens": 0,
            "had_event": False,
        },
    ]


@pytest.mark.asyncio
async def test_runtime_session_resets_event_budget_for_each_user_turn():
    request = DebugRunRequest(
        user_message="first turn",
        system_prompts=[],
        pre_messages=[],
        chat_agent=AgentParams(model="fake-model", temperature=0.1, max_tokens=200),
        event_agent=EventAgentParams(enabled_tools=["handoff_note"], max_event_loops=1),
    )
    client = TwoTurnSessionChatClient()
    runtime = DebugRuntime(client)

    queues = runtime.start_session(request)
    outputs: list[dict[str, Any]] = []
    while len([message for message in outputs if message["type"] == "turn_completed"]) < 1:
        outputs.append(await asyncio.wait_for(queues.output.get(), timeout=1))

    await queues.input.put(ChatMessage(role="user", content="second turn"))
    while len([message for message in outputs if message["type"] == "turn_completed"]) < 2:
        outputs.append(await asyncio.wait_for(queues.output.get(), timeout=1))

    await runtime.aclose()

    business_types = [message["type"] for message in _without_audit(outputs)]
    assert business_types.count("turn_started") == 2
    assert business_types.count("turn_completed") == 2
    assert client.calls == 4
    assert "Available events:" in client.messages_by_call[0][0].content
    assert "Available events:" in client.messages_by_call[2][0].content


@pytest.mark.asyncio
async def test_runtime_persists_session_turn_messages_audit_and_usage(tmp_path):
    store = SQLiteSessionStore(tmp_path / "agent_lab.sqlite3")
    request = DebugRunRequest(
        session_id="session-1",
        user_message="persist this",
        system_prompts=["You are a debugger."],
        pre_messages=[],
        chat_agent=AgentParams(model="chat-model", temperature=0.1, max_tokens=200),
        event_agent=EventAgentParams(enabled_tools=[], max_event_loops=1),
    )
    runtime = DebugRuntime(RoundStatsChatClient(), session_store=store)

    queues = runtime.start_session(request)
    outputs: list[dict[str, Any]] = []
    while not any(message["type"] == "turn_completed" for message in outputs):
        outputs.append(await asyncio.wait_for(queues.output.get(), timeout=1))
    await runtime.aclose()

    assert _without_audit(outputs)[0] == {
        "type": "session_started",
        "session_id": "session-1",
    }
    assert store.get_session("session-1")["config"]["chat_agent"]["model"] == "chat-model"
    assert [
        (message["turn_index"], message["role"], message["content"])
        for message in store.list_messages("session-1")
    ] == [
        (1, "user", "persist this"),
        (1, "assistant", "hello"),
    ]
    audit_events = [
        audit["event"]
        for audit in store.list_audit_logs("session-1")
    ]
    assert "session_started" in audit_events
    assert "chat_agent_request" in audit_events
    assert "turn_completed" in audit_events
    usage = store.usage_summary("session-1")
    assert usage["calls"][0]["total_tokens"] == 30
    assert usage["turns"][0]["turn_index"] == 1
    assert usage["session"]["total_tokens"] == 30


@pytest.mark.asyncio
async def test_runtime_continues_persisted_turn_indexes_for_existing_session(tmp_path):
    store = SQLiteSessionStore(tmp_path / "agent_lab.sqlite3")
    session_id = store.create_session(title="existing", config={})
    store.start_turn(session_id, turn_index=1, user_message="old")
    request = DebugRunRequest(
        session_id=session_id,
        user_message="new",
        system_prompts=[],
        pre_messages=[],
        chat_agent=AgentParams(model="chat-model", temperature=0.1, max_tokens=200),
        event_agent=EventAgentParams(enabled_tools=[], max_event_loops=1),
    )
    runtime = DebugRuntime(RoundStatsChatClient(), session_store=store)

    queues = runtime.start_session(request)
    outputs: list[dict[str, Any]] = []
    while not any(message["type"] == "turn_completed" for message in outputs):
        outputs.append(await asyncio.wait_for(queues.output.get(), timeout=1))
    await runtime.aclose()

    assert [
        (message["turn_index"], message["content"])
        for message in store.list_messages(session_id)
    ] == [
        (2, "new"),
        (2, "hello"),
    ]


class ScriptedChatClient:
    def __init__(self, rounds: list[list[StreamItem]]) -> None:
        self.rounds = rounds
        self.calls = 0
        self.messages_by_call: list[list[ChatMessage]] = []
        self.tools_by_call: list[list[dict[str, Any]]] = []

    async def stream_chat(
        self,
        messages: list[ChatMessage],
        tools: list[dict],
        params: AgentParams,
        tool_choice: dict[str, Any] | None = None,
    ) -> AsyncIterator[StreamItem]:
        self.messages_by_call.append(list(messages))
        self.tools_by_call.append(list(tools))
        round_items = self.rounds[self.calls]
        self.calls += 1
        for item in round_items:
            yield item


def _direct_request(
    *,
    enabled_tools: list[str],
    max_event_loops: int = 1,
    session_id: str | None = None,
) -> DebugRunRequest:
    return DebugRunRequest(
        session_id=session_id,
        user_message="debug direct tools",
        system_prompts=["You are a debugger."],
        pre_messages=[],
        chat_agent=AgentParams(model="chat-model"),
        event_agent=EventAgentParams(
            enabled_tools=enabled_tools,
            max_event_loops=max_event_loops,
        ),
        tool_invocation_mode="chat_agent_tools",
    )


@pytest.mark.asyncio
async def test_direct_mode_uses_enabled_schemas_and_valid_ordered_provider_transcript():
    executed: list[tuple[str, dict[str, Any], str]] = []

    def handle(event: ToolCallEvent) -> dict[str, Any]:
        executed.append((event.name, event.arguments, event.raw_arguments))
        return {"tool": event.name, "value": event.arguments["value"]}

    registry = ToolRegistry(
        [
            ToolDefinition(
                name="first_tool",
                description="First direct tool.",
                parameters={
                    "type": "object",
                    "properties": {"value": {"type": "string"}},
                    "required": ["value"],
                },
                handler=handle,
            ),
            ToolDefinition(
                name="second_tool",
                description="Second direct tool.",
                parameters={
                    "type": "object",
                    "properties": {"value": {"type": "string"}},
                    "required": ["value"],
                },
                handler=handle,
            ),
        ]
    )
    calls = [
        ToolCallEvent(
            id="provider-1",
            name="first_tool",
            arguments={"value": "one"},
            raw_arguments='{"value":"one"}',
        ),
        ToolCallEvent(
            id="provider-2",
            name="second_tool",
            arguments={"value": "two"},
            raw_arguments='{"value":"two"}',
        ),
    ]
    ignored_text_event = ToolCallEvent(
        id="text-ignored",
        name="first_tool",
        arguments={},
        raw_arguments="{}",
    )
    client = ScriptedChatClient(
        [
            [
                StreamItem.message_delta("Visible before tools."),
                StreamItem.provider_tool_call(calls[0]),
                StreamItem.text_event(ignored_text_event),
                StreamItem.provider_tool_call(calls[1]),
            ],
            [StreamItem.message_delta("Final answer.")],
        ]
    )

    outputs = await _collect_outputs(
        DebugRuntime(client, registry=registry).run(
            _direct_request(enabled_tools=["first_tool", "second_tool"])
        )
    )

    assert [tool["function"]["name"] for tool in client.tools_by_call[0]] == [
        "first_tool",
        "second_tool",
    ]
    assert client.tools_by_call[1] == []
    assert not any(
        "Available events:" in message.content
        for message in client.messages_by_call[0]
        if message.role == "system"
    )
    transcript = client.messages_by_call[1]
    assert [message.role for message in transcript[-4:]] == [
        "user",
        "assistant",
        "tool",
        "tool",
    ]
    assert transcript[-3].content == "Visible before tools."
    assert transcript[-3].tool_calls == calls
    assert [message.tool_call_id for message in transcript[-2:]] == [
        "provider-1",
        "provider-2",
    ]
    assert [message.name for message in transcript[-2:]] == [
        "first_tool",
        "second_tool",
    ]
    assert executed == [
        ("first_tool", {"value": "one"}, '{"value":"one"}'),
        ("second_tool", {"value": "two"}, '{"value":"two"}'),
    ]
    emitted_events = [
        message["event"] for message in outputs if message["type"] == "event"
    ]
    assert emitted_events == [call.model_dump() for call in calls]
    assert "event_agent_request" not in [
        message.get("event") for message in outputs if message["type"] == "audit"
    ]
    assert not any(
        message.get("event", {}).get("id") == "text-ignored"
        for message in outputs
        if message["type"] == "event"
    )
    assert all(
        message["details"]["tool_invocation_mode"] == "chat_agent_tools"
        for message in outputs
        if message["type"] == "audit"
        and message["event"] in {"chat_round_started", "chat_agent_request"}
    )


@pytest.mark.asyncio
async def test_direct_mode_does_not_execute_unknown_disabled_or_failed_handlers():
    executed: list[str] = []

    def disabled_handler(event: ToolCallEvent) -> dict[str, Any]:
        executed.append(event.name)
        return {"tool": event.name}

    def failing_handler(event: ToolCallEvent) -> dict[str, Any]:
        executed.append(event.name)
        raise RuntimeError("direct boom")

    registry = ToolRegistry(
        [
            ToolDefinition(
                name="enabled_tool",
                description="Enabled.",
                parameters={"type": "object"},
                handler=lambda event: {"tool": event.name},
            ),
            ToolDefinition(
                name="disabled_tool",
                description="Disabled.",
                parameters={"type": "object"},
                handler=disabled_handler,
            ),
            ToolDefinition(
                name="failing_tool",
                description="Fails.",
                parameters={"type": "object"},
                handler=failing_handler,
            ),
        ]
    )
    client = ScriptedChatClient(
        [
            [
                StreamItem.provider_tool_call(
                    ToolCallEvent(
                        id="unknown-1",
                        name="unknown_tool",
                        arguments={"kept": True},
                        raw_arguments='{"kept":true}',
                    )
                ),
                StreamItem.provider_tool_call(
                    ToolCallEvent(
                        id="disabled-1",
                        name="disabled_tool",
                        arguments={},
                        raw_arguments="{}",
                    )
                ),
                StreamItem.provider_tool_call(
                    ToolCallEvent(
                        id="failed-1",
                        name="failing_tool",
                        arguments={},
                        raw_arguments="{}",
                    )
                ),
            ],
            [StreamItem.message_delta("continued")],
        ]
    )

    outputs = await _collect_outputs(
        DebugRuntime(client, registry=registry).run(
            _direct_request(enabled_tools=["enabled_tool", "failing_tool"])
        )
    )

    assert [tool["function"]["name"] for tool in client.tools_by_call[0]] == [
        "enabled_tool",
        "failing_tool",
    ]
    assert executed == ["failing_tool"]
    results = [
        json.loads(message["message"]["content"])
        for message in outputs
        if message["type"] == "tool_result"
    ]
    assert results == [
        {"tool": "unknown_tool", "error": "unknown tool"},
        {"tool": "disabled_tool", "error": "tool disabled"},
        {"tool": "failing_tool", "error": "tool handler failed: direct boom"},
    ]


@pytest.mark.asyncio
async def test_direct_mode_ignores_provider_calls_after_event_budget_exhaustion():
    executed: list[str] = []
    registry = ToolRegistry(
        [
            ToolDefinition(
                name="once_tool",
                description="Run once.",
                parameters={"type": "object"},
                handler=lambda event: executed.append(event.id) or {"tool": event.name},
            )
        ]
    )
    client = ScriptedChatClient(
        [
            [
                StreamItem.provider_tool_call(
                    ToolCallEvent(
                        id="accepted",
                        name="once_tool",
                        arguments={},
                        raw_arguments="{}",
                    )
                )
            ],
            [
                StreamItem.provider_tool_call(
                    ToolCallEvent(
                        id="ignored",
                        name="once_tool",
                        arguments={},
                        raw_arguments="{}",
                    )
                ),
                StreamItem.message_delta("budget exhausted"),
            ],
        ]
    )

    outputs = await _collect_outputs(
        DebugRuntime(client, registry=registry).run(
            _direct_request(enabled_tools=["once_tool"], max_event_loops=1)
        )
    )

    assert client.tools_by_call == [
        [registry.tool_schema("once_tool")],
        [],
    ]
    assert executed == ["accepted"]
    assert [
        message["event"]["id"]
        for message in outputs
        if message["type"] == "event"
    ] == ["accepted"]
    assert outputs[-1] == {"type": "done"}


@pytest.mark.asyncio
async def test_direct_mode_session_turns_reset_budget_and_snapshot_mode(tmp_path):
    store = SQLiteSessionStore(tmp_path / "agent_lab.sqlite3")
    registry = ToolRegistry(
        [
            ToolDefinition(
                name="turn_tool",
                description="Per-turn tool.",
                parameters={"type": "object"},
                handler=lambda event: {"tool": event.name, "call_id": event.id},
            )
        ]
    )
    client = ScriptedChatClient(
        [
            [
                StreamItem.provider_tool_call(
                    ToolCallEvent(
                        id="turn-1-call",
                        name="turn_tool",
                        arguments={},
                        raw_arguments="{}",
                    )
                )
            ],
            [StreamItem.message_delta("turn one done")],
            [
                StreamItem.provider_tool_call(
                    ToolCallEvent(
                        id="turn-2-call",
                        name="turn_tool",
                        arguments={},
                        raw_arguments="{}",
                    )
                )
            ],
            [StreamItem.message_delta("turn two done")],
        ]
    )
    runtime = DebugRuntime(client, registry=registry, session_store=store)
    request = _direct_request(
        enabled_tools=["turn_tool"],
        max_event_loops=1,
        session_id="direct-session",
    )

    queues = runtime.start_session(request)
    outputs: list[dict[str, Any]] = []
    while sum(message["type"] == "turn_completed" for message in outputs) < 1:
        outputs.append(await asyncio.wait_for(queues.output.get(), timeout=1))
    await queues.input.put(ChatMessage(role="user", content="second direct turn"))
    while sum(message["type"] == "turn_completed" for message in outputs) < 2:
        outputs.append(await asyncio.wait_for(queues.output.get(), timeout=1))
    await runtime.aclose()

    assert [bool(tools) for tools in client.tools_by_call] == [True, False, True, False]
    assert [message.role for message in client.messages_by_call[1][-3:]] == [
        "user",
        "assistant",
        "tool",
    ]
    second_turn_transcript = client.messages_by_call[3]
    assert not any(message.name == "event_agent" for message in second_turn_transcript)
    assert [
        message.tool_call_id
        for message in second_turn_transcript
        if message.role == "tool"
    ] == ["turn-1-call", "turn-2-call"]
    session = store.get_session("direct-session")
    assert session is not None
    assert session["config"]["tool_invocation_mode"] == "chat_agent_tools"
    request_audits = [
        audit
        for audit in store.list_audit_logs("direct-session")
        if audit["event"] == "chat_agent_request"
    ]
    assert request_audits
    assert all(
        audit["details"]["tool_invocation_mode"] == "chat_agent_tools"
        for audit in request_audits
    )


def _single_tool_registry(
    handler: Any,
    *,
    name: str = "safe_tool",
) -> ToolRegistry:
    return ToolRegistry(
        [
            ToolDefinition(
                name=name,
                description="Tool used by execution-safety tests.",
                parameters={"type": "object"},
                handler=handler,
            )
        ]
    )


def _tool_call(call_id: str, *, name: str = "safe_tool") -> ToolCallEvent:
    return ToolCallEvent(
        id=call_id,
        name=name,
        arguments={},
        raw_arguments="{}",
    )


@pytest.mark.asyncio
async def test_direct_mode_rejects_provider_call_id_reused_across_one_shot_rounds():
    executed: list[str] = []
    registry = _single_tool_registry(
        lambda event: executed.append(event.id) or {"tool": event.name}
    )
    client = ScriptedChatClient(
        [
            [StreamItem.provider_tool_call(_tool_call("duplicate-call"))],
            [StreamItem.provider_tool_call(_tool_call("duplicate-call"))],
        ]
    )

    outputs = await _collect_outputs(
        DebugRuntime(client, registry=registry).run(
            _direct_request(enabled_tools=["safe_tool"], max_event_loops=2)
        )
    )

    assert executed == ["duplicate-call"]
    assert [
        message["message"]["tool_call_id"]
        for message in outputs
        if message["type"] == "tool_result"
    ] == ["duplicate-call"]
    assert outputs[-1]["type"] == "error"
    assert "duplicate provider tool-call ID: duplicate-call" in outputs[-1]["message"]


@pytest.mark.asyncio
async def test_direct_mode_rejects_provider_call_id_reused_in_session_rounds():
    executed: list[str] = []
    registry = _single_tool_registry(
        lambda event: executed.append(event.id) or {"tool": event.name}
    )
    client = ScriptedChatClient(
        [
            [StreamItem.provider_tool_call(_tool_call("session-duplicate"))],
            [StreamItem.provider_tool_call(_tool_call("session-duplicate"))],
        ]
    )
    runtime = DebugRuntime(client, registry=registry)

    queues = runtime.start_session(
        _direct_request(enabled_tools=["safe_tool"], max_event_loops=2)
    )
    outputs: list[dict[str, Any]] = []
    while not any(message["type"] == "error" for message in outputs):
        outputs.append(await asyncio.wait_for(queues.output.get(), timeout=1))
    await runtime.aclose()

    assert executed == ["session-duplicate"]
    assert [
        message["message"]["tool_call_id"]
        for message in outputs
        if message["type"] == "tool_result"
    ] == ["session-duplicate"]
    assert "duplicate provider tool-call ID: session-duplicate" in outputs[-1][
        "message"
    ]


@pytest.mark.asyncio
async def test_direct_mode_rejects_duplicate_provider_call_ids_within_batch():
    executed: list[str] = []
    registry = _single_tool_registry(
        lambda event: executed.append(event.id) or {"tool": event.name}
    )
    client = ScriptedChatClient(
        [
            [
                StreamItem.provider_tool_call(_tool_call("same-batch")),
                StreamItem.provider_tool_call(_tool_call("same-batch")),
            ]
        ]
    )

    outputs = await _collect_outputs(
        DebugRuntime(client, registry=registry).run(
            _direct_request(enabled_tools=["safe_tool"])
        )
    )

    assert executed == []
    assert outputs[-1]["type"] == "error"
    assert "duplicate provider tool-call ID: same-batch" in outputs[-1]["message"]


@pytest.mark.asyncio
async def test_direct_mode_executes_batch_concurrently_but_replies_in_provider_order():
    both_started = asyncio.Event()
    first_release = asyncio.Event()
    second_release = asyncio.Event()
    second_finished = asyncio.Event()
    started: list[str] = []

    async def handler(event: ToolCallEvent) -> dict[str, Any]:
        started.append(event.id)
        if len(started) == 2:
            both_started.set()
        if event.id == "first-call":
            await first_release.wait()
        else:
            await second_release.wait()
            second_finished.set()
        return {"tool": event.name, "call_id": event.id}

    registry = _single_tool_registry(handler)
    client = ScriptedChatClient(
        [
            [
                StreamItem.provider_tool_call(_tool_call("first-call")),
                StreamItem.provider_tool_call(_tool_call("second-call")),
            ],
            [StreamItem.message_delta("finished")],
        ]
    )
    runtime = DebugRuntime(client, registry=registry)
    run_task = asyncio.create_task(
        _collect_outputs(
            runtime.run(_direct_request(enabled_tools=["safe_tool"]))
        )
    )
    overlapped = False
    try:
        try:
            await asyncio.wait_for(both_started.wait(), timeout=1)
        except TimeoutError:
            pass
        else:
            overlapped = True
            second_release.set()
            await asyncio.wait_for(second_finished.wait(), timeout=1)
            assert not run_task.done()
    finally:
        first_release.set()
        second_release.set()
    outputs = await asyncio.wait_for(run_task, timeout=1)

    assert overlapped is True
    assert started == ["first-call", "second-call"]
    assert [
        message["message"]["tool_call_id"]
        for message in outputs
        if message["type"] == "tool_result"
    ] == ["first-call", "second-call"]


def _dual_budget_request() -> DebugRunRequest:
    return DebugRunRequest(
        user_message="dual budget",
        system_prompts=[],
        pre_messages=[],
        chat_agent=AgentParams(model="chat-model"),
        event_agent=EventAgentParams(
            enabled_tools=["safe_tool"],
            max_event_loops=1,
        ),
    )


@pytest.mark.asyncio
async def test_dual_mode_ignores_text_event_after_one_shot_budget_exhaustion():
    executed: list[str] = []
    registry = _single_tool_registry(
        lambda event: executed.append(event.id) or {"tool": event.name}
    )
    client = ScriptedChatClient(
        [
            [StreamItem.text_event(_tool_call("accepted-text-event"))],
            [
                StreamItem.text_event(_tool_call("ignored-text-event")),
                StreamItem.message_delta("final dual answer"),
            ],
        ]
    )

    outputs = await _collect_outputs(
        DebugRuntime(client, registry=registry).run(_dual_budget_request())
    )

    assert executed == ["accepted-text-event"]
    assert [
        message["event"]["id"]
        for message in outputs
        if message["type"] == "event"
    ] == ["accepted-text-event"]
    assert outputs[-1] == {"type": "done"}


@pytest.mark.asyncio
async def test_dual_mode_session_main_path_ignores_text_event_after_budget_exhaustion():
    executed: list[str] = []
    registry = _single_tool_registry(
        lambda event: executed.append(event.id) or {"tool": event.name}
    )
    client = ScriptedChatClient(
        [
            [StreamItem.text_event(_tool_call("session-accepted"))],
            [
                StreamItem.text_event(_tool_call("session-ignored")),
                StreamItem.message_delta("session final"),
            ],
        ]
    )
    runtime = DebugRuntime(client, registry=registry)

    queues = runtime.start_session(_dual_budget_request())
    outputs: list[dict[str, Any]] = []
    while not any(
        message["type"] in {"turn_completed", "error"} for message in outputs
    ):
        outputs.append(await asyncio.wait_for(queues.output.get(), timeout=1))
    await runtime.aclose()

    assert executed == ["session-accepted"]
    assert [
        message["event"]["id"]
        for message in outputs
        if message["type"] == "event"
    ] == ["session-accepted"]
    assert outputs[-1]["type"] == "turn_completed"