agent-lab/tests/test_event_batch.py

from __future__ import annotations

import asyncio
import threading
from collections.abc import Callable
from typing import Any

import pytest

import agent_lab.application.events as events_module
from agent_lab.application.events import (
    EventDefinition,
    EventKernel,
    EventRegistry,
    EventRequest,
    EventSource,
    EventStatus,
)
from agent_lab.application.tools import ToolDefinition, ToolRegistry
from agent_lab.domain.events import ToolCallEvent


def _executor_type():
    executor_type = getattr(events_module, "EventBatchExecutor", None)
    assert executor_type is not None, "EventBatchExecutor is not implemented"
    return executor_type


def _definition(
    name: str,
    handler: Callable[[EventRequest], Any],
    *,
    conflict_keys: tuple[str, ...] = (),
    timeout_seconds: float | None = None,
    terminal: bool = False,
    normalizer: Callable[[dict[str, Any]], dict[str, Any]] | None = None,
) -> EventDefinition:
    return EventDefinition(
        name=name,
        description=f"Execute {name}.",
        parameters={"type": "object"},
        handler=handler,
        conflict_keys=conflict_keys,
        timeout_seconds=timeout_seconds,
        terminal=terminal,
        normalizer=normalizer,
    )


def _request(event_id: str, name: str, arguments: dict[str, Any] | None = None):
    return EventRequest(
        id=event_id,
        name=name,
        arguments=arguments or {},
    )


def _executor(
    definitions: list[EventDefinition],
    *,
    max_parallel_events: int = 4,
    batch_timeout_seconds: float = 1.0,
    **options: Any,
):
    return _executor_type()(
        EventKernel(EventRegistry(definitions)),
        max_parallel_events=max_parallel_events,
        batch_timeout_seconds=batch_timeout_seconds,
        **options,
    )


@pytest.mark.asyncio
async def test_batch_executor_bounds_global_parallelism():
    active = 0
    max_active = 0

    async def handler(request: EventRequest) -> dict[str, Any]:
        nonlocal active, max_active
        active += 1
        max_active = max(max_active, active)
        await asyncio.sleep(0.02)
        active -= 1
        return {"event_id": request.id}

    executor = _executor(
        [_definition("example.work", handler)],
        max_parallel_events=2,
    )

    batch = await executor.execute(
        [
            _request(str(index), "example.work", {"index": index})
            for index in range(4)
        ],
        scope="session-1:turn-1",
    )

    assert max_active == 2
    assert [result.status for result in batch.results] == [
        EventStatus.SUCCESS,
    ] * 4


@pytest.mark.asyncio
async def test_batch_executor_preserves_input_order_when_completion_order_differs():
    release_first = asyncio.Event()
    second_finished = asyncio.Event()

    async def handler(request: EventRequest) -> dict[str, Any]:
        if request.id == "first":
            await release_first.wait()
        else:
            second_finished.set()
        return {"event_id": request.id}

    executor = _executor([_definition("example.work", handler)])
    task = asyncio.create_task(
        executor.execute(
            [
                _request("first", "example.work", {"order": 1}),
                _request("second", "example.work", {"order": 2}),
            ],
            scope="session-1:turn-1",
        )
    )
    await asyncio.wait_for(second_finished.wait(), timeout=0.2)
    release_first.set()

    batch = await asyncio.wait_for(task, timeout=0.2)

    assert [result.event_id for result in batch.results] == ["first", "second"]


@pytest.mark.asyncio
async def test_batch_executor_serializes_shared_conflict_keys():
    active = 0
    max_active = 0

    async def handler(request: EventRequest) -> dict[str, Any]:
        nonlocal active, max_active
        active += 1
        max_active = max(max_active, active)
        await asyncio.sleep(0.02)
        active -= 1
        return {"event_id": request.id}

    executor = _executor(
        [
            _definition(
                "example.write",
                handler,
                conflict_keys=("shared-resource",),
            )
        ]
    )

    await executor.execute(
        [
            _request("first", "example.write", {"order": 1}),
            _request("second", "example.write", {"order": 2}),
        ],
        scope="session-1:turn-1",
    )

    assert max_active == 1


@pytest.mark.asyncio
async def test_batch_executor_isolates_per_event_timeout_from_successful_sibling():
    async def slow_handler(request: EventRequest) -> dict[str, Any]:
        await asyncio.sleep(1)
        return {"event_id": request.id}

    executor = _executor(
        [
            _definition(
                "example.slow",
                slow_handler,
                timeout_seconds=0.01,
            ),
            _definition(
                "example.fast",
                lambda request: {"event_id": request.id},
            ),
        ],
        batch_timeout_seconds=0.2,
    )

    batch = await executor.execute(
        [
            _request("slow", "example.slow"),
            _request("fast", "example.fast"),
        ],
        scope="session-1:turn-1",
    )

    assert [result.status for result in batch.results] == [
        EventStatus.TIMEOUT,
        EventStatus.SUCCESS,
    ]
    assert batch.deadline_exceeded is False


@pytest.mark.asyncio
async def test_batch_executor_applies_overall_deadline_to_waiting_work():
    async def blocked_handler(request: EventRequest) -> dict[str, Any]:
        await asyncio.Event().wait()
        return {"event_id": request.id}

    executor = _executor(
        [_definition("example.blocked", blocked_handler)],
        max_parallel_events=1,
        batch_timeout_seconds=0.02,
    )

    batch = await executor.execute(
        [
            _request("first", "example.blocked"),
            _request("second", "example.blocked"),
        ],
        scope="session-1:turn-1",
    )

    assert [result.status for result in batch.results] == [
        EventStatus.TIMEOUT,
        EventStatus.TIMEOUT,
    ]
    assert batch.deadline_exceeded is True


@pytest.mark.asyncio
async def test_batch_executor_coalesces_exact_duplicates_inside_one_batch():
    calls = 0

    def handler(request: EventRequest) -> dict[str, Any]:
        nonlocal calls
        calls += 1
        return {"event_id": request.id, "arguments": request.arguments}

    executor = _executor([_definition("example.write", handler)])

    batch = await executor.execute(
        [
            _request("same", "example.write", {"a": 1, "b": 2}),
            _request("same", "example.write", {"b": 2, "a": 1}),
        ],
        scope="session-1:turn-1",
    )

    assert calls == 1
    assert batch.coalesced_count == 1
    assert batch.results[0] == batch.results[1]


@pytest.mark.asyncio
async def test_batch_executor_replay_key_includes_scope_event_id_and_request():
    calls = 0

    def handler(request: EventRequest) -> dict[str, Any]:
        nonlocal calls
        calls += 1
        return {"call": calls, "arguments": request.arguments}

    executor = _executor([_definition("example.write", handler)])
    original = _request("same", "example.write", {"a": 1, "b": 2})
    canonical_replay = _request(
        "same",
        "example.write",
        {"b": 2, "a": 1},
    )

    first = await executor.execute([original], scope="session-1:turn-1")
    replay = await executor.execute(
        [canonical_replay],
        scope="session-1:turn-1",
    )
    changed_request = await executor.execute(
        [_request("same", "example.write", {"a": 2})],
        scope="session-1:turn-1",
    )
    changed_id = await executor.execute(
        [_request("other", "example.write", {"a": 1, "b": 2})],
        scope="session-1:turn-1",
    )
    changed_scope = await executor.execute(
        [original],
        scope="session-1:turn-2",
    )

    assert first.results[0] == replay.results[0]
    assert replay.replayed_count == 1
    assert changed_request.replayed_count == 0
    assert changed_id.replayed_count == 0
    assert changed_scope.replayed_count == 0
    assert calls == 4


@pytest.mark.asyncio
async def test_batch_executor_runs_terminal_events_after_all_nonterminal_results():
    observed: list[str] = []

    async def nonterminal(request: EventRequest) -> dict[str, Any]:
        observed.append(f"start:{request.id}")
        await asyncio.sleep(0.01)
        observed.append(f"finish:{request.id}")
        return {"event_id": request.id}

    def terminal(request: EventRequest) -> dict[str, Any]:
        observed.append(f"terminal:{request.id}")
        return {"event_id": request.id}

    executor = _executor(
        [
            _definition("example.work", nonterminal),
            _definition("example.terminate", terminal, terminal=True),
        ]
    )

    batch = await executor.execute(
        [
            _request("terminate", "example.terminate"),
            _request("first", "example.work", {"order": 1}),
            _request("second", "example.work", {"order": 2}),
        ],
        scope="session-1:turn-1",
    )

    terminal_index = observed.index("terminal:terminate")
    assert observed.index("finish:first") < terminal_index
    assert observed.index("finish:second") < terminal_index
    assert [result.event_id for result in batch.results] == [
        "terminate",
        "first",
        "second",
    ]


@pytest.mark.asyncio
async def test_batch_executor_does_not_swallow_external_cancellation():
    started = asyncio.Event()

    async def handler(request: EventRequest) -> dict[str, Any]:
        started.set()
        await asyncio.Event().wait()
        return {"event_id": request.id}

    executor = _executor([_definition("example.blocked", handler)])
    task = asyncio.create_task(
        executor.execute(
            [_request("blocked", "example.blocked")],
            scope="session-1:turn-1",
        )
    )
    await asyncio.wait_for(started.wait(), timeout=0.2)

    task.cancel()

    with pytest.raises(asyncio.CancelledError):
        await task


@pytest.mark.asyncio
async def test_batch_executor_coalesces_normalized_provider_arguments_across_ids():
    calls: list[EventRequest] = []

    def normalize(arguments: dict[str, Any]) -> dict[str, Any]:
        value = arguments["value"]
        return {"value": int(value) if isinstance(value, float) else value}

    executor = _executor(
        [
            _definition(
                "example.normalized",
                lambda request: calls.append(request) or {"value": request.arguments["value"]},
                normalizer=normalize,
            )
        ]
    )
    requests = [
        EventRequest(
            id="integer-id",
            name="example.normalized",
            arguments={"value": 30},
            source=EventSource.PROVIDER_RESOLVED,
            raw_arguments='{"value":30}',
        ),
        EventRequest(
            id="float-id",
            name="example.normalized",
            arguments={"value": 30.0},
            source=EventSource.PROVIDER_RESOLVED,
            raw_arguments='{"value":30.0}',
        ),
    ]

    batch = await executor.execute(requests, scope="normalized-scope")

    assert len(calls) == 1
    assert batch.coalesced_count == 1
    assert [result.event_id for result in batch.results] == [
        "integer-id",
        "float-id",
    ]
    assert batch.results[0].payload == {"value": 30}
    assert batch.results[1].payload == {
        "value": 30,
        "deduplicated_from": "integer-id",
    }


@pytest.mark.asyncio
async def test_batch_executor_tool_replies_keep_each_coalesced_call_id():
    registry = ToolRegistry(
        [
            ToolDefinition(
                name="example.tool",
                description="Execute a coalesced tool.",
                parameters={"type": "object"},
                handler=lambda event: {"handled_by": event.id},
            )
        ]
    )
    executor = _executor_type()(
        EventKernel(registry.event_registry),
        batch_timeout_seconds=1,
    )
    events = [
        ToolCallEvent(id=event_id, name="example.tool", arguments={}, raw_arguments="{}")
        for event_id in ("first-id", "second-id")
    ]
    batch = await executor.execute(
        [
            registry.event_request(event, source=EventSource.PROVIDER_RESOLVED)
            for event in events
        ],
        scope="tool-scope",
    )

    replies = registry.tool_replies(events, batch)

    assert [reply.tool_call_id for reply in replies] == ["first-id", "second-id"]
    assert [result.event_id for result in batch.results] == [
        "first-id",
        "second-id",
    ]


@pytest.mark.asyncio
async def test_batch_key_falls_back_to_original_arguments_when_normalizer_fails():
    calls = 0

    def normalize(arguments: dict[str, Any]) -> dict[str, Any]:
        if isinstance(arguments.get("value"), str):
            raise ValueError("unsupported value")
        return {"value": int(arguments["value"])}

    def handler(request: EventRequest) -> dict[str, Any]:
        nonlocal calls
        calls += 1
        return {"value": request.arguments["value"]}

    executor = _executor(
        [_definition("example.normalized", handler, normalizer=normalize)]
    )

    batch = await executor.execute(
        [
            EventRequest(
                id="valid",
                name="example.normalized",
                arguments={"value": 30},
                source=EventSource.PROVIDER_RESOLVED,
            ),
            EventRequest(
                id="invalid",
                name="example.normalized",
                arguments={"value": "30"},
                source=EventSource.PROVIDER_RESOLVED,
            ),
        ],
        scope="normalizer-failure",
    )

    assert batch.coalesced_count == 0
    assert [result.status for result in batch.results] == [
        EventStatus.SUCCESS,
        EventStatus.INVALID_ARGUMENTS,
    ]
    assert calls == 1


@pytest.mark.asyncio
async def test_concurrent_replay_waiters_share_execution_and_cancel_independently():
    calls = 0
    started = asyncio.Event()
    release = asyncio.Event()

    async def handler(request: EventRequest) -> dict[str, Any]:
        nonlocal calls
        calls += 1
        started.set()
        await release.wait()
        return {"event_id": request.id}

    executor = _executor([_definition("example.shared", handler)])
    request = _request("same-id", "example.shared")
    first = asyncio.create_task(executor.execute([request], scope="shared-scope"))
    second = asyncio.create_task(executor.execute([request], scope="shared-scope"))
    await asyncio.wait_for(started.wait(), timeout=0.2)

    first.cancel()
    with pytest.raises(asyncio.CancelledError):
        await first
    release.set()
    second_batch = await asyncio.wait_for(second, timeout=0.2)

    assert calls == 1
    assert second_batch.results[0].status is EventStatus.SUCCESS


@pytest.mark.asyncio
async def test_release_scope_removes_replay_entries_for_reuse():
    calls = 0

    def handler(request: EventRequest) -> dict[str, Any]:
        nonlocal calls
        calls += 1
        return {"call": calls}

    executor = _executor([_definition("example.cached", handler)])
    request = _request("same-id", "example.cached")

    await executor.execute([request], scope="released-scope")
    executor.release_scope("released-scope")
    batch = await executor.execute([request], scope="released-scope")

    assert calls == 2
    assert batch.replayed_count == 0


@pytest.mark.asyncio
async def test_replay_cache_stays_bounded_across_250_scopes():
    executor = _executor(
        [_definition("example.cached", lambda request: {"id": request.id})],
        max_replay_entries=100,
    )
    request = _request("same-id", "example.cached")

    for index in range(250):
        await executor.execute([request], scope=f"scope-{index}")

    assert len(executor._replay_results) <= 100


@pytest.mark.asyncio
async def test_terminal_uses_independent_grace_after_sibling_batch_timeout():
    terminal_calls: list[str] = []

    async def blocked(request: EventRequest) -> dict[str, Any]:
        await asyncio.Event().wait()
        return {"event_id": request.id}

    def terminate(request: EventRequest) -> dict[str, Any]:
        terminal_calls.append(request.id)
        return {"terminated": True}

    executor = _executor(
        [
            _definition("example.blocked", blocked),
            _definition(
                "example.terminate",
                terminate,
                terminal=True,
                timeout_seconds=0.05,
            ),
        ],
        batch_timeout_seconds=0.01,
        terminal_grace_seconds=0.1,
    )

    batch = await executor.execute(
        [
            _request("sibling", "example.blocked"),
            _request("terminal", "example.terminate"),
        ],
        scope="terminal-grace",
    )

    assert [result.status for result in batch.results] == [
        EventStatus.TIMEOUT,
        EventStatus.SUCCESS,
    ]
    assert batch.deadline_exceeded is True
    assert terminal_calls == ["terminal"]


@pytest.mark.asyncio
async def test_terminal_grace_is_bounded_by_terminal_event_timeout():
    async def blocked_terminal(request: EventRequest) -> dict[str, Any]:
        await asyncio.Event().wait()
        return {"event_id": request.id}

    executor = _executor(
        [
            _definition(
                "example.terminate",
                blocked_terminal,
                terminal=True,
                timeout_seconds=0.01,
            )
        ],
        terminal_grace_seconds=0.2,
    )

    batch = await executor.execute(
        [_request("terminal", "example.terminate")],
        scope="terminal-timeout",
    )

    assert batch.results[0].status is EventStatus.TIMEOUT
    assert batch.results[0].error == "event timed out after 0.01 seconds"
    assert batch.deadline_exceeded is False


@pytest.mark.asyncio
async def test_blocking_sync_handler_timeout_returns_before_thread_finishes():
    started = threading.Event()
    release = threading.Event()
    finished = threading.Event()

    def blocking_handler(request: EventRequest) -> dict[str, Any]:
        started.set()
        release.wait(timeout=1)
        finished.set()
        return {"event_id": request.id}

    executor = _executor(
        [
            _definition(
                "example.blocking",
                blocking_handler,
                timeout_seconds=0.02,
            )
        ],
        batch_timeout_seconds=0.2,
    )
    loop = asyncio.get_running_loop()
    started_at = loop.time()

    batch = await executor.execute(
        [_request("blocking", "example.blocking")],
        scope="blocking-handler",
    )
    elapsed = loop.time() - started_at
    release.set()
    await asyncio.to_thread(finished.wait, 0.2)

    assert started.is_set()
    assert batch.results[0].status is EventStatus.TIMEOUT
    assert elapsed < 0.1
    assert finished.is_set()