|
|
@@ -1,6 +1,7 @@
|
|
|
from __future__ import annotations
|
|
|
|
|
|
import asyncio
|
|
|
+import threading
|
|
|
from collections.abc import Callable
|
|
|
from typing import Any
|
|
|
|
|
|
@@ -12,8 +13,11 @@ from agent_lab.application.events import (
|
|
|
EventKernel,
|
|
|
EventRegistry,
|
|
|
EventRequest,
|
|
|
+ EventSource,
|
|
|
EventStatus,
|
|
|
)
|
|
|
+from agent_lab.application.tools import ToolDefinition, ToolRegistry
|
|
|
+from agent_lab.domain.events import ToolCallEvent
|
|
|
|
|
|
|
|
|
def _executor_type():
|
|
|
@@ -29,6 +33,7 @@ def _definition(
|
|
|
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,
|
|
|
@@ -38,6 +43,7 @@ def _definition(
|
|
|
conflict_keys=conflict_keys,
|
|
|
timeout_seconds=timeout_seconds,
|
|
|
terminal=terminal,
|
|
|
+ normalizer=normalizer,
|
|
|
)
|
|
|
|
|
|
|
|
|
@@ -54,11 +60,13 @@ def _executor(
|
|
|
*,
|
|
|
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,
|
|
|
)
|
|
|
|
|
|
|
|
|
@@ -81,7 +89,10 @@ async def test_batch_executor_bounds_global_parallelism():
|
|
|
)
|
|
|
|
|
|
batch = await executor.execute(
|
|
|
- [_request(str(index), "example.work") for index in range(4)],
|
|
|
+ [
|
|
|
+ _request(str(index), "example.work", {"index": index})
|
|
|
+ for index in range(4)
|
|
|
+ ],
|
|
|
scope="session-1:turn-1",
|
|
|
)
|
|
|
|
|
|
@@ -107,8 +118,8 @@ async def test_batch_executor_preserves_input_order_when_completion_order_differ
|
|
|
task = asyncio.create_task(
|
|
|
executor.execute(
|
|
|
[
|
|
|
- _request("first", "example.work"),
|
|
|
- _request("second", "example.work"),
|
|
|
+ _request("first", "example.work", {"order": 1}),
|
|
|
+ _request("second", "example.work", {"order": 2}),
|
|
|
],
|
|
|
scope="session-1:turn-1",
|
|
|
)
|
|
|
@@ -146,8 +157,8 @@ async def test_batch_executor_serializes_shared_conflict_keys():
|
|
|
|
|
|
await executor.execute(
|
|
|
[
|
|
|
- _request("first", "example.write"),
|
|
|
- _request("second", "example.write"),
|
|
|
+ _request("first", "example.write", {"order": 1}),
|
|
|
+ _request("second", "example.write", {"order": 2}),
|
|
|
],
|
|
|
scope="session-1:turn-1",
|
|
|
)
|
|
|
@@ -309,8 +320,8 @@ async def test_batch_executor_runs_terminal_events_after_all_nonterminal_results
|
|
|
batch = await executor.execute(
|
|
|
[
|
|
|
_request("terminate", "example.terminate"),
|
|
|
- _request("first", "example.work"),
|
|
|
- _request("second", "example.work"),
|
|
|
+ _request("first", "example.work", {"order": 1}),
|
|
|
+ _request("second", "example.work", {"order": 2}),
|
|
|
],
|
|
|
scope="session-1:turn-1",
|
|
|
)
|
|
|
@@ -347,3 +358,305 @@ async def test_batch_executor_does_not_swallow_external_cancellation():
|
|
|
|
|
|
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()
|