# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from __future__ import annotations import asyncio import logging from dataclasses import FrozenInstanceError import msgspec import pytest from vllm_omni.engine.duplex.control_plane import ( DuplexControlPlane, DuplexOutputContext, DuplexRequestIdentity, DuplexStageRequestContext, DuplexStageSubmission, DuplexStageSubmissionResult, ) from vllm_omni.engine.duplex.lease import DuplexLeaseActivity, DuplexLeaseConfig from vllm_omni.engine.duplex.messages import ( AppendDuplexInputMessage, CloseDuplexSessionMessage, DuplexFence, DuplexSessionLifecycleMessage, OpenDuplexSessionMessage, ResumeDuplexSessionMessage, SignalDuplexTurnMessage, TouchDuplexSessionMessage, ) from vllm_omni.engine.duplex.runtime import ( DuplexAppendPlan, DuplexInputMode, DuplexRuntimeCapabilities, ) pytestmark = [pytest.mark.core_model, pytest.mark.cpu] class _Extension: def configure_sampling_params(self, *, runtime_config, defaults): del runtime_config return tuple(f"configured-{stage_id}" for stage_id, _ in enumerate(defaults)) def plan_append( self, *, request_id, fence, session_config, runtime_config, seq, turn_seq, mode, payload, final, sampling_params, ): del request_id, fence, session_config, runtime_config, seq, turn_seq, mode, payload, final assert sampling_params == "configured-0" return DuplexAppendPlan(prompt={"prompt_token_ids": [1, 2, 3]}) def decide_output(self, **kwargs): del kwargs return None class _TypedStagePort: stage_count = 2 def __init__(self) -> None: self.ensure_calls: list[DuplexStageRequestContext] = [] self.submit_calls: list[DuplexStageSubmission] = [] self.cleanup_calls: list[tuple[list[str], bool]] = [] def sampling_defaults(self) -> tuple[object, ...]: return ("default-0", "default-1") def ensure_request(self, context: DuplexStageRequestContext) -> None: self.ensure_calls.append(context) async def submit(self, submission: DuplexStageSubmission) -> DuplexStageSubmissionResult: self.submit_calls.append(submission) return DuplexStageSubmissionResult( request_id=submission.context.request_id, stage_id=submission.context.stage_id, replica_id=3, ) async def cleanup(self, request_ids: list[str], *, abort: bool = False) -> None: self.cleanup_calls.append((request_ids, abort)) class _Clock: def __init__(self) -> None: self.value = 0.0 def __call__(self) -> float: return self.value def advance(self, seconds: float) -> None: self.value += seconds @pytest.mark.asyncio async def test_control_plane_uses_frozen_typed_stage_context_without_request_state() -> None: stage_port = _TypedStagePort() result_sink: asyncio.Queue = asyncio.Queue() plane = DuplexControlPlane( extension=_Extension(), stage_port=stage_port, result_sink=result_sink, ) fence = DuplexFence("typed-port") await plane.handle( OpenDuplexSessionMessage( control_id="open", fence=fence, session_id=fence.session_id, capabilities={ "input_modes": [DuplexInputMode.APPEND_AUDIO_CHUNK.value], }, session_config={"voice": "test"}, runtime_config={"runtime": "test"}, ) ) assert (await result_sink.get()).ok is True context = stage_port.ensure_calls[-1] assert context.session_id == fence.session_id assert context.fence == fence assert context.stage_id == 0 assert context.final_stage_id == 1 assert context.sampling_params == ("configured-0", "configured-1") assert context.session_config == {"voice": "test"} assert context.runtime_config == {"runtime": "test"} with pytest.raises(FrozenInstanceError): context.stage_id = 1 # type: ignore[misc] await plane.handle( AppendDuplexInputMessage( control_id="append", fence=fence, session_id=fence.session_id, mode=DuplexInputMode.APPEND_AUDIO_CHUNK.value, payload={"audio": b"pcm"}, final=True, ) ) append_result = await result_sink.get() assert append_result.ok is True assert append_result.stage_results[0]["replica_id"] == 3 assert append_result.stage_results[0]["result"]["response_stage_id"] == 1 submission = stage_port.submit_calls[-1] assert submission.context == stage_port.ensure_calls[-1] assert submission.prompt == {"prompt_token_ids": [1, 2, 3]} assert submission.already_submitted is False assert not hasattr(submission, "request_state") def test_control_plane_accepts_only_typed_duplex_messages() -> None: plane = DuplexControlPlane( extension=None, stage_port=_TypedStagePort(), result_sink=asyncio.Queue(), ) assert plane.accepts( OpenDuplexSessionMessage( control_id="open", fence=DuplexFence("typed-message"), session_id="typed-message", capabilities=DuplexRuntimeCapabilities().__dict__, ) ) assert plane.accepts(type("Lookalike", (), {"type": "open_duplex_session"})()) is False @pytest.mark.parametrize( "message", [ TouchDuplexSessionMessage( control_id="touch-1", fence=DuplexFence("sid-message"), session_id="sid-message", activity=DuplexLeaseActivity.HEARTBEAT.value, ), ResumeDuplexSessionMessage( control_id="resume-1", fence=DuplexFence("sid-message"), session_id="sid-message", expected_lease_generation=3, ), DuplexSessionLifecycleMessage( fence=DuplexFence("sid-message"), session_id="sid-message", event="expired", reason="idle_ttl_expired", lease_generation=4, submitted_request_ids=["req-submitted"], reserved_request_ids=["req-reserved"], ), ], ) def test_duplex_lease_messages_round_trip(message) -> None: encoded = msgspec.json.encode(message) decoded = msgspec.json.decode(encoded, type=type(message)) assert decoded == message @pytest.mark.asyncio async def test_control_plane_touches_append_signal_and_explicit_activity() -> None: clock = _Clock() stage_port = _TypedStagePort() result_sink: asyncio.Queue = asyncio.Queue() lifecycle_sink: asyncio.Queue = asyncio.Queue() plane = DuplexControlPlane( extension=_Extension(), stage_port=stage_port, result_sink=result_sink, lifecycle_sink=lifecycle_sink, lease_config=DuplexLeaseConfig(idle_ttl_s=30.0, disconnect_grace_s=5.0), clock=clock, ) fence = DuplexFence("sid-touch") await plane.handle( OpenDuplexSessionMessage( control_id="open", fence=fence, session_id=fence.session_id, capabilities={"input_modes": [DuplexInputMode.APPEND_AUDIO_CHUNK.value]}, ) ) await result_sink.get() session = plane.sessions.require(fence.session_id) clock.advance(1.0) await plane.handle( AppendDuplexInputMessage( control_id="append", fence=fence, session_id=fence.session_id, mode=DuplexInputMode.APPEND_AUDIO_CHUNK.value, payload={"audio": b"pcm"}, ) ) assert (await result_sink.get()).ok is True assert session.lease.last_activity == 1.0 clock.advance(1.0) await plane.handle( SignalDuplexTurnMessage( control_id="signal", fence=fence, session_id=fence.session_id, event="session.update", ) ) assert (await result_sink.get()).ok is True assert session.lease.last_activity == 2.0 submit_count = len(stage_port.submit_calls) clock.advance(1.0) await plane.handle( TouchDuplexSessionMessage( control_id="heartbeat", fence=fence, session_id=fence.session_id, activity=DuplexLeaseActivity.HEARTBEAT.value, ) ) touch_result = await result_sink.get() assert touch_result.ok is True assert session.lease.last_activity == 3.0 assert len(stage_port.submit_calls) == submit_count assert lifecycle_sink.empty() @pytest.mark.asyncio async def test_control_plane_resume_requires_expected_lease_generation() -> None: clock = _Clock() result_sink: asyncio.Queue = asyncio.Queue() plane = DuplexControlPlane( extension=None, stage_port=_TypedStagePort(), result_sink=result_sink, lifecycle_sink=asyncio.Queue(), lease_config=DuplexLeaseConfig(), clock=clock, ) fence = DuplexFence("sid-resume-control") await plane.handle( OpenDuplexSessionMessage( control_id="open", fence=fence, session_id=fence.session_id, capabilities={}, ) ) await result_sink.get() await plane.handle( TouchDuplexSessionMessage( control_id="detach", fence=fence, session_id=fence.session_id, activity=DuplexLeaseActivity.DETACH.value, ) ) await result_sink.get() await plane.handle( ResumeDuplexSessionMessage( control_id="resume", fence=fence, session_id=fence.session_id, expected_lease_generation=0, ) ) result = await result_sink.get() assert result.ok is True assert result.stage_results[0]["result"]["lease_generation"] == 1 await plane.handle( ResumeDuplexSessionMessage( control_id="stale-resume", fence=fence, session_id=fence.session_id, expected_lease_generation=0, ) ) stale_result = await result_sink.get() assert stale_result.ok is False assert stale_result.accepted_fence == fence assert stale_result.lease_generation == 1 @pytest.mark.asyncio async def test_control_plane_reaps_one_session_through_cleanup_and_lifecycle_sink() -> None: clock = _Clock() stage_port = _TypedStagePort() result_sink: asyncio.Queue = asyncio.Queue() lifecycle_sink: asyncio.Queue = asyncio.Queue() plane = DuplexControlPlane( extension=None, stage_port=stage_port, result_sink=result_sink, lifecycle_sink=lifecycle_sink, lease_config=DuplexLeaseConfig(idle_ttl_s=2.0, disconnect_grace_s=1.0), clock=clock, ) for session_id in ("sid-a", "sid-b"): fence = DuplexFence(session_id) await plane.handle( OpenDuplexSessionMessage( control_id=f"open-{session_id}", fence=fence, session_id=session_id, capabilities={}, ) ) await result_sink.get() clock.advance(1.0) plane.sessions.require("sid-b").touch(DuplexFence("sid-b"), DuplexLeaseActivity.HEARTBEAT) clock.advance(1.1) expired_count = await plane.reap_expired() assert expired_count == 1 assert plane.sessions.get("sid-a") is None assert plane.sessions.get("sid-b") is not None assert stage_port.cleanup_calls == [] lifecycle = await lifecycle_sink.get() assert isinstance(lifecycle, DuplexSessionLifecycleMessage) assert lifecycle.session_id == "sid-a" assert lifecycle.reason == "idle_ttl_expired" assert result_sink.empty() @pytest.mark.asyncio async def test_control_plane_retries_expired_cleanup_before_publishing_lifecycle() -> None: class _FailOnceStagePort(_TypedStagePort): def __init__(self) -> None: super().__init__() self.failures_remaining = 1 async def cleanup(self, request_ids: list[str], *, abort: bool = False) -> None: if self.failures_remaining: self.failures_remaining -= 1 raise RuntimeError("transient cleanup failure") await super().cleanup(request_ids, abort=abort) clock = _Clock() stage_port = _FailOnceStagePort() result_sink: asyncio.Queue = asyncio.Queue() lifecycle_sink: asyncio.Queue = asyncio.Queue() plane = DuplexControlPlane( extension=None, stage_port=stage_port, result_sink=result_sink, lifecycle_sink=lifecycle_sink, lease_config=DuplexLeaseConfig(idle_ttl_s=1.0, disconnect_grace_s=1.0), clock=clock, ) fence = DuplexFence("sid-retry-expiry") await plane.handle( OpenDuplexSessionMessage( control_id="open-retry", fence=fence, session_id=fence.session_id, capabilities={}, ) ) await result_sink.get() request_id = plane.stage_request_id(fence, stage_id=0) plane.sessions.require(fence.session_id).reserve_stage_request(0, request_id, fence=fence) clock.advance(2.0) assert await plane.reap_expired() == 0 assert lifecycle_sink.empty() assert await plane.reap_expired() == 1 assert (await lifecycle_sink.get()).session_id == fence.session_id @pytest.mark.asyncio async def test_expired_session_holds_admission_slot_while_cleanup_is_blocked() -> None: class _BlockedCleanupStagePort(_TypedStagePort): def __init__(self) -> None: super().__init__() self.cleanup_started = asyncio.Event() self.release_cleanup = asyncio.Event() async def cleanup(self, request_ids: list[str], *, abort: bool = False) -> None: self.cleanup_started.set() await self.release_cleanup.wait() await super().cleanup(request_ids, abort=abort) clock = _Clock() stage_port = _BlockedCleanupStagePort() result_sink: asyncio.Queue = asyncio.Queue() plane = DuplexControlPlane( extension=None, stage_port=stage_port, result_sink=result_sink, lease_config=DuplexLeaseConfig(idle_ttl_s=1.0, disconnect_grace_s=1.0), clock=clock, max_sessions=1, ) old_fence = DuplexFence("sid-expired-blocked") await plane.handle( OpenDuplexSessionMessage( control_id="open-expired-blocked", fence=old_fence, session_id=old_fence.session_id, capabilities={}, ) ) assert (await result_sink.get()).ok is True old_session = plane.sessions.require(old_fence.session_id) old_session.bind_stage_request(0, "req-expired-blocked", fence=old_fence) clock.advance(2.0) reap_task = asyncio.create_task(plane.reap_expired()) await asyncio.wait_for(stage_port.cleanup_started.wait(), timeout=1.0) replacement_fence = DuplexFence("sid-replacement-blocked") await plane.handle( OpenDuplexSessionMessage( control_id="open-replacement-blocked", fence=replacement_fence, session_id=replacement_fence.session_id, capabilities={}, ) ) blocked_result = await result_sink.get() assert blocked_result.ok is False assert blocked_result.error is not None assert blocked_result.error.code == "resource_exhausted" assert plane.sessions.get(old_fence.session_id) is old_session stage_port.release_cleanup.set() assert await reap_task == 1 assert plane.sessions.get(old_fence.session_id) is None await plane.handle( OpenDuplexSessionMessage( control_id="open-replacement-after-cleanup", fence=replacement_fence, session_id=replacement_fence.session_id, capabilities={}, ) ) assert (await result_sink.get()).ok is True @pytest.mark.asyncio async def test_failed_expiry_cleanup_retains_admission_slot_for_retry() -> None: class _FailingCleanupStagePort(_TypedStagePort): async def cleanup(self, request_ids: list[str], *, abort: bool = False) -> None: del request_ids, abort raise RuntimeError("cleanup failed") clock = _Clock() result_sink: asyncio.Queue = asyncio.Queue() plane = DuplexControlPlane( extension=None, stage_port=_FailingCleanupStagePort(), result_sink=result_sink, lease_config=DuplexLeaseConfig(idle_ttl_s=1.0, disconnect_grace_s=1.0), clock=clock, max_sessions=1, ) old_fence = DuplexFence("sid-expired-failed") await plane.handle( OpenDuplexSessionMessage( control_id="open-expired-failed", fence=old_fence, session_id=old_fence.session_id, capabilities={}, ) ) assert (await result_sink.get()).ok is True old_session = plane.sessions.require(old_fence.session_id) old_session.bind_stage_request(0, "req-expired-failed", fence=old_fence) clock.advance(2.0) assert await plane.reap_expired() == 0 replacement_fence = DuplexFence("sid-replacement-failed") await plane.handle( OpenDuplexSessionMessage( control_id="open-replacement-failed", fence=replacement_fence, session_id=replacement_fence.session_id, capabilities={}, ) ) blocked_result = await result_sink.get() assert blocked_result.ok is False assert blocked_result.error is not None assert blocked_result.error.code == "resource_exhausted" assert plane.sessions.get(old_fence.session_id) is old_session @pytest.mark.asyncio async def test_capacity_rejection_is_an_info_event(caplog: pytest.LogCaptureFixture) -> None: result_sink: asyncio.Queue = asyncio.Queue() plane = DuplexControlPlane( extension=None, stage_port=_TypedStagePort(), result_sink=result_sink, max_sessions=1, ) first_fence = DuplexFence("sid-capacity-first") await plane.handle( OpenDuplexSessionMessage( control_id="open-capacity-first", fence=first_fence, session_id=first_fence.session_id, capabilities={}, ) ) assert (await result_sink.get()).ok is True caplog.clear() rejected_fence = DuplexFence("sid-capacity-rejected") await plane.handle( OpenDuplexSessionMessage( control_id="open-capacity-rejected", fence=rejected_fence, session_id=rejected_fence.session_id, capabilities={}, ) ) result = await result_sink.get() assert result.error.code == "resource_exhausted" assert not [record for record in caplog.records if record.levelno >= logging.ERROR] @pytest.mark.asyncio async def test_failed_request_cleanup_retries_before_releasing_admission() -> None: class _FailOnceCleanupStagePort(_TypedStagePort): def __init__(self) -> None: super().__init__() self.attempts = 0 async def cleanup(self, request_ids: list[str], *, abort: bool = False) -> None: self.cleanup_calls.append((list(request_ids), abort)) self.attempts += 1 if self.attempts == 1: raise RuntimeError("transient request cleanup failure") stage_port = _FailOnceCleanupStagePort() result_sink: asyncio.Queue = asyncio.Queue() plane = DuplexControlPlane( extension=None, stage_port=stage_port, result_sink=result_sink, max_sessions=1, ) old_fence = DuplexFence("sid-request-cleanup") await plane.handle( OpenDuplexSessionMessage( control_id="open-request-cleanup", fence=old_fence, session_id=old_fence.session_id, capabilities={}, ) ) assert (await result_sink.get()).ok is True old_session = plane.sessions.require(old_fence.session_id) old_session.bind_stage_request(0, "req-request-cleanup", fence=old_fence) closed = plane.close_sessions_for_request_ids( ["req-request-cleanup"], abort=True, ) assert closed == {old_fence.session_id: ["req-request-cleanup"]} assert plane.sessions.get(old_fence.session_id) is old_session with pytest.raises(RuntimeError, match="transient request cleanup failure"): await stage_port.cleanup(closed[old_fence.session_id], abort=True) replacement_fence = DuplexFence("sid-request-cleanup-replacement") await plane.handle( OpenDuplexSessionMessage( control_id="open-request-cleanup-replacement-blocked", fence=replacement_fence, session_id=replacement_fence.session_id, capabilities={}, ) ) assert (await result_sink.get()).error.code == "resource_exhausted" assert await plane.reap_expired() == 1 assert plane.sessions.get(old_fence.session_id) is None assert stage_port.cleanup_calls == [ (["req-request-cleanup"], True), (["req-request-cleanup"], True), ] await plane.handle( OpenDuplexSessionMessage( control_id="open-request-cleanup-replacement", fence=replacement_fence, session_id=replacement_fence.session_id, capabilities={}, ) ) assert (await result_sink.get()).ok is True @pytest.mark.asyncio async def test_stale_request_cleanup_does_not_finalize_reopened_incarnation() -> None: stage_port = _TypedStagePort() result_sink: asyncio.Queue = asyncio.Queue() plane = DuplexControlPlane( extension=None, stage_port=stage_port, result_sink=result_sink, max_sessions=1, ) old_fence = DuplexFence("sid-request-cleanup-reopen") await plane.handle( OpenDuplexSessionMessage( control_id="open-request-cleanup-old", fence=old_fence, session_id=old_fence.session_id, capabilities={}, ) ) assert (await result_sink.get()).ok is True old_session = plane.sessions.require(old_fence.session_id) old_session.bind_stage_request(0, "req-request-cleanup-old", fence=old_fence) closed = plane.close_sessions_for_request_ids( ["req-request-cleanup-old"], abort=True, cleanup_in_progress=True, ) assert closed == {old_fence.session_id: ["req-request-cleanup-old"]} plane.defer_request_cleanups([old_fence.session_id]) await plane.handle( CloseDuplexSessionMessage( control_id="close-request-cleanup-old", fence=old_fence, session_id=old_fence.session_id, reason="client_close", ) ) assert (await result_sink.get()).ok is True assert plane.sessions.get(old_fence.session_id) is None new_fence = DuplexFence(old_fence.session_id, incarnation=1) await plane.handle( OpenDuplexSessionMessage( control_id="open-request-cleanup-new", fence=new_fence, session_id=new_fence.session_id, capabilities={}, ) ) assert (await result_sink.get()).ok is True new_session = plane.sessions.require(new_fence.session_id) assert new_session.resume(new_fence, expected_lease_generation=0) == 1 assert await plane.reap_expired() == 1 assert plane.sessions.get(new_fence.session_id) is new_session @pytest.mark.asyncio async def test_failed_open_cleanup_is_retried_before_releasing_admission() -> None: class _FailOpenAndFirstCleanupStagePort(_TypedStagePort): def __init__(self) -> None: super().__init__() self.ensure_attempts = 0 self.cleanup_attempts = 0 def ensure_request(self, context: DuplexStageRequestContext) -> None: self.ensure_attempts += 1 if self.ensure_attempts == 1: raise RuntimeError("initial stage setup failed") super().ensure_request(context) async def cleanup(self, request_ids: list[str], *, abort: bool = False) -> None: self.cleanup_calls.append((list(request_ids), abort)) self.cleanup_attempts += 1 if self.cleanup_attempts == 1: raise RuntimeError("transient open rollback failure") stage_port = _FailOpenAndFirstCleanupStagePort() result_sink: asyncio.Queue = asyncio.Queue() plane = DuplexControlPlane( extension=_Extension(), stage_port=stage_port, result_sink=result_sink, max_sessions=1, ) failed_fence = DuplexFence("sid-open-rollback") await plane.handle( OpenDuplexSessionMessage( control_id="open-fails-during-stage-setup", fence=failed_fence, session_id=failed_fence.session_id, capabilities={"input_modes": ["append_audio_chunk"]}, ) ) failed_result = await result_sink.get() assert failed_result.ok is False assert plane.sessions.get(failed_fence.session_id) is not None replacement_fence = DuplexFence("sid-open-rollback-replacement") await plane.handle( OpenDuplexSessionMessage( control_id="open-replacement-before-rollback", fence=replacement_fence, session_id=replacement_fence.session_id, capabilities={"input_modes": ["append_audio_chunk"]}, ) ) assert (await result_sink.get()).error.code == "resource_exhausted" assert await plane.reap_expired() == 0 assert plane.sessions.get(failed_fence.session_id) is None await plane.handle( OpenDuplexSessionMessage( control_id="open-replacement-after-rollback", fence=replacement_fence, session_id=replacement_fence.session_id, capabilities={"input_modes": ["append_audio_chunk"]}, ) ) assert (await result_sink.get()).ok is True @pytest.mark.asyncio async def test_expired_cleanup_failure_does_not_block_other_sessions() -> None: class _OneStuckExpiryStagePort(_TypedStagePort): async def cleanup(self, request_ids: list[str], *, abort: bool = False) -> None: if request_ids == [stuck_request_id]: raise RuntimeError("stuck expiry cleanup") await super().cleanup(request_ids, abort=abort) clock = _Clock() stage_port = _OneStuckExpiryStagePort() result_sink: asyncio.Queue = asyncio.Queue() lifecycle_sink: asyncio.Queue = asyncio.Queue() plane = DuplexControlPlane( extension=None, stage_port=stage_port, result_sink=result_sink, lifecycle_sink=lifecycle_sink, lease_config=DuplexLeaseConfig(idle_ttl_s=1.0, disconnect_grace_s=1.0), clock=clock, ) stuck_fence = DuplexFence("sid-expiry-stuck") ready_fence = DuplexFence("sid-expiry-ready") stuck_request_id = plane.stage_request_id(stuck_fence, stage_id=0) ready_request_id = plane.stage_request_id(ready_fence, stage_id=0) for fence in (stuck_fence, ready_fence): await plane.handle( OpenDuplexSessionMessage( control_id=f"open-{fence.session_id}", fence=fence, session_id=fence.session_id, capabilities={}, ) ) assert (await result_sink.get()).ok is True plane.sessions.require(fence.session_id).reserve_stage_request( 0, plane.stage_request_id(fence, stage_id=0), fence=fence, ) clock.advance(2.0) assert await plane.reap_expired() == 1 assert (await lifecycle_sink.get()).session_id == ready_fence.session_id assert stage_port.cleanup_calls == [([ready_request_id], False)] @pytest.mark.asyncio async def test_cancel_retains_request_resources_until_abort_succeeds() -> None: class _FailOnceStagePort(_TypedStagePort): def __init__(self) -> None: super().__init__() self.failures_remaining = 1 async def cleanup(self, request_ids: list[str], *, abort: bool = False) -> None: if self.failures_remaining: self.failures_remaining -= 1 raise RuntimeError("transient abort failure") await super().cleanup(request_ids, abort=abort) stage_port = _FailOnceStagePort() result_sink: asyncio.Queue = asyncio.Queue() plane = DuplexControlPlane(extension=None, stage_port=stage_port, result_sink=result_sink) cancelled_fence = DuplexFence("sid-cancel-retry") next_fence = DuplexFence("sid-cancel-retry", epoch=1) session = plane.sessions.open_session(cancelled_fence) session.bind_stage_request(0, "req-cancel-retry", fence=cancelled_fence) def signal(control_id: str) -> SignalDuplexTurnMessage: return SignalDuplexTurnMessage( control_id=control_id, fence=cancelled_fence, next_fence=next_fence, session_id=cancelled_fence.session_id, event="input.cancel", ) await plane.handle(signal("cancel-first")) assert (await result_sink.get()).ok is False assert session.fence == next_fence assert session.resource_request_ids(cancelled_fence) == ["req-cancel-retry"] await plane.handle(signal("cancel-retry")) assert (await result_sink.get()).ok is True assert session.resource_request_ids(cancelled_fence) == [] assert stage_port.cleanup_calls == [(["req-cancel-retry"], True)] @pytest.mark.asyncio async def test_close_retains_session_resources_until_cleanup_succeeds() -> None: class _FailOnceStagePort(_TypedStagePort): def __init__(self) -> None: super().__init__() self.failures_remaining = 1 async def cleanup(self, request_ids: list[str], *, abort: bool = False) -> None: if self.failures_remaining: self.failures_remaining -= 1 raise RuntimeError("transient close cleanup failure") await super().cleanup(request_ids, abort=abort) stage_port = _FailOnceStagePort() result_sink: asyncio.Queue = asyncio.Queue() plane = DuplexControlPlane(extension=None, stage_port=stage_port, result_sink=result_sink) fence = DuplexFence("sid-close-retry") session = plane.sessions.open_session(fence) session.bind_stage_request(0, "req-close-retry", fence=fence) def close(control_id: str) -> CloseDuplexSessionMessage: return CloseDuplexSessionMessage( control_id=control_id, fence=fence, session_id=fence.session_id, reason="client_close", ) await plane.handle(close("close-first")) assert (await result_sink.get()).ok is False assert plane.sessions.get(fence.session_id) is session assert session.resource_request_ids() == ["req-close-retry"] await plane.handle(close("close-retry")) assert (await result_sink.get()).ok is True assert plane.sessions.get(fence.session_id) is None assert stage_port.cleanup_calls == [(["req-close-retry"], True)] @pytest.mark.asyncio async def test_control_cleanup_failure_does_not_block_other_sessions() -> None: class _IndependentFailureStagePort(_TypedStagePort): def __init__(self) -> None: super().__init__() self.attempts: dict[str, int] = {} async def cleanup(self, request_ids: list[str], *, abort: bool = False) -> None: request_id = request_ids[0] self.attempts[request_id] = self.attempts.get(request_id, 0) + 1 if request_id == "req-stuck" or self.attempts[request_id] == 1: raise RuntimeError(f"cleanup failed for {request_id}") await super().cleanup(request_ids, abort=abort) stage_port = _IndependentFailureStagePort() result_sink: asyncio.Queue = asyncio.Queue() plane = DuplexControlPlane(extension=None, stage_port=stage_port, result_sink=result_sink) for session_id, request_id in (("sid-stuck", "req-stuck"), ("sid-ready", "req-ready")): fence = DuplexFence(session_id) session = plane.sessions.open_session(fence) session.bind_stage_request(0, request_id, fence=fence) await plane.handle( CloseDuplexSessionMessage( control_id=f"close-{session_id}", fence=fence, session_id=session_id, reason="client_close", ) ) assert (await result_sink.get()).ok is False await plane.reap_expired() assert plane.sessions.get("sid-stuck") is not None assert plane.sessions.get("sid-ready") is None assert stage_port.cleanup_calls == [(["req-ready"], True)] @pytest.mark.asyncio async def test_control_cleanup_is_single_flight_with_concurrent_reaper() -> None: class _SlowCleanupStagePort(_TypedStagePort): def __init__(self) -> None: super().__init__() self.cleanup_started = asyncio.Event() self.release_cleanup = asyncio.Event() self.cleanup_attempts = 0 async def cleanup(self, request_ids: list[str], *, abort: bool = False) -> None: self.cleanup_attempts += 1 self.cleanup_started.set() await self.release_cleanup.wait() await super().cleanup(request_ids, abort=abort) stage_port = _SlowCleanupStagePort() result_sink: asyncio.Queue = asyncio.Queue() plane = DuplexControlPlane(extension=None, stage_port=stage_port, result_sink=result_sink) fence = DuplexFence("sid-concurrent-cleanup") session = plane.sessions.open_session(fence) session.bind_stage_request(0, "req-concurrent-cleanup", fence=fence) close_task = asyncio.create_task( plane.handle( CloseDuplexSessionMessage( control_id="close-concurrent-cleanup", fence=fence, session_id=fence.session_id, reason="client_close", ) ) ) await asyncio.wait_for(stage_port.cleanup_started.wait(), timeout=1) reaper_task = asyncio.create_task(plane.reap_expired()) await asyncio.sleep(0) await asyncio.sleep(0) assert stage_port.cleanup_attempts == 1 stage_port.release_cleanup.set() await asyncio.gather(close_task, reaper_task) assert (await result_sink.get()).ok is True assert plane.sessions.get(fence.session_id) is None assert stage_port.cleanup_calls == [(["req-concurrent-cleanup"], True)] @pytest.mark.asyncio async def test_unknown_input_mode_rejects_open_without_registering_session() -> None: result_sink: asyncio.Queue = asyncio.Queue() plane = DuplexControlPlane(extension=None, stage_port=_TypedStagePort(), result_sink=result_sink) fence = DuplexFence("sid-unknown-mode") await plane.handle( OpenDuplexSessionMessage( control_id="open-unknown-mode", fence=fence, session_id=fence.session_id, capabilities={"input_modes": ["future_mode"]}, ) ) result = await result_sink.get() assert result.ok is False assert result.error is not None assert result.error.code == "invalid_capability" assert plane.sessions.get(fence.session_id) is None @pytest.mark.asyncio async def test_extension_free_turn_commit_does_not_allocate_or_submit_stage_request() -> None: stage_port = _TypedStagePort() result_sink: asyncio.Queue = asyncio.Queue() plane = DuplexControlPlane(extension=None, stage_port=stage_port, result_sink=result_sink) fence = DuplexFence("sid-turn-commit") await plane.handle( OpenDuplexSessionMessage( control_id="open-turn-commit", fence=fence, session_id=fence.session_id, capabilities={"input_modes": [DuplexInputMode.TURN_COMMIT_ONLY.value]}, ) ) assert (await result_sink.get()).ok is True await plane.handle( AppendDuplexInputMessage( control_id="append-turn-commit", fence=fence, session_id=fence.session_id, mode=DuplexInputMode.TURN_COMMIT_ONLY.value, payload=None, final=True, ) ) result = await result_sink.get() assert result.ok is True assert result.stage_results[0]["result"]["data_plane_append"] is False assert stage_port.ensure_calls == [] assert stage_port.submit_calls == [] @pytest.mark.asyncio async def test_session_update_replaces_both_configs_with_one_generation() -> None: result_sink: asyncio.Queue = asyncio.Queue() plane = DuplexControlPlane(extension=_Extension(), stage_port=_TypedStagePort(), result_sink=result_sink) fence = DuplexFence("sid-atomic-config") session = plane.sessions.open_session(fence) await plane.handle( SignalDuplexTurnMessage( control_id="update-config", fence=fence, session_id=fence.session_id, event="session.update", session_config={"voice": "new"}, runtime_config={"temperature": 0.5}, ) ) assert (await result_sink.get()).ok is True assert session.config_generation == 1 assert session.session_config == {"voice": "new"} assert session.runtime_config == {"temperature": 0.5} @pytest.mark.asyncio async def test_stage_submission_is_compensated_when_local_commit_fails() -> None: stage_port = _TypedStagePort() result_sink: asyncio.Queue = asyncio.Queue() plane = DuplexControlPlane(extension=_Extension(), stage_port=stage_port, result_sink=result_sink) fence = DuplexFence("sid-submit-compensation") session = plane.sessions.open_session( fence, capabilities=DuplexRuntimeCapabilities(input_modes={DuplexInputMode.APPEND_AUDIO_CHUNK}), ) request_id = plane.stage_request_id(fence, stage_id=0) def fail_commit(_reservation): raise RuntimeError("local commit failed") session.commit_append = fail_commit # type: ignore[method-assign] await plane.handle( AppendDuplexInputMessage( control_id="append-compensation", fence=fence, session_id=fence.session_id, mode=DuplexInputMode.APPEND_AUDIO_CHUNK.value, payload={"audio": b"pcm"}, ) ) assert (await result_sink.get()).ok is False assert stage_port.cleanup_calls == [([request_id], True)] @pytest.mark.asyncio async def test_failed_submission_compensation_blocks_append_until_reaper_cleans_it() -> None: class _FailingCleanupStagePort(_TypedStagePort): def __init__(self) -> None: super().__init__() self.cleanup_failures = 2 async def cleanup(self, request_ids: list[str], *, abort: bool = False) -> None: await super().cleanup(request_ids, abort=abort) if self.cleanup_failures > 0: self.cleanup_failures -= 1 raise RuntimeError("stage abort failed") stage_port = _FailingCleanupStagePort() result_sink: asyncio.Queue = asyncio.Queue() plane = DuplexControlPlane(extension=_Extension(), stage_port=stage_port, result_sink=result_sink) fence = DuplexFence("sid-durable-submit-compensation") session = plane.sessions.open_session( fence, capabilities=DuplexRuntimeCapabilities(input_modes={DuplexInputMode.APPEND_AUDIO_CHUNK}), ) original_commit = session.commit_append def fail_commit(_reservation): raise RuntimeError("local commit failed") session.commit_append = fail_commit # type: ignore[method-assign] first = AppendDuplexInputMessage( control_id="append-compensation-1", fence=fence, session_id=fence.session_id, mode=DuplexInputMode.APPEND_AUDIO_CHUNK.value, payload={"audio": b"pcm"}, ) await plane.handle(first) assert (await result_sink.get()).ok is False assert plane.pending_submission_cleanup_count == 1 assert len(stage_port.submit_calls) == 1 await plane.handle( AppendDuplexInputMessage( control_id="append-compensation-2", fence=fence, session_id=fence.session_id, mode=DuplexInputMode.APPEND_AUDIO_CHUNK.value, payload={"audio": b"pcm"}, ) ) assert (await result_sink.get()).ok is False assert plane.pending_submission_cleanup_count == 1 assert len(stage_port.submit_calls) == 1 session.commit_append = original_commit # type: ignore[method-assign] stage_port.cleanup_failures = 0 await plane.reap_expired() assert plane.pending_submission_cleanup_count == 0 assert session.resource_request_ids() == [] def test_stale_output_is_rejected_before_extension_decision() -> None: class _CountingExtension(_Extension): def __init__(self) -> None: self.decision_calls = 0 def decide_output(self, **kwargs): del kwargs self.decision_calls += 1 return None extension = _CountingExtension() plane = DuplexControlPlane(extension=extension, stage_port=_TypedStagePort(), result_sink=asyncio.Queue()) current = DuplexFence("sid-stale-output", epoch=1) plane.sessions.open_session(current) context = DuplexOutputContext( identity=DuplexRequestIdentity( session_id=current.session_id, fence=DuplexFence(current.session_id, epoch=0), ), final_stage_id=1, segment_finished=True, ) assert plane.decide_output(0, object(), context) is None assert extension.decision_calls == 0 @pytest.mark.asyncio async def test_control_dispatch_is_ordered_per_session_without_blocking_other_sessions() -> None: class _BlockingStagePort(_TypedStagePort): def __init__(self) -> None: super().__init__() self.started = asyncio.Event() self.release = asyncio.Event() async def submit(self, submission: DuplexStageSubmission) -> DuplexStageSubmissionResult: if submission.context.session_id == "sid-blocked": self.started.set() await self.release.wait() return await super().submit(submission) stage_port = _BlockingStagePort() result_sink: asyncio.Queue = asyncio.Queue() plane = DuplexControlPlane(extension=_Extension(), stage_port=stage_port, result_sink=result_sink) blocked = DuplexFence("sid-blocked") independent = DuplexFence("sid-independent") capabilities = DuplexRuntimeCapabilities(input_modes={DuplexInputMode.APPEND_AUDIO_CHUNK}) plane.sessions.open_session(blocked, capabilities=capabilities) plane.sessions.open_session(independent, capabilities=capabilities) plane.dispatch( AppendDuplexInputMessage( control_id="blocked-append", fence=blocked, session_id=blocked.session_id, mode=DuplexInputMode.APPEND_AUDIO_CHUNK.value, payload={"audio": b"pcm"}, operation_id="blocked-operation", ) ) await asyncio.wait_for(stage_port.started.wait(), timeout=1) plane.dispatch( AppendDuplexInputMessage( control_id="blocked-append-retry", fence=blocked, session_id=blocked.session_id, mode=DuplexInputMode.APPEND_AUDIO_CHUNK.value, payload={"audio": b"pcm"}, operation_id="blocked-operation", ) ) plane.dispatch( TouchDuplexSessionMessage( control_id="blocked-touch", fence=blocked, session_id=blocked.session_id, activity=DuplexLeaseActivity.HEARTBEAT.value, ) ) plane.dispatch( TouchDuplexSessionMessage( control_id="independent-touch", fence=independent, session_id=independent.session_id, activity=DuplexLeaseActivity.HEARTBEAT.value, ) ) independent_result = await asyncio.wait_for(result_sink.get(), timeout=1) assert independent_result.control_id == "independent-touch" stage_port.release.set() await plane.drain() assert [result_sink.get_nowait().control_id for _ in range(3)] == [ "blocked-append", "blocked-append-retry", "blocked-touch", ] assert [ submission.context.session_id for submission in stage_port.submit_calls if submission.context.session_id == blocked.session_id ] == [blocked.session_id]