Files
opensquilla--opensquilla/tests/test_engine/test_concurrent_run_serialized.py
2026-07-13 13:12:33 +08:00

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."
)