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 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 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


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=[],
        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_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_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.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.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)


@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[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 == "passed"
    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 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