113 lines
3.7 KiB
Python
113 lines
3.7 KiB
Python
"""Bug C2 invariant test — AC-C2-2.
|
|
|
|
Verifies that two concurrent asyncio Tasks calling TurnRunner.run() for the
|
|
same session are properly serialized (max in-flight == 1) rather than running
|
|
concurrently due to the broken lock.locked() owner-blind check.
|
|
|
|
The fix replaces lock.locked() with a ContextVar-based owner check so Task B
|
|
cannot skip lock acquisition just because Task A holds it.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from collections.abc import AsyncIterator
|
|
from typing import Any
|
|
from unittest.mock import AsyncMock, MagicMock
|
|
|
|
import pytest
|
|
|
|
from opensquilla.engine.runtime import TurnRunner
|
|
from opensquilla.engine.types import DoneEvent
|
|
|
|
|
|
def _make_runner_with_lock(
|
|
shared_lock: asyncio.Lock,
|
|
) -> TurnRunner:
|
|
"""Build a minimal TurnRunner whose session lock is the provided shared_lock."""
|
|
provider = MagicMock()
|
|
provider.provider_name = "stub"
|
|
|
|
async def _chat(*args: Any, **kwargs: Any) -> AsyncIterator[Any]:
|
|
# Yield a small delay so concurrent tasks have time to overlap if
|
|
# serialization is broken.
|
|
await asyncio.sleep(0)
|
|
yield DoneEvent(stop_reason="end_turn", usage={})
|
|
|
|
provider.chat = _chat
|
|
|
|
selector = MagicMock()
|
|
selector.resolve.return_value = provider
|
|
selector.clone.return_value = selector
|
|
selector.current_config = MagicMock(model="stub-model")
|
|
|
|
session_manager = MagicMock()
|
|
session_manager.get = AsyncMock(return_value=None)
|
|
session_manager.append_message = AsyncMock(return_value=None)
|
|
session_manager.update = AsyncMock(return_value=None)
|
|
session_manager.get_compaction_summary = AsyncMock(return_value=None)
|
|
|
|
return TurnRunner(
|
|
provider_selector=selector,
|
|
session_manager=session_manager,
|
|
session_lock_provider=lambda key: shared_lock,
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_two_concurrent_runs_serialize() -> None:
|
|
"""Two concurrent TurnRunner.run() calls for the same session must serialize.
|
|
|
|
AC-C2-2: max tasks in-flight simultaneously must be 1, not 2.
|
|
|
|
If the bug is present (lock.locked() used instead of ContextVar owner check),
|
|
Task B sees locked()=True while Task A holds the lock and skips acquire,
|
|
causing both to run concurrently (max_in_flight == 2).
|
|
|
|
With the fix, Task B correctly waits for Task A to release the lock.
|
|
"""
|
|
session_key = "agent:main:concurrent-run-test"
|
|
shared_lock = asyncio.Lock()
|
|
runner = _make_runner_with_lock(shared_lock)
|
|
|
|
from opensquilla.tools.types import ToolContext
|
|
|
|
tool_ctx = ToolContext(session_key=session_key)
|
|
|
|
in_flight = 0
|
|
max_in_flight = 0
|
|
|
|
# Patch _run_turn to track concurrency
|
|
original_run_turn = runner._run_turn
|
|
|
|
async def _instrumented_run_turn(*args: Any, **kwargs: Any) -> AsyncIterator[Any]:
|
|
nonlocal in_flight, max_in_flight
|
|
in_flight += 1
|
|
if in_flight > max_in_flight:
|
|
max_in_flight = in_flight
|
|
try:
|
|
async for event in original_run_turn(*args, **kwargs):
|
|
yield event
|
|
finally:
|
|
in_flight -= 1
|
|
|
|
runner._run_turn = _instrumented_run_turn # type: ignore[method-assign]
|
|
|
|
async def _run_one() -> None:
|
|
async for _ in runner.run(
|
|
message="hello",
|
|
session_key=session_key,
|
|
tool_context=tool_ctx,
|
|
):
|
|
pass
|
|
|
|
# Launch two concurrent tasks for the same session
|
|
await asyncio.gather(_run_one(), _run_one())
|
|
|
|
assert max_in_flight == 1, (
|
|
f"Expected max 1 concurrent _run_turn execution per session, "
|
|
f"got {max_in_flight}. "
|
|
"Two concurrent TurnRunner.run() calls are not serialized — "
|
|
"the lock.locked() owner-blind check is still present."
|
|
)
|