agent-lab/tests/test_event_batch.py

from __future__ import annotations

import asyncio
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,
    EventStatus,
)


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,
) -> EventDefinition:
    return EventDefinition(
        name=name,
        description=f"Execute {name}.",
        parameters={"type": "object"},
        handler=handler,
        conflict_keys=conflict_keys,
        timeout_seconds=timeout_seconds,
        terminal=terminal,
    )


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,
):
    return _executor_type()(
        EventKernel(EventRegistry(definitions)),
        max_parallel_events=max_parallel_events,
        batch_timeout_seconds=batch_timeout_seconds,
    )


@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") 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"),
                _request("second", "example.work"),
            ],
            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"),
            _request("second", "example.write"),
        ],
        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"),
            _request("second", "example.work"),
        ],
        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