agent-lab/tests/test_benchmark_runner.py

import asyncio
from collections.abc import AsyncIterator
from typing import Any

import pytest
from pydantic import ValidationError

from agent_lab.application import benchmark
from agent_lab.application.contracts import AgentParams
from agent_lab.application.events import (
    InMemoryCalendarScheduleAdapter,
    InMemoryDeviceVolumeAdapter,
    ResultPolicy,
)
from agent_lab.application.tools import ToolDefinition, ToolRegistry
from agent_lab.domain.events import ToolCallEvent
from agent_lab.domain.messages import ChatMessage, StreamItem
from agent_lab.domain.messages import TokenUsage
from agent_lab.infrastructure.sqlite_store import SQLiteSessionStore


class ManualClock:
    def __init__(self) -> None:
        self.now = 0.0

    def __call__(self) -> float:
        return self.now


class ScriptedTimingClient:
    def __init__(self, clock: ManualClock, *, fail: bool = False) -> None:
        self.clock = clock
        self.fail = fail
        self.closed = False

    async def stream_chat(
        self,
        messages: list[ChatMessage],
        tools: list[dict[str, Any]],
        params: AgentParams,
        tool_choice: dict[str, Any] | None = None,
    ) -> AsyncIterator[StreamItem]:
        del messages, tools, params, tool_choice
        self.clock.now = 0.011
        yield StreamItem.raw_response_chunk({"provider": "first"})
        if self.fail:
            self.clock.now = 0.023
            raise RuntimeError("stream failed")
        self.clock.now = 0.037
        yield StreamItem.message_delta("visible")
        self.clock.now = 0.041
        yield StreamItem.usage_item(TokenUsage(total_tokens=3))
        self.clock.now = 0.059
        yield StreamItem.usage_item(TokenUsage(total_tokens=5))
        self.clock.now = 0.071

    async def aclose(self) -> None:
        self.closed = True


class WhitespaceTimingClient:
    def __init__(self, clock: ManualClock) -> None:
        self.clock = clock

    async def stream_chat(
        self,
        messages: list[ChatMessage],
        tools: list[dict[str, Any]],
        params: AgentParams,
        tool_choice: dict[str, Any] | None = None,
    ) -> AsyncIterator[StreamItem]:
        del messages, tools, params, tool_choice
        self.clock.now = 0.011
        yield StreamItem.raw_response_chunk({"provider": "first"})
        self.clock.now = 0.019
        yield StreamItem.message_delta("")
        self.clock.now = 0.023
        yield StreamItem.message_delta("   \n")
        self.clock.now = 0.037
        yield StreamItem.message_delta("visible")
        self.clock.now = 0.050
        yield StreamItem.usage_item(TokenUsage(total_tokens=5))

    async def aclose(self) -> None:
        return None


class FallbackBenchmarkClient:
    def __init__(self) -> None:
        self.calls = 0
        self.closed = False

    async def stream_chat(
        self,
        messages: list[ChatMessage],
        tools: list[dict[str, Any]],
        params: AgentParams,
        tool_choice: dict[str, Any] | None = None,
    ) -> AsyncIterator[StreamItem]:
        del messages, tools, params
        self.calls += 1
        if self.calls == 1:
            assert tool_choice is None
            yield StreamItem.raw_response_chunk({"chat": 1})
            yield StreamItem.message_delta("I will set the volume to 30.")
            yield StreamItem.text_event(
                ToolCallEvent(
                    id="volume-1",
                    name="device.volume.adjust",
                    arguments={"mode": "absolute"},
                    raw_arguments='{"mode":"absolute"}',
                )
            )
            yield StreamItem.usage_item(TokenUsage(total_tokens=5))
            return
        if self.calls == 2:
            assert tool_choice is not None
            yield StreamItem.raw_response_chunk({"fallback": 1})
            yield StreamItem.provider_tool_call(
                ToolCallEvent(
                    id="fallback-volume",
                    name="device.volume.adjust",
                    arguments={"mode": "absolute", "value": 30},
                    raw_arguments='{"mode":"absolute","value":30}',
                )
            )
            yield StreamItem.usage_item(TokenUsage(total_tokens=7))
            return
        raise RuntimeError("unexpected fallback client call")

    async def aclose(self) -> None:
        self.closed = True


class RecordingSQLiteSessionStore(SQLiteSessionStore):
    instances: list["RecordingSQLiteSessionStore"] = []

    def __init__(self, database_path: str) -> None:
        super().__init__(database_path)
        self.instances.append(self)


class BlockingTimingClient:
    def __init__(self, clock: ManualClock) -> None:
        self.clock = clock
        self.blocked = asyncio.Event()
        self.closed = False

    async def stream_chat(
        self,
        messages: list[ChatMessage],
        tools: list[dict[str, Any]],
        params: AgentParams,
        tool_choice: dict[str, Any] | None = None,
    ) -> AsyncIterator[StreamItem]:
        del messages, tools, params, tool_choice
        self.clock.now = 0.010
        yield StreamItem.raw_response_chunk({"first": True})
        self.blocked.set()
        await asyncio.Event().wait()

    async def aclose(self) -> None:
        self.closed = True


class NoAnswerClient:
    def __init__(self) -> None:
        self.closed = False

    async def stream_chat(
        self,
        messages: list[ChatMessage],
        tools: list[dict[str, Any]],
        params: AgentParams,
        tool_choice: dict[str, Any] | None = None,
    ) -> AsyncIterator[StreamItem]:
        del messages, tools, params, tool_choice
        yield StreamItem.raw_response_chunk({"no_answer": True})
        yield StreamItem.usage_item(TokenUsage(total_tokens=1))

    async def aclose(self) -> None:
        self.closed = True


class ConcurrentTimingClient:
    def __init__(self) -> None:
        self.calls = 0
        self.first_started = asyncio.Event()
        self.release_first = asyncio.Event()

    async def stream_chat(
        self,
        messages: list[ChatMessage],
        tools: list[dict[str, Any]],
        params: AgentParams,
        tool_choice: dict[str, Any] | None = None,
    ) -> AsyncIterator[StreamItem]:
        del messages, tools, params, tool_choice
        self.calls += 1
        call_number = self.calls
        yield StreamItem.raw_response_chunk({"call": call_number})
        if call_number == 1:
            self.first_started.set()
            await self.release_first.wait()
        yield StreamItem.usage_item(TokenUsage(total_tokens=call_number))

    async def aclose(self) -> None:
        return None


class IncrementalTimingClient:
    def __init__(self, clock: ManualClock) -> None:
        self.clock = clock

    async def stream_chat(
        self,
        messages: list[ChatMessage],
        tools: list[dict[str, Any]],
        params: AgentParams,
        tool_choice: dict[str, Any] | None = None,
    ) -> AsyncIterator[StreamItem]:
        del messages, tools, params, tool_choice
        for item in (
            StreamItem.raw_response_chunk({"provider": "first"}),
            StreamItem.message_delta("visible"),
            StreamItem.usage_item(TokenUsage(total_tokens=5)),
        ):
            self.clock.now += 0.010
            yield item
        self.clock.now += 0.010

    async def aclose(self) -> None:
        return None


class CloseFailingProviderIterator:
    def __init__(
        self,
        *,
        next_error: BaseException | None = None,
        items: tuple[StreamItem, ...] = (),
    ) -> None:
        self.next_error = next_error
        self.items = list(items)

    def __aiter__(self) -> "CloseFailingProviderIterator":
        return self

    async def __anext__(self) -> StreamItem:
        if self.next_error is not None:
            raise self.next_error
        if self.items:
            return self.items.pop(0)
        raise StopAsyncIteration

    async def aclose(self) -> None:
        raise RuntimeError("close failed")


class IteratorTimingClient:
    def __init__(self, iterator: AsyncIterator[StreamItem]) -> None:
        self.iterator = iterator

    def stream_chat(
        self,
        messages: list[ChatMessage],
        tools: list[dict[str, Any]],
        params: AgentParams,
        tool_choice: dict[str, Any] | None = None,
    ) -> AsyncIterator[StreamItem]:
        del messages, tools, params, tool_choice
        return self.iterator


class SelectiveUsageBenchmarkClient:
    def __init__(
        self,
        case: benchmark.BenchmarkCase,
        mode: benchmark.BenchmarkMode,
        *,
        missing_usage_calls: set[int],
    ) -> None:
        self.rounds = list(benchmark.build_mock_rounds(case, mode))
        self.missing_usage_calls = missing_usage_calls
        self.calls = 0

    async def stream_chat(
        self,
        messages: list[ChatMessage],
        tools: list[dict[str, Any]],
        params: AgentParams,
        tool_choice: dict[str, Any] | None = None,
    ) -> AsyncIterator[StreamItem]:
        del messages, tools, params, tool_choice
        self.calls += 1
        yield StreamItem.raw_response_chunk({"call": self.calls})
        for item in self.rounds[self.calls - 1]:
            yield item
        if self.calls not in self.missing_usage_calls:
            yield StreamItem.usage_item(
                TokenUsage(
                    prompt_tokens=10,
                    completion_tokens=5,
                    total_tokens=15,
                    cached_tokens=2,
                )
            )

    async def aclose(self) -> None:
        return None


class NonOverlappingParallelProbe:
    def __init__(self) -> None:
        self.overlapped = False
        self.volume = InMemoryDeviceVolumeAdapter()
        self.calendar = InMemoryCalendarScheduleAdapter()

    async def adjust(
        self,
        event_id: str,
        *,
        mode: str,
        value: int | None = None,
        delta: int | None = None,
    ) -> dict[str, Any]:
        return self.volume.adjust(
            event_id,
            mode=mode,
            value=value,
            delta=delta,
        )

    async def create(
        self,
        event_id: str,
        *,
        title: str,
        start_at: str,
        timezone: str,
        recurrence: str | None = None,
        reminder_minutes: int | None = None,
    ) -> dict[str, Any]:
        return self.calendar.create(
            event_id,
            title=title,
            start_at=start_at,
            timezone=timezone,
            recurrence=recurrence,
            reminder_minutes=reminder_minutes,
        )


def test_benchmark_result_models_are_strict_and_allow_nullable_metrics():
    timing_type = getattr(benchmark, "BenchmarkModelCallTiming")
    result_type = getattr(benchmark, "BenchmarkRunResult")

    timing = timing_type(
        call_index=1,
        call_kind="chat_completion",
        first_item_kind="raw_chunk",
        provider_ttft_ms=11,
        visible_ttft_ms=37,
        elapsed_ms=59,
        usage=TokenUsage(
            prompt_tokens=2,
            completion_tokens=3,
            total_tokens=5,
            cached_tokens=1,
        ),
    )
    result = result_type(
        case_id=benchmark.BenchmarkCaseId.ORDINARY_CHAT,
        mode=benchmark.BenchmarkMode.DUAL_AGENT,
        iteration=1,
        status="failed",
        initial_provider_ttft_ms=None,
        visible_ttft_ms=None,
        turn_wall_time_ms=None,
        prompt_tokens=None,
        completion_tokens=None,
        total_tokens=None,
        cached_tokens=None,
        model_call_count=None,
        fallback_count=None,
        tool_count=None,
        event_names=[],
        batch_event_names=[],
        tool_event_names=[],
        event_sources=[],
        tool_statuses=[],
        tool_latencies_ms=[],
        semantic_failures=["missing visible answer"],
        error=None,
        model_calls=[timing],
    )

    assert result.model_calls == [timing]
    assert result.initial_provider_ttft_ms is None
    with pytest.raises(ValidationError):
        timing_type(
            call_index="1",
            call_kind="chat_completion",
            first_item_kind=None,
            provider_ttft_ms=None,
            visible_ttft_ms=None,
            elapsed_ms=None,
            usage=None,
        )
    with pytest.raises(ValidationError):
        result_type.model_validate(result.model_dump() | {"unexpected": True})


@pytest.mark.asyncio
async def test_timing_client_keeps_provider_visible_and_elapsed_timings_separate():
    clock = ManualClock()
    inner = ScriptedTimingClient(clock)
    client = benchmark.TimingChatClient(inner, clock=clock)

    items = [
        item
        async for item in client.stream_chat(
            messages=[],
            tools=[],
            params=AgentParams(model="benchmark-model"),
        )
    ]

    assert [item.kind for item in items] == [
        "raw_chunk",
        "message_delta",
        "usage",
        "usage",
    ]
    assert client.timings == [
        benchmark.BenchmarkModelCallTiming(
            call_index=1,
            call_kind="chat_completion",
            first_item_kind="raw_chunk",
            provider_ttft_ms=11,
            visible_ttft_ms=37,
            elapsed_ms=71,
            usage=TokenUsage(total_tokens=5),
        )
    ]


@pytest.mark.asyncio
async def test_timing_client_excludes_consumer_backpressure_from_model_elapsed():
    clock = ManualClock()
    client = benchmark.TimingChatClient(
        IncrementalTimingClient(clock),
        clock=clock,
    )

    items = []
    async for item in client.stream_chat(
        messages=[],
        tools=[],
        params=AgentParams(model="benchmark-model"),
    ):
        items.append(item)
        clock.now += 0.100

    assert [item.kind for item in items] == ["raw_chunk", "message_delta", "usage"]
    assert client.timings[0].provider_ttft_ms == 10
    assert client.timings[0].visible_ttft_ms == 120
    assert client.timings[0].elapsed_ms == 40


@pytest.mark.asyncio
async def test_timing_client_ignores_blank_deltas_for_visible_ttft():
    clock = ManualClock()
    client = benchmark.TimingChatClient(
        WhitespaceTimingClient(clock),
        clock=clock,
    )

    async for _ in client.stream_chat(
        messages=[],
        tools=[],
        params=AgentParams(model="benchmark-model"),
    ):
        pass

    assert client.timings[0].provider_ttft_ms == 11
    assert client.timings[0].visible_ttft_ms == 37


@pytest.mark.asyncio
async def test_benchmark_output_queue_ignores_blank_deltas_for_visible_ttft():
    clock = ManualClock()
    queue = benchmark._BenchmarkOutputQueue(clock)

    clock.now = 0.019
    await queue.put({"type": "message_delta", "content": ""})
    clock.now = 0.023
    await queue.put({"type": "message_delta", "content": " \n\t "})
    assert queue.first_message_delta_at is None

    clock.now = 0.037
    await queue.put({"type": "message_delta", "content": "visible"})

    assert queue.first_message_delta_at == 0.037


@pytest.mark.asyncio
async def test_timing_client_records_failed_fallback_and_closes_inner_client():
    clock = ManualClock()
    inner = ScriptedTimingClient(clock, fail=True)
    client = benchmark.TimingChatClient(inner, clock=clock)

    with pytest.raises(RuntimeError, match="stream failed"):
        async for _ in client.stream_chat(
            messages=[],
            tools=[{"type": "function", "function": {"name": "mock"}}],
            params=AgentParams(model="benchmark-model"),
            tool_choice={"type": "function", "function": {"name": "mock"}},
        ):
            pass
    await client.aclose()

    assert inner.closed is True
    assert client.timings == [
        benchmark.BenchmarkModelCallTiming(
            call_index=1,
            call_kind="argument_fallback",
            first_item_kind="raw_chunk",
            provider_ttft_ms=11,
            visible_ttft_ms=None,
            elapsed_ms=23,
            usage=None,
        )
    ]


@pytest.mark.asyncio
async def test_timing_client_finishes_timing_when_stream_is_cancelled():
    clock = ManualClock()
    inner = BlockingTimingClient(clock)
    client = benchmark.TimingChatClient(inner, clock=clock)

    async def consume() -> None:
        async for _ in client.stream_chat(
            messages=[],
            tools=[],
            params=AgentParams(model="benchmark-model"),
        ):
            pass

    task = asyncio.create_task(consume())
    await inner.blocked.wait()
    clock.now = 0.025
    task.cancel()
    await asyncio.gather(task, return_exceptions=True)
    await client.aclose()

    assert client.timings[0].first_item_kind == "raw_chunk"
    assert client.timings[0].provider_ttft_ms == 10
    assert client.timings[0].visible_ttft_ms is None
    assert client.timings[0].elapsed_ms == 25
    assert inner.closed is True


@pytest.mark.asyncio
async def test_timing_client_preserves_provider_cancellation_when_close_fails():
    iterator = CloseFailingProviderIterator(next_error=asyncio.CancelledError())
    client = benchmark.TimingChatClient(IteratorTimingClient(iterator))

    with pytest.raises(asyncio.CancelledError):
        async for _ in client.stream_chat(
            messages=[],
            tools=[],
            params=AgentParams(model="benchmark-model"),
        ):
            pass

    assert len(client.timings) == 1
    assert client.timings[0].first_item_kind is None
    assert client.timings[0].usage is None


@pytest.mark.asyncio
async def test_timing_client_preserves_provider_error_when_close_fails():
    iterator = CloseFailingProviderIterator(
        next_error=RuntimeError("provider failed")
    )
    client = benchmark.TimingChatClient(IteratorTimingClient(iterator))

    with pytest.raises(RuntimeError, match="provider failed"):
        async for _ in client.stream_chat(
            messages=[],
            tools=[],
            params=AgentParams(model="benchmark-model"),
        ):
            pass

    assert len(client.timings) == 1


@pytest.mark.asyncio
async def test_timing_client_propagates_close_error_after_normal_completion():
    iterator = CloseFailingProviderIterator()
    client = benchmark.TimingChatClient(IteratorTimingClient(iterator))

    with pytest.raises(RuntimeError, match="close failed"):
        async for _ in client.stream_chat(
            messages=[],
            tools=[],
            params=AgentParams(model="benchmark-model"),
        ):
            pass

    assert len(client.timings) == 1


@pytest.mark.asyncio
async def test_timing_client_propagates_close_error_when_consumer_closes_stream():
    iterator = CloseFailingProviderIterator(
        items=(StreamItem.raw_response_chunk({"provider": "first"}),)
    )
    client = benchmark.TimingChatClient(IteratorTimingClient(iterator))
    stream = client.stream_chat(
        messages=[],
        tools=[],
        params=AgentParams(model="benchmark-model"),
    )

    assert (await anext(stream)).kind == "raw_chunk"
    with pytest.raises(RuntimeError, match="close failed"):
        await stream.aclose()

    assert len(client.timings) == 1
    assert client.timings[0].first_item_kind == "raw_chunk"


@pytest.mark.asyncio
async def test_timing_client_indexes_concurrent_calls_by_start_order():
    inner = ConcurrentTimingClient()
    client = benchmark.TimingChatClient(inner)

    async def collect() -> None:
        async for _ in client.stream_chat(
            messages=[],
            tools=[],
            params=AgentParams(model="benchmark-model"),
        ):
            pass

    first = asyncio.create_task(collect())
    await inner.first_started.wait()
    await collect()
    inner.release_first.set()
    await first

    assert [timing.call_index for timing in client.timings] == [1, 2]
    assert [timing.usage.total_tokens for timing in client.timings] == [1, 2]


@pytest.mark.asyncio
@pytest.mark.parametrize("mode", benchmark.BENCHMARK_MODES)
@pytest.mark.parametrize("case_id", benchmark.BENCHMARK_CASE_IDS)
async def test_mock_benchmark_client_emits_raw_semantics_and_fixed_usage(
    case_id: benchmark.BenchmarkCaseId,
    mode: benchmark.BenchmarkMode,
):
    case = benchmark.BENCHMARK_CASE_CATALOG[case_id]
    client = benchmark.MockBenchmarkChatClient(case, mode)
    expected_rounds = benchmark.build_mock_rounds(case, mode)

    emitted_rounds = []
    for _ in expected_rounds:
        emitted_rounds.append(
            [
                item
                async for item in client.stream_chat(
                    messages=[],
                    tools=[],
                    params=AgentParams(model="benchmark-model"),
                )
            ]
        )
    await client.aclose()

    assert [items[0].kind for items in emitted_rounds] == [
        "raw_chunk"
    ] * len(expected_rounds)
    assert [items[1:-1] for items in emitted_rounds] == expected_rounds
    assert [items[-1].usage for items in emitted_rounds] == [
        TokenUsage(
            prompt_tokens=10,
            completion_tokens=5,
            total_tokens=15,
            cached_tokens=2,
        )
    ] * len(expected_rounds)

    with pytest.raises(RuntimeError, match=f"mock benchmark stream exhausted: {case_id}"):
        async for _ in client.stream_chat(
            messages=[],
            tools=[],
            params=AgentParams(model="benchmark-model"),
        ):
            pass


@pytest.mark.asyncio
async def test_mock_runner_is_ordered_ignores_factory_and_reports_event_ledgers():
    config = benchmark.BenchmarkConfig(
        schema_version=1,
        base_url="https://provider.example/v1",
        model="benchmark-model",
        runs_per_case=2,
        cases=[
            benchmark.BenchmarkCaseId.WEB_SEARCH_TWO_ANSWERS,
            benchmark.BenchmarkCaseId.ORDINARY_CHAT,
        ],
        modes=[
            benchmark.BenchmarkMode.CHAT_AGENT_TOOLS,
            benchmark.BenchmarkMode.DUAL_AGENT,
        ],
    )
    factory_calls = 0

    def forbidden_factory(config: benchmark.BenchmarkConfig, api_key: str):
        del config, api_key
        nonlocal factory_calls
        factory_calls += 1
        raise AssertionError("mock mode must not call client_factory")

    results = await benchmark.BenchmarkRunner(
        config,
        api_key=None,
        mock=True,
        client_factory=forbidden_factory,
    ).run()

    assert factory_calls == 0
    assert [(item.case_id, item.mode, item.iteration) for item in results] == [
        (case_id, mode, iteration)
        for case_id in config.cases
        for mode in config.modes
        for iteration in range(1, 3)
    ]
    assert {item.status for item in results} == {"passed"}
    search_results = [
        item
        for item in results
        if item.case_id is benchmark.BenchmarkCaseId.WEB_SEARCH_TWO_ANSWERS
    ]
    assert {item.model_call_count for item in search_results} == {2}
    assert {item.fallback_count for item in search_results} == {0}
    assert {item.tool_count for item in search_results} == {1}
    assert {tuple(item.event_names) for item in search_results} == {
        ("knowledge.web.search",)
    }
    assert {tuple(item.batch_event_names) for item in search_results} == {
        ("knowledge.web.search",)
    }
    assert {tuple(item.tool_event_names) for item in search_results} == {
        ("knowledge.web.search",)
    }
    assert {tuple(item.tool_statuses) for item in search_results} == {("success",)}
    assert {
        tuple(item.event_sources) for item in search_results
    } == {("provider_resolved",), ("text_event",)}
    assert all(len(item.tool_latencies_ms) == 1 for item in search_results)
    assert all(
        latency is not None and latency >= 0
        for item in search_results
        for latency in item.tool_latencies_ms
    )
    ordinary_results = [
        item
        for item in results
        if item.case_id is benchmark.BenchmarkCaseId.ORDINARY_CHAT
    ]
    assert all(item.event_names == [] for item in ordinary_results)
    assert all(item.batch_event_names == [] for item in ordinary_results)
    assert all(item.tool_event_names == [] for item in ordinary_results)
    assert all(item.tool_count == 0 for item in ordinary_results)


@pytest.mark.asyncio
async def test_mock_runner_passes_all_catalog_cases_in_both_modes():
    config = benchmark.BenchmarkConfig(
        schema_version=1,
        base_url="https://provider.example/v1",
        model="benchmark-model",
    )

    results = await benchmark.BenchmarkRunner(config, api_key=None, mock=True).run()

    assert len(results) == 12
    assert all(result.status == "passed" for result in results)
    assert all(result.semantic_failures == [] for result in results)
    parallel_results = [
        result
        for result in results
        if result.case_id is benchmark.BenchmarkCaseId.PARALLEL_VOLUME_SCHEDULE
    ]
    assert {result.parallel_events_overlapped for result in parallel_results} == {True}
    assert all(
        result.parallel_events_overlapped is None
        for result in results
        if result.case_id is not benchmark.BenchmarkCaseId.PARALLEL_VOLUME_SCHEDULE
    )


@pytest.mark.asyncio
async def test_runner_fails_parallel_case_when_successful_handlers_do_not_overlap():
    probe = NonOverlappingParallelProbe()
    config = benchmark.BenchmarkConfig(
        schema_version=1,
        base_url="https://provider.example/v1",
        model="benchmark-model",
        cases=[benchmark.BenchmarkCaseId.PARALLEL_VOLUME_SCHEDULE],
        modes=[benchmark.BenchmarkMode.DUAL_AGENT],
    )

    result = (
        await benchmark.BenchmarkRunner(
            config,
            api_key=None,
            mock=True,
            parallel_probe_factory=lambda: probe,
        ).run()
    )[0]

    assert result.event_names == [
        "device.volume.adjust",
        "calendar.schedule.create",
    ]
    assert result.tool_statuses == ["success", "success"]
    assert result.tool_count == 2
    assert result.parallel_events_overlapped is False
    assert result.status == "failed"
    assert "parallel event handlers did not overlap" in result.semantic_failures


@pytest.mark.asyncio
@pytest.mark.parametrize(
    ("case_id", "missing_usage_calls", "expected_missing"),
    [
        (benchmark.BenchmarkCaseId.ORDINARY_CHAT, {1}, 1),
        (benchmark.BenchmarkCaseId.WEB_SEARCH_TWO_ANSWERS, {2}, 1),
    ],
)
async def test_runner_rejects_missing_or_partial_model_usage(
    case_id: benchmark.BenchmarkCaseId,
    missing_usage_calls: set[int],
    expected_missing: int,
):
    case = benchmark.BENCHMARK_CASE_CATALOG[case_id]
    client = SelectiveUsageBenchmarkClient(
        case,
        benchmark.BenchmarkMode.DUAL_AGENT,
        missing_usage_calls=missing_usage_calls,
    )
    config = benchmark.BenchmarkConfig(
        schema_version=1,
        base_url="https://provider.example/v1",
        model="benchmark-model",
        cases=[case_id],
        modes=[benchmark.BenchmarkMode.DUAL_AGENT],
    )

    result = (
        await benchmark.BenchmarkRunner(
            config,
            api_key="test-key",
            client_factory=lambda config, api_key: client,
        ).run()
    )[0]

    assert result.status == "failed"
    assert (
        result.prompt_tokens,
        result.completion_tokens,
        result.total_tokens,
        result.cached_tokens,
    ) == (None, None, None, None)
    assert f"model usage missing for {expected_missing} call(s)" in (
        result.semantic_failures
    )


@pytest.mark.asyncio
async def test_runner_uses_a_fresh_in_memory_store_and_unique_session_per_run(monkeypatch):
    RecordingSQLiteSessionStore.instances = []
    monkeypatch.setattr(benchmark, "SQLiteSessionStore", RecordingSQLiteSessionStore)
    config = benchmark.BenchmarkConfig(
        schema_version=1,
        base_url="https://provider.example/v1",
        model="benchmark-model",
        runs_per_case=2,
        cases=[benchmark.BenchmarkCaseId.ORDINARY_CHAT],
        modes=[benchmark.BenchmarkMode.DUAL_AGENT],
    )

    results = await benchmark.BenchmarkRunner(config, api_key=None, mock=True).run()

    assert len(results) == 2
    assert len(RecordingSQLiteSessionStore.instances) == 2
    sessions = [store.list_sessions() for store in RecordingSQLiteSessionStore.instances]
    assert all(store.database_path == ":memory:" for store in RecordingSQLiteSessionStore.instances)
    assert [len(records) for records in sessions] == [1, 1]
    assert len({records[0]["id"] for records in sessions}) == 2


@pytest.mark.asyncio
async def test_runner_keeps_provider_visible_and_turn_wall_time_distinct():
    clock = ManualClock()
    client = ScriptedTimingClient(clock)
    config = benchmark.BenchmarkConfig(
        schema_version=1,
        base_url="https://provider.example/v1",
        model="benchmark-model",
        cases=[benchmark.BenchmarkCaseId.ORDINARY_CHAT],
        modes=[benchmark.BenchmarkMode.DUAL_AGENT],
    )

    result = (
        await benchmark.BenchmarkRunner(
            config,
            api_key="test-key",
            client_factory=lambda config, api_key: client,
            clock=clock,
        ).run()
    )[0]

    assert result.status == "passed"
    assert result.initial_provider_ttft_ms == 11
    assert result.visible_ttft_ms == 37
    assert result.turn_wall_time_ms == 71
    assert result.model_calls[0].elapsed_ms == 71
    assert client.closed is True


@pytest.mark.asyncio
async def test_runner_continues_after_one_factory_failure():
    config = benchmark.BenchmarkConfig(
        schema_version=1,
        base_url="https://provider.example/v1",
        model="benchmark-model",
        cases=[benchmark.BenchmarkCaseId.ORDINARY_CHAT],
        modes=[
            benchmark.BenchmarkMode.DUAL_AGENT,
            benchmark.BenchmarkMode.CHAT_AGENT_TOOLS,
        ],
    )
    clients: list[ScriptedTimingClient] = []
    calls = 0

    def factory(config: benchmark.BenchmarkConfig, api_key: str):
        del config, api_key
        nonlocal calls
        calls += 1
        if calls == 1:
            raise RuntimeError("first factory failed")
        clock = ManualClock()
        client = ScriptedTimingClient(clock)
        clients.append(client)
        return client

    results = await benchmark.BenchmarkRunner(
        config,
        api_key="test-key",
        client_factory=factory,
    ).run()

    assert [result.status for result in results] == ["failed", "passed"]
    assert results[0].error == "first factory failed"
    assert results[0].model_call_count == 0
    assert results[0].fallback_count == 0
    assert results[0].prompt_tokens is None
    assert results[0].completion_tokens is None
    assert results[0].total_tokens is None
    assert results[0].cached_tokens is None
    assert results[1].error is None
    assert clients[0].closed is True


@pytest.mark.asyncio
async def test_runner_marks_semantic_mismatch_failed_and_continues():
    config = benchmark.BenchmarkConfig(
        schema_version=1,
        base_url="https://provider.example/v1",
        model="benchmark-model",
        cases=[benchmark.BenchmarkCaseId.ORDINARY_CHAT],
        modes=[
            benchmark.BenchmarkMode.DUAL_AGENT,
            benchmark.BenchmarkMode.CHAT_AGENT_TOOLS,
        ],
    )
    failed_client = NoAnswerClient()
    calls = 0

    def factory(config: benchmark.BenchmarkConfig, api_key: str):
        del config, api_key
        nonlocal calls
        calls += 1
        if calls == 1:
            return failed_client
        return benchmark.MockBenchmarkChatClient(
            benchmark.BENCHMARK_CASE_CATALOG[benchmark.BenchmarkCaseId.ORDINARY_CHAT],
            benchmark.BenchmarkMode.CHAT_AGENT_TOOLS,
        )

    results = await benchmark.BenchmarkRunner(
        config,
        api_key="test-key",
        client_factory=factory,
    ).run()

    assert [result.status for result in results] == ["failed", "passed"]
    assert results[0].error is None
    assert results[0].semantic_failures == ["answer_count expected 1, got 0"]
    assert failed_client.closed is True


@pytest.mark.asyncio
async def test_runner_reports_argument_fallback_model_and_tool_ledgers(monkeypatch):
    client = FallbackBenchmarkClient()
    registry = ToolRegistry(
        [
            ToolDefinition(
                name="device.volume.adjust",
                description="Force fallback for benchmark ledger coverage.",
                parameters={
                    "type": "object",
                    "properties": {
                        "mode": {"type": "string"},
                        "value": {"type": "integer"},
                    },
                    "required": ["mode", "value"],
                },
                handler=lambda event: {
                    "tool": event.name,
                    "status": "payload-status-is-not-kernel-status",
                },
                argument_resolver=lambda event, context: {},
                result_policy=ResultPolicy.SILENT_SUCCESS,
            )
        ]
    )
    monkeypatch.setattr(benchmark, "build_default_tool_registry", lambda: registry)
    config = benchmark.BenchmarkConfig(
        schema_version=1,
        base_url="https://provider.example/v1",
        model="benchmark-model",
        cases=[benchmark.BenchmarkCaseId.DEVICE_VOLUME_SILENT],
        modes=[benchmark.BenchmarkMode.DUAL_AGENT],
    )

    result = (
        await benchmark.BenchmarkRunner(
            config,
            api_key="test-key",
            client_factory=lambda config, api_key: client,
        ).run()
    )[0]

    assert result.status == "failed"
    assert result.model_call_count == 1
    assert result.fallback_count == 1
    assert result.tool_count == 1
    assert result.total_tokens == 12
    assert [timing.call_kind for timing in result.model_calls] == [
        "chat_completion",
        "argument_fallback",
    ]
    assert result.event_names == ["device.volume.adjust"]
    assert result.event_sources == ["text_event"]
    assert result.tool_statuses == ["success"]
    assert result.semantic_failures == [
        "fallback_count expected 0, ledger got 1",
        "fallback_count expected 0, timing got 1",
    ]
    assert client.calls == 2
    assert client.closed is True


def test_live_runner_requires_an_api_key():
    config = benchmark.BenchmarkConfig(
        schema_version=1,
        base_url="https://provider.example/v1",
        model="benchmark-model",
        cases=[benchmark.BenchmarkCaseId.ORDINARY_CHAT],
        modes=[benchmark.BenchmarkMode.DUAL_AGENT],
    )

    with pytest.raises(ValueError, match="api_key"):
        benchmark.BenchmarkRunner(config, api_key=None)


def test_visible_answer_order_ignores_empty_deltas_before_a_tool_result():
    answers, first_answer_index, first_tool_index = (
        benchmark.BenchmarkRunner._visible_answers(
            [
                {"type": "message_delta", "content": "  "},
                {"type": "tool_result", "message": {}},
                {"type": "message_delta", "content": "late answer"},
            ]
        )
    )

    assert answers == ["late answer"]
    assert first_answer_index == 2
    assert first_tool_index == 1


def _semantic_failures_for_counts(
    *,
    ledger_model_count: int,
    timing_model_count: int,
    ledger_fallback_count: int = 0,
    timing_fallback_count: int = 0,
) -> list[str]:
    config = benchmark.BenchmarkConfig(
        schema_version=1,
        base_url="https://provider.example/v1",
        model="benchmark-model",
        cases=[benchmark.BenchmarkCaseId.ORDINARY_CHAT],
        modes=[benchmark.BenchmarkMode.DUAL_AGENT],
    )
    runner = benchmark.BenchmarkRunner(config, api_key=None, mock=True)
    timings = [
        benchmark.BenchmarkModelCallTiming(
            call_index=index,
            call_kind=call_kind,
            first_item_kind="raw_chunk",
            provider_ttft_ms=0,
            visible_ttft_ms=0 if call_kind == "chat_completion" else None,
            elapsed_ms=0,
            usage=TokenUsage(),
        )
        for index, call_kind in enumerate(
            ["chat_completion"] * timing_model_count
            + ["argument_fallback"] * timing_fallback_count,
            start=1,
        )
    ]
    return runner._semantic_failures(
        case=benchmark.BENCHMARK_CASE_CATALOG[
            benchmark.BenchmarkCaseId.ORDINARY_CHAT
        ],
        mode=benchmark.BenchmarkMode.DUAL_AGENT,
        outputs=[{"type": "message_delta", "content": "Ordinary answer."}],
        audits=[],
        event_names=[],
        batch_event_names=[],
        tool_event_names=[],
        event_sources=[],
        tool_statuses=[],
        model_call_count=ledger_model_count,
        fallback_count=ledger_fallback_count,
        tool_count=0,
        timings=timings,
        check_mock_text=True,
    )


@pytest.mark.parametrize("actual_count", [0, 2])
def test_semantics_reject_missing_or_extra_model_calls_even_when_ledgers_agree(
    actual_count: int,
):
    failures = _semantic_failures_for_counts(
        ledger_model_count=actual_count,
        timing_model_count=actual_count,
    )

    assert failures == [
        f"model_call_count expected 1, ledger got {actual_count}",
        f"model_call_count expected 1, timing got {actual_count}",
    ]


def test_semantics_reject_extra_fallback_calls_even_when_ledgers_agree():
    failures = _semantic_failures_for_counts(
        ledger_model_count=1,
        timing_model_count=1,
        ledger_fallback_count=1,
        timing_fallback_count=1,
    )

    assert failures == [
        "fallback_count expected 0, ledger got 1",
        "fallback_count expected 0, timing got 1",
    ]


def test_semantics_keep_timing_vs_ledger_consistency_failure():
    failures = _semantic_failures_for_counts(
        ledger_model_count=1,
        timing_model_count=2,
    )

    assert failures == [
        "model_call_count expected 1, timing got 2",
        "model_call_count ledger=1, timing=2",
    ]


def test_build_result_rejects_wrong_event_identity_in_all_ledgers_and_pairs_status():
    config = benchmark.BenchmarkConfig(
        schema_version=1,
        base_url="https://provider.example/v1",
        model="benchmark-model",
        cases=[benchmark.BenchmarkCaseId.DEVICE_VOLUME_SILENT],
        modes=[benchmark.BenchmarkMode.DUAL_AGENT],
    )
    runner = benchmark.BenchmarkRunner(config, api_key=None, mock=True)
    timing = benchmark.BenchmarkModelCallTiming(
        call_index=1,
        call_kind="chat_completion",
        first_item_kind="raw_chunk",
        provider_ttft_ms=1,
        visible_ttft_ms=2,
        elapsed_ms=3,
        usage=TokenUsage(total_tokens=5),
    )

    result = runner._build_result(
        case=benchmark.BENCHMARK_CASE_CATALOG[
            benchmark.BenchmarkCaseId.DEVICE_VOLUME_SILENT
        ],
        mode=benchmark.BenchmarkMode.DUAL_AGENT,
        iteration=1,
        outputs=[
            {
                "type": "message_delta",
                "content": "I will set the volume to 30.",
            },
            {"type": "tool_result", "message": {}},
        ],
        audits=[
            {
                "event": "chat_event_detected",
                "details": {
                    "event_name": "wrong.detected",
                    "event_source": "text_event",
                },
            },
            {
                "event": "event_batch_results",
                "details": {
                    "results": [
                        {"event_name": "wrong.batch", "status": "success"}
                    ]
                },
            },
        ],
        usage={
            "calls": [
                {
                    "call_kind": "chat_completion",
                    "event_name": None,
                    "tool_latency_ms": None,
                },
                {
                    "call_kind": "tool_execution",
                    "event_name": "wrong.tool",
                    "tool_latency_ms": 4,
                },
            ],
            "session": {
                "prompt_tokens": 2,
                "completion_tokens": 3,
                "total_tokens": 5,
                "cached_tokens": 0,
                "fallback_count": 0,
                "tool_count": 1,
                "turn_wall_time_ms": 10,
            },
        },
        timings=[timing],
        visible_ttft_ms=2,
        runtime_error=None,
    )

    assert result.status == "failed"
    assert result.event_names == ["wrong.detected"]
    assert result.batch_event_names == ["wrong.batch"]
    assert result.tool_event_names == ["wrong.tool"]
    assert result.tool_statuses == ["success"]
    assert result.semantic_failures == [
        "event names expected ('device.volume.adjust',), got ('wrong.detected',)",
        "batch event names expected ('device.volume.adjust',), got ('wrong.batch',)",
        "tool event names expected ('device.volume.adjust',), got ('wrong.tool',)",
        "event batch statuses expected "
        "(('device.volume.adjust', 'success'),), got (('wrong.batch', 'success'),)",
    ]


def test_build_result_reports_timing_attempt_counts_and_flags_ledger_mismatch():
    config = benchmark.BenchmarkConfig(
        schema_version=1,
        base_url="https://provider.example/v1",
        model="benchmark-model",
        cases=[benchmark.BenchmarkCaseId.ORDINARY_CHAT],
        modes=[benchmark.BenchmarkMode.DUAL_AGENT],
    )
    runner = benchmark.BenchmarkRunner(config, api_key=None, mock=True)
    timing = benchmark.BenchmarkModelCallTiming(
        call_index=1,
        call_kind="chat_completion",
        first_item_kind="raw_chunk",
        provider_ttft_ms=1,
        visible_ttft_ms=2,
        elapsed_ms=3,
        usage=TokenUsage(total_tokens=5),
    )

    result = runner._build_result(
        case=benchmark.BENCHMARK_CASE_CATALOG[
            benchmark.BenchmarkCaseId.ORDINARY_CHAT
        ],
        mode=benchmark.BenchmarkMode.DUAL_AGENT,
        iteration=1,
        outputs=[{"type": "message_delta", "content": "Ordinary answer."}],
        audits=[],
        usage={
            "calls": [],
            "session": {
                "prompt_tokens": 0,
                "completion_tokens": 0,
                "total_tokens": 0,
                "cached_tokens": 0,
                "fallback_count": 0,
                "tool_count": 0,
                "turn_wall_time_ms": 3,
            },
        },
        timings=[timing],
        visible_ttft_ms=2,
        runtime_error=None,
    )

    assert result.status == "failed"
    assert result.model_call_count == 1
    assert result.fallback_count == 0
    assert result.semantic_failures == [
        "model_call_count expected 1, ledger got 0",
        "model_call_count ledger=0, timing=1",
    ]


def test_failed_result_uses_timing_attempt_counts_and_returned_usage():
    timings = [
        benchmark.BenchmarkModelCallTiming(
            call_index=1,
            call_kind="chat_completion",
            first_item_kind="raw_chunk",
            provider_ttft_ms=1,
            visible_ttft_ms=2,
            elapsed_ms=3,
            usage=TokenUsage(
                prompt_tokens=2,
                completion_tokens=3,
                total_tokens=5,
                cached_tokens=1,
            ),
        ),
        benchmark.BenchmarkModelCallTiming(
            call_index=2,
            call_kind="argument_fallback",
            first_item_kind="raw_chunk",
            provider_ttft_ms=4,
            visible_ttft_ms=None,
            elapsed_ms=6,
            usage=TokenUsage(
                prompt_tokens=7,
                completion_tokens=11,
                total_tokens=18,
                cached_tokens=2,
            ),
        ),
    ]

    result = benchmark.BenchmarkRunner._failed_result(
        benchmark.BenchmarkCaseId.ORDINARY_CHAT,
        benchmark.BenchmarkMode.DUAL_AGENT,
        1,
        error="failed before persistence",
        timings=timings,
    )

    assert result.model_call_count == 1
    assert result.fallback_count == 1
    assert (
        result.prompt_tokens,
        result.completion_tokens,
        result.total_tokens,
        result.cached_tokens,
    ) == (9, 14, 23, 3)
    assert result.batch_event_names == []
    assert result.tool_event_names == []


def test_failed_result_keeps_tokens_nullable_when_no_timing_usage_returned():
    timing = benchmark.BenchmarkModelCallTiming(
        call_index=1,
        call_kind="chat_completion",
        first_item_kind="raw_chunk",
        provider_ttft_ms=1,
        visible_ttft_ms=None,
        elapsed_ms=2,
        usage=None,
    )

    result = benchmark.BenchmarkRunner._failed_result(
        benchmark.BenchmarkCaseId.ORDINARY_CHAT,
        benchmark.BenchmarkMode.DUAL_AGENT,
        1,
        error="failed without usage",
        timings=[timing],
    )

    assert result.model_call_count == 1
    assert result.fallback_count == 0
    assert (
        result.prompt_tokens,
        result.completion_tokens,
        result.total_tokens,
        result.cached_tokens,
    ) == (None, None, None, None)
    assert result.semantic_failures == ["model usage missing for 1 call(s)"]


def test_failed_result_does_not_sum_partial_timing_usage():
    timings = [
        benchmark.BenchmarkModelCallTiming(
            call_index=1,
            call_kind="chat_completion",
            first_item_kind="raw_chunk",
            provider_ttft_ms=1,
            visible_ttft_ms=2,
            elapsed_ms=3,
            usage=TokenUsage(
                prompt_tokens=2,
                completion_tokens=3,
                total_tokens=5,
                cached_tokens=1,
            ),
        ),
        benchmark.BenchmarkModelCallTiming(
            call_index=2,
            call_kind="argument_fallback",
            first_item_kind="raw_chunk",
            provider_ttft_ms=4,
            visible_ttft_ms=None,
            elapsed_ms=6,
            usage=None,
        ),
    ]

    result = benchmark.BenchmarkRunner._failed_result(
        benchmark.BenchmarkCaseId.ORDINARY_CHAT,
        benchmark.BenchmarkMode.DUAL_AGENT,
        1,
        error="failed with partial usage",
        timings=timings,
    )

    assert (
        result.prompt_tokens,
        result.completion_tokens,
        result.total_tokens,
        result.cached_tokens,
    ) == (None, None, None, None)
    assert result.semantic_failures == ["model usage missing for 1 call(s)"]