瀏覽代碼

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 周之前
父節點
當前提交
317a2861f0
共有 2 個文件被更改,包括 198 次插入15 次删除
  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.abc import Hashable, Iterable, Sequence
 from dataclasses import dataclass, replace
+from itertools import count
 from typing import Any
 
 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]
+_SCOPE_TOKEN_COUNTER = count(1)
 
 
 @dataclass(frozen=True)
@@ -63,7 +65,7 @@ class EventBatchExecutor:
             OrderedDict()
         )
         self._inflight: dict[_ReplayKey, asyncio.Task[_ExecutionOutcome]] = {}
-        self._scope_generations: dict[Hashable, int] = {}
+        self._scope_tokens: dict[Hashable, int] = {}
         self._state_lock = asyncio.Lock()
 
     async def execute(
@@ -76,7 +78,10 @@ class EventBatchExecutor:
     ) -> EventBatchResult:
         if not requests:
             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] = []
         indexes_by_key: dict[str, list[int]] = {}
@@ -118,7 +123,7 @@ class EventBatchExecutor:
                     self._execute_replayable(
                         item,
                         scope=scope,
-                        generation=generation,
+                        scope_token=scope_token,
                         enabled_names=enabled_names,
                         context=context,
                         deadline=deadline,
@@ -151,7 +156,7 @@ class EventBatchExecutor:
 
         await self._cache_ordered_results(
             scope,
-            generation,
+            scope_token,
             requests,
             indexes_by_key,
             outcomes_by_key,
@@ -166,13 +171,14 @@ class EventBatchExecutor:
         )
 
     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):
-            if key[0] == scope and key[1] == generation:
+            if key[0] == scope and key[1] == scope_token:
                 self._replay_results.pop(key, None)
         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)
 
     async def _execute_replayable(
@@ -180,13 +186,13 @@ class EventBatchExecutor:
         item: _BatchItem,
         *,
         scope: Hashable,
-        generation: int,
+        scope_token: int,
         enabled_names: Iterable[str] | None,
         context: EventExecutionContext | None,
         deadline: float,
         count_deadline_exceeded: 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:
             replayed = self._replay_results.get(replay_key)
             if replayed is not None:
@@ -233,7 +239,7 @@ class EventBatchExecutor:
                     self._inflight.pop(replay_key, None)
                 if (
                     outcome is not None
-                    and self._scope_generations.get(replay_key[0], 0)
+                    and self._scope_tokens.get(replay_key[0])
                     == replay_key[1]
                 ):
                     self._store_replay_locked(replay_key, outcome)
@@ -243,14 +249,14 @@ class EventBatchExecutor:
     async def _cache_ordered_results(
         self,
         scope: Hashable,
-        generation: int,
+        scope_token: int,
         requests: Sequence[EventRequest],
         indexes_by_key: dict[str, list[int]],
         outcomes_by_key: dict[str, _ExecutionOutcome],
         ordered: list[EventResult | None],
     ) -> None:
         async with self._state_lock:
-            if self._scope_generations.get(scope, 0) != generation:
+            if self._scope_tokens.get(scope) != scope_token:
                 return
             for key, indexes in indexes_by_key.items():
                 outcome = outcomes_by_key[key]
@@ -258,7 +264,7 @@ class EventBatchExecutor:
                     result = ordered[index]
                     assert result is not None
                     self._store_replay_locked(
-                        (scope, generation, requests[index].id, key),
+                        (scope, scope_token, requests[index].id, key),
                         _ExecutionOutcome(
                             result,
                             batch_timed_out=outcome.batch_timed_out,
@@ -466,7 +472,10 @@ class EventBatchExecutor:
             copied = json.loads(serialized)
         except (TypeError, ValueError):
             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 serialized
 
@@ -478,3 +487,18 @@ class EventBatchExecutor:
         except (TypeError, ValueError):
             return {}
         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 threading
 from collections.abc import Callable
+from enum import IntEnum
 from typing import Any
 
 import pytest
@@ -719,6 +720,111 @@ async def test_non_json_normalizer_output_has_distinct_canonical_key():
     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
 async def test_concurrent_replay_waiters_share_execution_and_cancel_independently():
     calls = 0
@@ -746,6 +852,38 @@ async def test_concurrent_replay_waiters_share_execution_and_cancel_independentl
 
     assert calls == 1
     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
@@ -868,6 +1006,27 @@ async def test_replay_cache_stays_bounded_across_250_scopes():
     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
 async def test_terminal_uses_independent_grace_after_sibling_batch_timeout():
     terminal_calls: list[str] = []