Browse Source

fix: finalize batch canonical scope safety

Problem:
Canonical JSON round-trip checks relied on Python equality, allowing IntEnum values to collide with integers, while released scopes retained per-scope generation tombstones indefinitely.

Risk:
Scope tokens are process-local monotonic ownership identifiers rather than durable IDs. Released or reused scopes invalidate old ownership, so late in-flight work may finish but cannot repopulate replay or in-flight mappings.
zhenyu.hu 2 weeks ago
parent
commit
317a2861f0
2 changed files with 198 additions and 15 deletions
  1. 39 15
      src/agent_lab/application/events/batch.py
  2. 159 0
      tests/test_event_batch.py

+ 39 - 15
src/agent_lab/application/events/batch.py

@@ -6,6 +6,7 @@ import json
 from collections import OrderedDict
 from collections import OrderedDict
 from collections.abc import Hashable, Iterable, Sequence
 from collections.abc import Hashable, Iterable, Sequence
 from dataclasses import dataclass, replace
 from dataclasses import dataclass, replace
+from itertools import count
 from typing import Any
 from typing import Any
 
 
 from agent_lab.application.events.kernel import EventKernel
 from agent_lab.application.events.kernel import EventKernel
@@ -20,6 +21,7 @@ from agent_lab.application.events.models import (
 
 
 
 
 _ReplayKey = tuple[Hashable, int, str, str]
 _ReplayKey = tuple[Hashable, int, str, str]
+_SCOPE_TOKEN_COUNTER = count(1)
 
 
 
 
 @dataclass(frozen=True)
 @dataclass(frozen=True)
@@ -63,7 +65,7 @@ class EventBatchExecutor:
             OrderedDict()
             OrderedDict()
         )
         )
         self._inflight: dict[_ReplayKey, asyncio.Task[_ExecutionOutcome]] = {}
         self._inflight: dict[_ReplayKey, asyncio.Task[_ExecutionOutcome]] = {}
-        self._scope_generations: dict[Hashable, int] = {}
+        self._scope_tokens: dict[Hashable, int] = {}
         self._state_lock = asyncio.Lock()
         self._state_lock = asyncio.Lock()
 
 
     async def execute(
     async def execute(
@@ -76,7 +78,10 @@ class EventBatchExecutor:
     ) -> EventBatchResult:
     ) -> EventBatchResult:
         if not requests:
         if not requests:
             return EventBatchResult()
             return EventBatchResult()
-        generation = self._scope_generations.get(scope, 0)
+        scope_token = self._scope_tokens.get(scope)
+        if scope_token is None:
+            scope_token = next(_SCOPE_TOKEN_COUNTER)
+            self._scope_tokens[scope] = scope_token
 
 
         items: list[_BatchItem] = []
         items: list[_BatchItem] = []
         indexes_by_key: dict[str, list[int]] = {}
         indexes_by_key: dict[str, list[int]] = {}
@@ -118,7 +123,7 @@ class EventBatchExecutor:
                     self._execute_replayable(
                     self._execute_replayable(
                         item,
                         item,
                         scope=scope,
                         scope=scope,
-                        generation=generation,
+                        scope_token=scope_token,
                         enabled_names=enabled_names,
                         enabled_names=enabled_names,
                         context=context,
                         context=context,
                         deadline=deadline,
                         deadline=deadline,
@@ -151,7 +156,7 @@ class EventBatchExecutor:
 
 
         await self._cache_ordered_results(
         await self._cache_ordered_results(
             scope,
             scope,
-            generation,
+            scope_token,
             requests,
             requests,
             indexes_by_key,
             indexes_by_key,
             outcomes_by_key,
             outcomes_by_key,
@@ -166,13 +171,14 @@ class EventBatchExecutor:
         )
         )
 
 
     def release_scope(self, scope: Hashable) -> None:
     def release_scope(self, scope: Hashable) -> None:
-        generation = self._scope_generations.get(scope, 0)
-        self._scope_generations[scope] = generation + 1
+        scope_token = self._scope_tokens.pop(scope, None)
+        if scope_token is None:
+            return
         for key in tuple(self._replay_results):
         for key in tuple(self._replay_results):
-            if key[0] == scope and key[1] == generation:
+            if key[0] == scope and key[1] == scope_token:
                 self._replay_results.pop(key, None)
                 self._replay_results.pop(key, None)
         for key in tuple(self._inflight):
         for key in tuple(self._inflight):
-            if key[0] == scope and key[1] == generation:
+            if key[0] == scope and key[1] == scope_token:
                 self._inflight.pop(key, None)
                 self._inflight.pop(key, None)
 
 
     async def _execute_replayable(
     async def _execute_replayable(
@@ -180,13 +186,13 @@ class EventBatchExecutor:
         item: _BatchItem,
         item: _BatchItem,
         *,
         *,
         scope: Hashable,
         scope: Hashable,
-        generation: int,
+        scope_token: int,
         enabled_names: Iterable[str] | None,
         enabled_names: Iterable[str] | None,
         context: EventExecutionContext | None,
         context: EventExecutionContext | None,
         deadline: float,
         deadline: float,
         count_deadline_exceeded: bool,
         count_deadline_exceeded: bool,
     ) -> tuple[_ExecutionOutcome, bool]:
     ) -> tuple[_ExecutionOutcome, bool]:
-        replay_key = (scope, generation, item.request.id, item.key)
+        replay_key = (scope, scope_token, item.request.id, item.key)
         async with self._state_lock:
         async with self._state_lock:
             replayed = self._replay_results.get(replay_key)
             replayed = self._replay_results.get(replay_key)
             if replayed is not None:
             if replayed is not None:
@@ -233,7 +239,7 @@ class EventBatchExecutor:
                     self._inflight.pop(replay_key, None)
                     self._inflight.pop(replay_key, None)
                 if (
                 if (
                     outcome is not None
                     outcome is not None
-                    and self._scope_generations.get(replay_key[0], 0)
+                    and self._scope_tokens.get(replay_key[0])
                     == replay_key[1]
                     == replay_key[1]
                 ):
                 ):
                     self._store_replay_locked(replay_key, outcome)
                     self._store_replay_locked(replay_key, outcome)
@@ -243,14 +249,14 @@ class EventBatchExecutor:
     async def _cache_ordered_results(
     async def _cache_ordered_results(
         self,
         self,
         scope: Hashable,
         scope: Hashable,
-        generation: int,
+        scope_token: int,
         requests: Sequence[EventRequest],
         requests: Sequence[EventRequest],
         indexes_by_key: dict[str, list[int]],
         indexes_by_key: dict[str, list[int]],
         outcomes_by_key: dict[str, _ExecutionOutcome],
         outcomes_by_key: dict[str, _ExecutionOutcome],
         ordered: list[EventResult | None],
         ordered: list[EventResult | None],
     ) -> None:
     ) -> None:
         async with self._state_lock:
         async with self._state_lock:
-            if self._scope_generations.get(scope, 0) != generation:
+            if self._scope_tokens.get(scope) != scope_token:
                 return
                 return
             for key, indexes in indexes_by_key.items():
             for key, indexes in indexes_by_key.items():
                 outcome = outcomes_by_key[key]
                 outcome = outcomes_by_key[key]
@@ -258,7 +264,7 @@ class EventBatchExecutor:
                     result = ordered[index]
                     result = ordered[index]
                     assert result is not None
                     assert result is not None
                     self._store_replay_locked(
                     self._store_replay_locked(
-                        (scope, generation, requests[index].id, key),
+                        (scope, scope_token, requests[index].id, key),
                         _ExecutionOutcome(
                         _ExecutionOutcome(
                             result,
                             result,
                             batch_timed_out=outcome.batch_timed_out,
                             batch_timed_out=outcome.batch_timed_out,
@@ -466,7 +472,10 @@ class EventBatchExecutor:
             copied = json.loads(serialized)
             copied = json.loads(serialized)
         except (TypeError, ValueError):
         except (TypeError, ValueError):
             return None
             return None
-        if not isinstance(copied, dict) or copied != arguments:
+        if not isinstance(copied, dict) or not _json_values_equal(
+            copied,
+            arguments,
+        ):
             return None
             return None
         return serialized
         return serialized
 
 
@@ -478,3 +487,18 @@ class EventBatchExecutor:
         except (TypeError, ValueError):
         except (TypeError, ValueError):
             return {}
             return {}
         return value if isinstance(value, dict) else {}
         return value if isinstance(value, dict) else {}
+
+
+def _json_values_equal(left: Any, right: Any) -> bool:
+    if type(left) is not type(right):
+        return False
+    if isinstance(left, dict):
+        return left.keys() == right.keys() and all(
+            _json_values_equal(left[key], right[key]) for key in left
+        )
+    if isinstance(left, list):
+        return len(left) == len(right) and all(
+            _json_values_equal(left_item, right_item)
+            for left_item, right_item in zip(left, right, strict=True)
+        )
+    return left == right

+ 159 - 0
tests/test_event_batch.py

@@ -4,6 +4,7 @@ import asyncio
 import json
 import json
 import threading
 import threading
 from collections.abc import Callable
 from collections.abc import Callable
+from enum import IntEnum
 from typing import Any
 from typing import Any
 
 
 import pytest
 import pytest
@@ -719,6 +720,111 @@ async def test_non_json_normalizer_output_has_distinct_canonical_key():
     assert calls == 1
     assert calls == 1
 
 
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "wrap",
+    [
+        pytest.param(lambda value: {"value": value}, id="top-level"),
+        pytest.param(
+            lambda value: {"outer": {"value": value}},
+            id="nested-dict",
+        ),
+        pytest.param(
+            lambda value: {"outer": [value]},
+            id="nested-list",
+        ),
+    ],
+)
+async def test_canonical_json_equality_requires_recursive_exact_types(wrap):
+    class NumericEnum(IntEnum):
+        ONE = 1
+
+    calls: list[str] = []
+
+    def handler(request: EventRequest) -> dict[str, Any]:
+        calls.append(request.id)
+        return {"event_id": request.id}
+
+    executor = _executor([_definition("example.strict-types", handler)])
+    values = [NumericEnum.ONE, True, 1, 1.0]
+    requests = [
+        EventRequest(
+            id=event_id,
+            name="example.strict-types",
+            arguments=wrap(value),
+            source=EventSource.PROVIDER_RESOLVED,
+            raw_arguments=json.dumps(wrap(value)),
+        )
+        for event_id, value in zip(
+            ("enum", "bool", "int", "float"),
+            values,
+            strict=True,
+        )
+    ]
+
+    batch = await executor.execute(requests, scope=f"strict-types-{id(wrap)}")
+
+    assert batch.coalesced_count == 0
+    assert [result.status for result in batch.results] == [
+        EventStatus.INVALID_ARGUMENTS,
+        EventStatus.SUCCESS,
+        EventStatus.SUCCESS,
+        EventStatus.SUCCESS,
+    ]
+    assert calls == ["bool", "int", "float"]
+
+
+@pytest.mark.asyncio
+async def test_normalizer_explicitly_coalesces_valid_numeric_types_to_int():
+    class NumericEnum(IntEnum):
+        ONE = 1
+
+    calls = 0
+
+    def normalize(arguments: dict[str, Any]) -> dict[str, Any]:
+        return {"value": int(arguments["value"])}
+
+    def handler(request: EventRequest) -> dict[str, Any]:
+        nonlocal calls
+        calls += 1
+        return {"event_id": request.id, "value": request.arguments["value"]}
+
+    executor = _executor(
+        [
+            _definition(
+                "example.normalized-types",
+                handler,
+                normalizer=normalize,
+            )
+        ]
+    )
+    requests = [
+        EventRequest(
+            id=event_id,
+            name="example.normalized-types",
+            arguments={"value": value},
+            source=EventSource.PROVIDER_RESOLVED,
+            raw_arguments=json.dumps({"value": value}),
+        )
+        for event_id, value in zip(
+            ("enum", "bool", "int", "float"),
+            (NumericEnum.ONE, True, 1, 1.0),
+            strict=True,
+        )
+    ]
+
+    batch = await executor.execute(requests, scope="normalized-strict-types")
+
+    assert batch.coalesced_count == 2
+    assert [result.status for result in batch.results] == [
+        EventStatus.INVALID_ARGUMENTS,
+        EventStatus.SUCCESS,
+        EventStatus.SUCCESS,
+        EventStatus.SUCCESS,
+    ]
+    assert calls == 1
+
+
 @pytest.mark.asyncio
 @pytest.mark.asyncio
 async def test_concurrent_replay_waiters_share_execution_and_cancel_independently():
 async def test_concurrent_replay_waiters_share_execution_and_cancel_independently():
     calls = 0
     calls = 0
@@ -746,6 +852,38 @@ async def test_concurrent_replay_waiters_share_execution_and_cancel_independentl
 
 
     assert calls == 1
     assert calls == 1
     assert second_batch.results[0].status is EventStatus.SUCCESS
     assert second_batch.results[0].status is EventStatus.SUCCESS
+    scope_tokens = getattr(executor, "_scope_tokens", None)
+    assert scope_tokens is not None
+    assert len(scope_tokens) == 1
+    assert scope_tokens["shared-scope"] > 0
+
+
+@pytest.mark.asyncio
+async def test_scope_tokens_are_globally_unique_removed_and_never_reused():
+    async def handler(request: EventRequest) -> dict[str, Any]:
+        return {"event_id": request.id}
+
+    executor = _executor([_definition("example.token", handler)])
+    request = _request("same-id", "example.token")
+
+    await executor.execute([request], scope="scope-a")
+    await executor.execute([request], scope="scope-b")
+    scope_tokens = getattr(executor, "_scope_tokens", None)
+
+    assert scope_tokens is not None
+    first_a = scope_tokens["scope-a"]
+    first_b = scope_tokens["scope-b"]
+    assert first_a != first_b
+
+    executor.release_scope("scope-a")
+    assert "scope-a" not in scope_tokens
+    await executor.execute([request], scope="scope-a")
+    second_a = scope_tokens["scope-a"]
+
+    assert second_a not in {first_a, first_b}
+    executor.release_scope("scope-a")
+    executor.release_scope("scope-b")
+    assert scope_tokens == {}
 
 
 
 
 @pytest.mark.asyncio
 @pytest.mark.asyncio
@@ -868,6 +1006,27 @@ async def test_replay_cache_stays_bounded_across_250_scopes():
     assert len(executor._replay_results) <= 100
     assert len(executor._replay_results) <= 100
 
 
 
 
+@pytest.mark.asyncio
+async def test_thousand_released_scopes_leave_no_token_or_replay_state():
+    async def handler(request: EventRequest) -> dict[str, Any]:
+        return {"event_id": request.id}
+
+    executor = _executor([_definition("example.cleanup", handler)])
+    request = _request("same-id", "example.cleanup")
+
+    for index in range(1000):
+        scope = f"released-{index}"
+        await executor.execute([request], scope=scope)
+        executor.release_scope(scope)
+
+    scope_tokens = getattr(executor, "_scope_tokens", None)
+    retained_generations = getattr(executor, "_scope_generations", {})
+    assert scope_tokens == {}
+    assert retained_generations == {}
+    assert executor._replay_results == {}
+    assert executor._inflight == {}
+
+
 @pytest.mark.asyncio
 @pytest.mark.asyncio
 async def test_terminal_uses_independent_grace_after_sibling_batch_timeout():
 async def test_terminal_uses_independent_grace_after_sibling_batch_timeout():
     terminal_calls: list[str] = []
     terminal_calls: list[str] = []