From 9e7e3a802cee2b09d8ce88ed3d2cad32a414dcba Mon Sep 17 00:00:00 2001 From: chickeyton Date: Fri, 25 Sep 2026 10:23:08 +0800 Subject: [PATCH] [Bugfix] Compensate cancelled duplex open, append and resume; fix serving README drift (#7636 Issues 1, 2, 3, 21) (#7785) Signed-off-by: chickeyton Signed-off-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com> Co-authored-by: Claude Fable 5.1 Co-authored-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com> --- docs/design/fullduplex.md | 9 +- docs/serving/README.md | 27 +- tests/engine/duplex/test_duplex_lease.py | 26 ++ tests/engine/duplex/test_session_manager.py | 198 +++++++- tests/engine/duplex/test_session_runner.py | 140 +++++- .../entrypoints/duplex/test_duplex_serving.py | 425 +++++++++++++++++- .../openai/test_duplex_session_attachment.py | 87 ++++ tests/entrypoints/test_duplex_omni.py | 178 +++++++- vllm_omni/engine/duplex/messages.py | 5 + .../engine/duplex/session/append_task.py | 70 ++- vllm_omni/engine/duplex/session/context.py | 23 +- .../engine/duplex/session/engine_session.py | 4 +- vllm_omni/engine/duplex/session/lease.py | 14 +- vllm_omni/engine/duplex/session/manager.py | 86 +++- vllm_omni/engine/duplex/session/runner.py | 1 + vllm_omni/engine/duplex_omni_engine.py | 17 +- vllm_omni/entrypoints/duplex/serving.py | 191 +++++++- .../entrypoints/duplex/session_attachment.py | 129 +++++- vllm_omni/entrypoints/duplex_omni.py | 163 ++++++- 19 files changed, 1697 insertions(+), 96 deletions(-) diff --git a/docs/design/fullduplex.md b/docs/design/fullduplex.md index df698880e..2733ecf08 100644 --- a/docs/design/fullduplex.md +++ b/docs/design/fullduplex.md @@ -283,8 +283,13 @@ command handle.submit(DuplexCommand) -> DuplexSessionManager.dispatch: unknown_session / input_backpressure checks, then runner mailbox output DuplexOrchestrator._intercept_stage_output -> runner.on_stage_output -> mailbox -> typed events -> output_sink (DuplexSessionEventMessage) -> DuplexOmni._route_engine_message -> handle.events() -detach DuplexOmni.detach_session -> touch(DETACH): engine-owned disconnect grace; expiry -> SessionExpired -resume DuplexOmni.resume_session(expected_lease_generation) -> lease CAS; the existing handle is re-entered +detach DuplexOmni.detach_session(expected_lease_generation) -> touch(DETACH): engine-owned disconnect + grace, refused for a lease newer than the one the caller held; expiry -> SessionExpired +resume DuplexOmni.resume_session(expected_lease_generation) -> lease CAS keyed by the control id (a replay + answers with the generation it produced); the existing handle is re-entered. A caller cancelled + mid-RPC observes the outcome afterwards (replaying until the engine answers) and settles a landed + resume: the generation goes to the resume waiting to activate, else to the connection still + attached, else the lease is detached again close DuplexOmni.close_session -> close RPC; the manager tears the runner down, then the stage cleanup (abort submitted requests, release reserved ids), then SessionClosed, then the RPC result. SessionClosed is emitted after the cleanup attempt (in a finally), so it is also sent when the diff --git a/docs/serving/README.md b/docs/serving/README.md index 472ab2a49..06ab24637 100644 --- a/docs/serving/README.md +++ b/docs/serving/README.md @@ -150,17 +150,22 @@ interchangeably. | Robot policy inference | `WS /v1/realtime/robot/openpi` | Send MessagePack observations and receive action arrays | [OpenPI Robot Policy](openpi_api.md) | All six routes are model- or configuration-dependent. In particular, -`/v1/realtime` is full duplex when the model is a duplex model (its pipeline -declares a `duplex_plugin` and the deploy configuration sets -`session_mode: duplex`); a stock Realtime client needs no vendor query -parameter, and `duplex=0` is the explicit opt-out. Such a server serves the -websocket route plus `POST /v1/chat/completions`, and no other turn-based -HTTP route. On a duplex server that chat route is not the turn-based path: -each request runs on a short-lived duplex session, so it holds one of -`duplex_session.max_sessions` for its lifetime and answers at the model's -real-time pace -- see [Full Duplex](full_duplex_api.md). Clients should -also verify the duplex capability payload because the query-parameter form -falls back to the ordinary realtime handler when duplex is unavailable. +`/v1/realtime` is full duplex when the server is a duplex server: the model's +pipeline declares a `duplex_plugin` and its deploy configuration sets +`session_mode: duplex` (with `session_mode: turn` the same model boots the +ordinary turn-based serving stack instead). On a duplex server a stock +Realtime client needs no vendor query parameter; `duplex=0` selects the +turn-based Realtime handler, which a duplex server does not mount, so that +connection is refused with `Realtime API is not available`. Such a server +serves the websocket route plus `POST /v1/chat/completions`, `/v1/models` +and `/health`, and no other turn-based HTTP route. The chat route is the +ordinary chat service running on the duplex engine: a request is a +turn-based generation on the same stages, served alongside the live +websocket sessions; it opens no duplex session and holds no +`duplex_session.max_sessions` slot -- see [Full Duplex](full_duplex_api.md). +On a server that is not duplex, an explicit `?duplex=1` is refused rather +than answered by the turn-based handler, so a client that requested duplex +never silently gets the other protocol. ## Related Endpoints diff --git a/tests/engine/duplex/test_duplex_lease.py b/tests/engine/duplex/test_duplex_lease.py index bd26b5740..326ea9309 100644 --- a/tests/engine/duplex/test_duplex_lease.py +++ b/tests/engine/duplex/test_duplex_lease.py @@ -85,6 +85,32 @@ def test_detach_grace_and_resume_advance_the_lease_generation() -> None: session.resume_lease(expected_lease_generation=0) +def test_a_resume_replayed_under_its_control_id_answers_without_resuming_again() -> None: + """A caller that lost the answer (cancelled or timed-out waiter) asks again with the same id. + + The engine must tell it the generation that resume produced, not bump the + lease a second time, and must still refuse any other resume against the + old generation. + """ + clock = FakeMonotonicClock(10.0) + session = _session("sid-replay", clock) + session.detach_lease() + + assert session.resume_lease(expected_lease_generation=0, control_id="rpc-a") == 1 + clock.advance(5.0) + session.detach_lease() + + assert session.resume_lease(expected_lease_generation=0, control_id="rpc-a") == 1 + assert session.lease_generation == 1 + assert session.lease.detached_at == 15.0, "a replay is an answer, not a new resume" + with pytest.raises(ValueError, match="lease generation mismatch"): + session.resume_lease(expected_lease_generation=0, control_id="rpc-b") + + assert session.resume_lease(expected_lease_generation=1, control_id="rpc-b") == 2 + with pytest.raises(ValueError, match="lease generation mismatch"): + session.resume_lease(expected_lease_generation=0, control_id="rpc-a") + + def test_active_operation_prevents_mid_transaction_expiry() -> None: clock = FakeMonotonicClock(0.0) session = _session("sid-operation", clock, idle_ttl_s=5.0) diff --git a/tests/engine/duplex/test_session_manager.py b/tests/engine/duplex/test_session_manager.py index 2f44a9713..5fe07d6bf 100644 --- a/tests/engine/duplex/test_session_manager.py +++ b/tests/engine/duplex/test_session_manager.py @@ -296,6 +296,7 @@ class Harness: clock: FakeClock output_sink: asyncio.Queue = field(default_factory=asyncio.Queue) result_sink: asyncio.Queue = field(default_factory=asyncio.Queue) + _control_seq: int = 0 @classmethod def create( @@ -344,20 +345,30 @@ class Harness: ) return await self.result() - async def resume(self, session_id: str, *, expected_lease_generation: int) -> DuplexControlResultMessage: + async def resume( + self, session_id: str, *, expected_lease_generation: int, control_id: str | None = None + ) -> DuplexControlResultMessage: + # Distinct by default: the engine answers a repeated control id as a + # replay of the same resume, which only the replay test wants. + self._control_seq += 1 await self.manager.handle( ResumeDuplexSessionMessage( - control_id=f"resume-{session_id}-{expected_lease_generation}", + control_id=control_id or f"resume-{session_id}-{expected_lease_generation}-{self._control_seq}", session_id=session_id, expected_lease_generation=expected_lease_generation, ) ) return await self.result() - async def touch(self, session_id: str, activity: str) -> DuplexControlResultMessage: + async def touch( + self, session_id: str, activity: str, *, expected_lease_generation: int | None = None + ) -> DuplexControlResultMessage: await self.manager.handle( TouchDuplexSessionMessage( - control_id=f"touch-{session_id}-{activity}", session_id=session_id, activity=activity + control_id=f"touch-{session_id}-{activity}-{expected_lease_generation}", + session_id=session_id, + activity=activity, + expected_lease_generation=expected_lease_generation, ) ) return await self.result() @@ -1656,6 +1667,168 @@ async def test_reaper_loop_survives_one_cleanup_failure(first_cleanup_delay: flo await asyncio.wait_for(task, timeout=5.0) +# --------------------------------------------------------------------------- # +# Cancelled open (#7636 Issue 1) # +# --------------------------------------------------------------------------- # + + +class _BlockingResultSink(asyncio.Queue): + """A result sink whose ``put`` parks, so an open can be cancelled after it admitted the session.""" + + def __init__(self) -> None: + super().__init__() + self.entered = asyncio.Event() + self.gate = asyncio.Event() + + async def put(self, item: object) -> None: + self.entered.set() + await self.gate.wait() + await super().put(item) + + +async def test_an_open_cancelled_while_awaiting_the_plugin_frees_the_admission_slot() -> None: + """``CancelledError`` is not an ``Exception``: the rollback used to be skipped. + + Admission counts ``runners | closing | admitting``, so every open cancelled + mid-flight (engine teardown, task cancellation) burned one slot for good; + after ``max_sessions`` of them every new session was refused. + """ + async with Harness.create(max_sessions=1) as harness: + harness.plugin.blocked_session_ids.add("blocked") + config = DuplexSessionConfig(model="fake-model", instructions="blocked") + open_task = asyncio.create_task(harness.open("sid-cancelled", config)) + await asyncio.wait_for(harness.plugin.runtime_config_started.wait(), timeout=1.0) + + open_task.cancel() + with pytest.raises(asyncio.CancelledError): + await open_task + + assert harness.manager.get("sid-cancelled") is None + assert harness.manager.active_count() == 0 + assert harness.result_sink.empty(), "a cancelled open answers nobody" + assert (await harness.open("sid-replacement")).ok is True + + +async def test_an_open_cancelled_after_admission_releases_the_runner_and_its_stage_reservation() -> None: + """The worst landing spot for the cancel: the runner is registered and Stage0 is reserved. + + Both have to be undone, or the runner stays in ``runners`` (one slot gone) + and the Stage0 request stays reserved in the orchestrator. + """ + async with Harness.create(max_sessions=1) as harness: + sink = _BlockingResultSink() + original_sink = harness.manager._result_sink + harness.manager._result_sink = sink + open_task = asyncio.create_task( + harness.manager.handle( + OpenDuplexSessionMessage( + control_id="open-cancelled", + session_id="sid-cancelled", + session_config=DuplexSessionConfig(model="fake-model"), + ) + ) + ) + await asyncio.wait_for(sink.entered.wait(), timeout=1.0) + assert "sid-cancelled" in harness.manager.runners + assert [context.request_id for context in harness.stage_port.ensure_calls] == [ + stage0_request_id("sid-cancelled") + ] + + open_task.cancel() + with pytest.raises(asyncio.CancelledError): + await open_task + harness.manager._result_sink = original_sink + + assert "sid-cancelled" not in harness.manager.runners + assert harness.manager.active_count() == 0 + assert harness.stage_port.cleanup_calls == [([stage0_request_id("sid-cancelled")], False)] + assert stage0_request_id("sid-cancelled") not in harness.manager._request_index + assert (await harness.open("sid-replacement")).ok is True + + +async def test_an_open_cancelled_again_during_its_rollback_leaves_the_stage_cleanup_to_the_reaper() -> None: + """A second cancellation landing in the rollback's awaits must not orphan the Stage0 request. + + The runner, the reservations and the request index are already gone by + then, so nothing would ever clean the orchestrator's request state up. + The rollback records the ids as a pending request cleanup before it + awaits anything, and the reaper's retry finishes the job. + """ + async with Harness.create(max_sessions=1) as harness: + sink = _BlockingResultSink() + original_sink = harness.manager._result_sink + harness.manager._result_sink = sink + open_task = asyncio.create_task( + harness.manager.handle( + OpenDuplexSessionMessage( + control_id="open-cancelled-twice", + session_id="sid-cancelled", + session_config=DuplexSessionConfig(model="fake-model"), + ) + ) + ) + await asyncio.wait_for(sink.entered.wait(), timeout=1.0) + runner = harness.manager.runners["sid-cancelled"] + shutdown_started = asyncio.Event() + + async def parked_shutdown() -> None: + shutdown_started.set() + await asyncio.Event().wait() + + runner.shutdown = parked_shutdown # type: ignore[method-assign] + + open_task.cancel() + await asyncio.wait_for(shutdown_started.wait(), timeout=1.0) + open_task.cancel() # lands inside the rollback, before the stage cleanup + with pytest.raises(asyncio.CancelledError): + await open_task + harness.manager._result_sink = original_sink + + assert "sid-cancelled" not in harness.manager.runners + assert harness.manager.active_count() == 0 + assert harness.stage_port.cleanup_calls == [], "the cleanup await was never reached" + key = ("sid-cancelled", 0) + assert key in harness.manager._pending_request_cleanups + + await harness.manager.reap_expired() + assert harness.stage_port.cleanup_calls == [([stage0_request_id("sid-cancelled")], False)] + assert key not in harness.manager._pending_request_cleanups + assert (await harness.open("sid-replacement")).ok is True + + +# --------------------------------------------------------------------------- # +# Detach fenced on the lease generation # +# --------------------------------------------------------------------------- # + + +async def test_detach_is_refused_for_a_lease_the_caller_no_longer_holds() -> None: + """A connection giving up its lease must not start the grace for the lease a later resume owns. + + Resume is a CAS on the lease generation; detach is now fenced the same + way, so a stale caller (its resume was superseded by another connection's) + is refused instead of detaching the winner. + """ + async with Harness.create() as harness: + await harness.open("sid-fence") + session = harness.session("sid-fence") + assert (await harness.resume("sid-fence", expected_lease_generation=0)).ok is True + assert session.lease_generation == 1 + + stale = await harness.touch("sid-fence", "detach", expected_lease_generation=0) + assert stale.ok is False + assert stale.error_code == "session_resume_conflict" + assert session.lease.detached_at is None, "the current lease is untouched" + + current = await harness.touch("sid-fence", "detach", expected_lease_generation=1) + assert current.ok is True + assert session.lease.detached_at is not None + + # An unfenced detach keeps its meaning: whatever generation is current. + assert (await harness.resume("sid-fence", expected_lease_generation=1)).ok is True + assert (await harness.touch("sid-fence", "detach")).ok is True + assert session.lease.detached_at is not None + + # --------------------------------------------------------------------------- # # Cancellation waits # # --------------------------------------------------------------------------- # @@ -1685,3 +1858,20 @@ async def test_cancel_append_tasks_absorbs_a_task_that_outlives_the_wait() -> No assert await tasks.cancel_append_tasks(timeout_s=0.05) is True assert tasks.append_tail is None await asyncio.wait_for(asyncio.gather(task, return_exceptions=True), timeout=5.0) + + +async def test_a_resume_replayed_under_its_control_id_reports_the_generation_it_produced() -> None: + """The client side of an abandoned resume: it lost the answer and asks again with the same id.""" + async with Harness.create() as harness: + await harness.open("sid-replay") + session = harness.session("sid-replay") + first = await harness.resume("sid-replay", expected_lease_generation=0, control_id="rpc-a") + assert first.ok is True and first.lease_generation == 1 + + replay = await harness.resume("sid-replay", expected_lease_generation=0, control_id="rpc-a") + assert replay.ok is True + assert replay.lease_generation == 1 + assert session.lease_generation == 1, "a replay does not resume again" + + other = await harness.resume("sid-replay", expected_lease_generation=0, control_id="rpc-b") + assert other.ok is False and other.error_code == "session_resume_conflict" diff --git a/tests/engine/duplex/test_session_runner.py b/tests/engine/duplex/test_session_runner.py index 92fba2793..c51fc243c 100644 --- a/tests/engine/duplex/test_session_runner.py +++ b/tests/engine/duplex/test_session_runner.py @@ -73,6 +73,9 @@ class RecordingStagePort(DuplexStagePort): self.cleanups: list[tuple[list[str], bool]] = [] self.aborts: list[list[str]] = [] self.fail_submit: Exception | None = None + #: When set, ``submit`` parks on it after signalling ``submit_started``. + self.submit_gate: asyncio.Event | None = None + self.submit_started = asyncio.Event() @property def stage_count(self) -> int: @@ -87,6 +90,9 @@ class RecordingStagePort(DuplexStagePort): async def submit(self, submission: DuplexStageSubmission) -> DuplexStageSubmissionResult: if self.fail_submit is not None: raise self.fail_submit + if self.submit_gate is not None: + self.submit_started.set() + await self.submit_gate.wait() self.submissions.append(submission) return DuplexStageSubmissionResult( request_id=submission.context.request_id, @@ -163,7 +169,7 @@ class Harness: def deliver( self, - output: object, + output: SimpleNamespace, *, stage_id: int = 1, segment_finished: bool = False, @@ -186,7 +192,7 @@ class Harness: raise AssertionError("stage output is missing request_id") return self.runner.on_stage_output(stage_id, output, metrics, request_id=request_id, context=context) - async def deliver_and_settle(self, output: object, **kwargs: Any) -> list[DuplexEvent]: + async def deliver_and_settle(self, output: SimpleNamespace, **kwargs: Any) -> list[DuplexEvent]: self.deliver(output, **kwargs) return await self.settle() @@ -1624,3 +1630,133 @@ async def test_server_vad_speech_stopped_still_commits_a_turn_mode_session() -> assert len(_final_submissions(h)) == 1, "the detector's stop commits the turn and starts the response" finally: await close_harness(h) + + +# --------------------------------------------------------------------------- # +# Cancelled append and its precreated response (#7636 Issue 2) # +# --------------------------------------------------------------------------- # + + +async def _commit_with_pending_response(h: Harness) -> str: + """Commit a turn whose final append parks in the stage port; return the precreated response id.""" + await h.run(append_audio()) + h.port.submit_gate = asyncio.Event() + h.submit(commands.Commit(create_response=True)) + await asyncio.wait_for(h.port.submit_started.wait(), timeout=2.0) + # The parked append keeps the runner "busy": settle only until the mailbox is quiet. + events = await h.settle(timeout_s=0.3) + assert "response.created" in types(events) + response_id = h.session.active_response_id + assert response_id is not None + assert h.runner.tasks.append_tail is not None and not h.runner.tasks.append_tail.done() + return response_id + + +@pytest.mark.asyncio +async def test_an_append_cancelled_from_outside_fails_the_response_it_precreated() -> None: + """A cancelled append gives back the response it precreated, as its docstring promises. + + The ``CancelledError`` branch used to roll back only the PCM reservation, + so ``active_response_id`` stayed pinned to a response nobody would ever + finish: the client had seen ``response.created`` and waited forever. + """ + h = await open_harness(auto_response=False) + try: + response_id = await _commit_with_pending_response(h) + + pending = h.runner.tasks.append_tail + assert pending is not None + pending.cancel() + events = await h.settle() + + done = find(events, "response.done") + assert done.response_id == response_id + assert done.status == "failed" + assert done.response["status_details"]["reason"] == "append_cancelled" + assert h.session.active_response_id is None + assert h.session.state == DuplexSessionState.OPEN + + # The cancelled append counts as a failed one: the chain behind it is + # refused with the documented error, not left hanging. + h.port.submit_gate = None + await h.run(append_audio()) + events = await h.run(commands.Commit(create_response=True)) + assert "response.created" not in types(events) + assert find(events, "error").code == "commit_aborted" + assert h.session.active_response_id is None + finally: + await close_harness(h) + + +@pytest.mark.asyncio +async def test_response_cancel_of_a_pending_append_ends_the_response_once_as_cancelled() -> None: + """When the runner cancels the append itself, it owns the response: one terminal, status cancelled. + + The append must not also fail the response on its way out, or the client + would get a failed ``response.done`` for a response it cancelled. + """ + h = await open_harness(auto_response=False) + try: + response_id = await _commit_with_pending_response(h) + + events = await h.run(commands.CancelResponse()) + + terminals = [event for event in events if event.type == "response.done"] + assert [event.response_id for event in terminals] == [response_id] + assert terminals[0].status == "cancelled" + assert h.session.active_response_id is None + assert h.session.state == DuplexSessionState.OPEN + finally: + await close_harness(h) + + +@pytest.mark.asyncio +async def test_closing_the_session_with_a_pending_append_emits_no_failed_response() -> None: + """A close ends the active response itself; the append it cancels stays quiet.""" + h = await open_harness(auto_response=False) + try: + await _commit_with_pending_response(h) + + events = await h.run(commands.CloseSession(reason="client_close")) + events += await h.settle() + + assert [event.type for event in events if event.type == "response.done"] == [] + assert "session.closed" in types(events) + assert h.session.active_response_id is None + finally: + await close_harness(h) + + +@pytest.mark.asyncio +async def test_input_clear_racing_a_queued_append_drops_it_as_the_clients_doing() -> None: + """An ``input_audio_buffer.clear`` that lands while an append waits its turn is not a runtime failure. + + The clear deactivates the queued append's reservation; when its turn + comes it gives everything back and stops, without an error and without + failing the session, and its reason is the client's clear rather than + ``runtime_append_failed``. + """ + h = await open_harness() + try: + h.port.submit_gate = asyncio.Event() + h.submit(append_audio()) + await asyncio.wait_for(h.port.submit_started.wait(), timeout=2.0) + h.submit(append_audio()) + await h.settle(timeout_s=0.3) + assert len(h.runner.tasks.append_tasks) == 2, "the second append is queued behind the parked one" + queued = h.runner.tasks.append_tail + assert queued is not None and not queued.done() + + h.submit(commands.ClearInput()) + events = await h.settle(timeout_s=0.3) + assert types(events) == ["input_audio_buffer.cleared"] + + h.port.submit_gate.set() + events = await h.settle() + assert queued.done() and queued.result() is False, "the cleared append never ran" + assert len(h.port.submissions) == 1, "only the append that was already in flight reached the stage" + assert [event.type for event in events if event.type in {"error", "response.done", "session.closed"}] == [] + assert h.session.pending_input_bytes == 0 + assert h.session.state == DuplexSessionState.OPEN + finally: + await close_harness(h) diff --git a/tests/entrypoints/duplex/test_duplex_serving.py b/tests/entrypoints/duplex/test_duplex_serving.py index dde616d94..70a0e9d48 100644 --- a/tests/entrypoints/duplex/test_duplex_serving.py +++ b/tests/entrypoints/duplex/test_duplex_serving.py @@ -8,6 +8,7 @@ from __future__ import annotations import asyncio import base64 import json +from collections.abc import Awaitable, Callable from contextlib import suppress from typing import Any @@ -131,7 +132,15 @@ class FakeOmni: self.handles: dict[str, FakeHandle] = {} self.resumed: list[tuple[str, int]] = [] self.detached: list[str] = [] + #: Detaches the engine refused because the caller's lease generation was stale. + self.detach_refused: list[tuple[str, int]] = [] self.open_error: DuplexSessionError | None = None + #: When set, ``resume_session`` parks on it once the engine applied the + #: resume: the lease generation is bumped, the result is still in flight. + self.resume_gate: asyncio.Event | None = None + self.resume_started = asyncio.Event() + #: Settlements of resumes whose caller was cancelled mid-RPC (DuplexOmni's contract). + self.compensations: list[asyncio.Task[None]] = [] async def open_session(self, config: Any) -> FakeHandle: self.opened.append(dict(config)) @@ -146,13 +155,46 @@ class FakeOmni: def get_session(self, session_id: str) -> FakeHandle | None: return self.handles.get(session_id) - async def resume_session(self, session_id: str, *, expected_lease_generation: int) -> FakeHandle: + async def resume_session( + self, + session_id: str, + *, + expected_lease_generation: int, + on_abandoned: Callable[[int], Awaitable[None]] | None = None, + ) -> FakeHandle: self.resumed.append((session_id, expected_lease_generation)) handle = self.handles[session_id] - handle.lease_generation = expected_lease_generation + 1 + lease_generation = expected_lease_generation + 1 + handle.lease_generation = lease_generation + if self.resume_gate is not None: + self.resume_started.set() + gate = self.resume_gate + try: + await gate.wait() + except asyncio.CancelledError: + # DuplexOmni's contract: the engine applied the resume; once + # its answer arrives the landed generation is settled off the + # cancelled task, through the caller's callback if it gave one. + async def settle() -> None: + await gate.wait() + if on_abandoned is not None: + await on_abandoned(lease_generation) + else: + with suppress(DuplexSessionError): + await self.detach_session(session_id, expected_lease_generation=lease_generation) + + self.compensations.append(asyncio.create_task(settle())) + raise return handle - async def detach_session(self, session_id: str) -> None: + async def detach_session(self, session_id: str, *, expected_lease_generation: int | None = None) -> None: + handle = self.handles[session_id] + if expected_lease_generation is not None and expected_lease_generation != handle.lease_generation: + # The engine's fence: a detach of a lease the caller no longer holds is refused. + self.detach_refused.append((session_id, expected_lease_generation)) + raise DuplexSessionError( + "duplex lease generation mismatch", code="session_resume_conflict", session_id=session_id + ) self.detached.append(session_id) @@ -607,3 +649,380 @@ async def test_a_resume_that_fails_to_activate_does_not_detach_the_live_attachme pending.cancel() with suppress(asyncio.CancelledError, Exception): await pending + + +# --------------------------------------------------------------------------- # +# Cancelled resume and the engine disconnect grace (#7636 Issue 3) # +# --------------------------------------------------------------------------- # + + +@pytest.mark.asyncio +async def test_a_resume_cancelled_during_activation_puts_the_engine_lease_back_into_grace() -> None: + """``_resume`` resumes the engine lease first; a cancel mid-activation must detach it again. + + The rollback used to catch ``Exception`` only. A handler task cancelled + while sending ``session.resumed`` or a replay entry left the engine with + ``detached_at=None`` and the outer handler with no attachment to clean up: + the reaper reclaimed nothing after the grace, and with one slot the next + open was refused with ``resource_exhausted``. + """ + omni = FakeOmni() + handler = _handler(omni) + ws, handle, task = await _open(handler, omni) + token = ws.sent[0]["resume_token"] + ws.disconnect() + await asyncio.wait_for(task, timeout=2.0) + assert omni.detached == [handle.session_id] + registry = handler._attachment_registry + + activating = asyncio.Event() + + async def parked_resume(*args, **kwargs): + activating.set() + await asyncio.Event().wait() + + original_resume = registry.resume + registry.resume = parked_resume # type: ignore[method-assign] + try: + ws2, task2 = await _resume(handler, handle.session_id, token) + await asyncio.wait_for(activating.wait(), timeout=2.0) + assert omni.resumed == [(handle.session_id, 0)], "the engine lease was resumed before activation" + + task2.cancel() + with suppress(asyncio.CancelledError): + await task2 + finally: + registry.resume = original_resume # type: ignore[method-assign] + + # The lease is detached again, so the disconnect grace runs for it. + assert omni.detached == [handle.session_id, handle.session_id] + assert not await registry.has_attachment(handle.session_id) + assert handle.close_reasons == [] + # And the session is still resumable with the token the cancelled attempt presented. + ws3, task3 = await _resume(handler, handle.session_id, token) + resumed = await ws3.wait_for("session.resumed") + assert resumed["session_id"] == handle.session_id + ws3.disconnect() + await asyncio.wait_for(task3, timeout=2.0) + + +@pytest.mark.asyncio +async def test_a_resume_cancelled_after_activation_detaches_the_attachment_it_made() -> None: + """A cancel after the registry activated the new attachment must undo that attachment too. + + Otherwise the session stays attached to a socket nobody serves, with the + engine lease resumed: no grace, no expiry, until the idle TTL. + """ + omni = FakeOmni() + handler = _handler(omni) + ws, handle, task = await _open(handler, omni) + token = ws.sent[0]["resume_token"] + registry = handler._attachment_registry + + # The takeover notifies the replaced socket after activation; park there. + notifying = asyncio.Event() + + async def parked_send(payload: dict[str, Any]) -> None: + notifying.set() + await asyncio.Event().wait() + + ws.send_json = parked_send # type: ignore[method-assign] + ws2, task2 = await _resume(handler, handle.session_id, token) + await asyncio.wait_for(notifying.wait(), timeout=2.0) + assert await registry.is_current_attachment(handle.session_id, 2), "the new socket is attached" + + task2.cancel() + with suppress(asyncio.CancelledError): + await task2 + + assert not await registry.has_attachment(handle.session_id) + assert omni.detached == [handle.session_id] + assert handle.close_reasons == [] + + task.cancel() + with suppress(asyncio.CancelledError, Exception): + await task + + +@pytest.mark.asyncio +async def test_an_abandoned_resume_never_detaches_the_lease_a_later_resume_owns() -> None: + """Two connections resuming the same session: the loser's rollback must not touch the winner's lease. + + Interleaving: A has activated and is notifying the socket it replaced; B, + holding A's rotated token, has completed its engine resume but its RPC + result has not reached the handler yet, so nothing is attached for B. + Cancelling A detaches A's registry generation, which used to make the + engine detach unconditional: it landed on the lease B had just resumed, + and the reaper expired B after the grace even though it heartbeats. The + engine now refuses a detach fenced on a generation it has moved past. + """ + omni = FakeOmni() + handler = _handler(omni) + ws, handle, task = await _open(handler, omni) + token = ws.sent[0]["resume_token"] + registry = handler._attachment_registry + + # A takes over and parks while notifying the replaced socket (after activation). + notifying = asyncio.Event() + + async def parked_send(payload: dict[str, Any]) -> None: + notifying.set() + await asyncio.Event().wait() + + ws.send_json = parked_send # type: ignore[method-assign] + ws_a, task_a = await _resume(handler, handle.session_id, token) + await asyncio.wait_for(notifying.wait(), timeout=2.0) + resumed_a = await ws_a.wait_for("session.resumed") + assert handle.lease_generation == 1 + + # B resumes with A's rotated token: the engine applies it, the result is still in flight. + omni.resume_gate = asyncio.Event() + ws_b, task_b = await _resume(handler, handle.session_id, resumed_a["resume_token"]) + await asyncio.wait_for(omni.resume_started.wait(), timeout=2.0) + assert handle.lease_generation == 2 + + task_a.cancel() + with suppress(asyncio.CancelledError): + await task_a + assert omni.detached == [], "A's rollback must not detach the lease B resumed" + assert omni.detach_refused == [(handle.session_id, 1)] + assert not await registry.has_attachment(handle.session_id), "A's attachment is gone" + + omni.resume_gate.set() + resumed_b = await ws_b.wait_for("session.resumed") + assert resumed_b["attachment_generation"] == 3 + assert await registry.is_current_attachment(handle.session_id, 3) + assert omni.detached == [] + + # B's own disconnect detaches the lease B holds. + ws_b.disconnect() + await asyncio.wait_for(task_b, timeout=2.0) + assert omni.detached == [handle.session_id] + assert omni.detach_refused == [(handle.session_id, 1)] + task.cancel() + with suppress(asyncio.CancelledError, Exception): + await task + + +@pytest.mark.asyncio +async def test_an_abandoned_takeover_leaves_the_live_connection_serving_and_owning_the_lease() -> None: + """A takeover cancelled mid-RPC must not put the connection it never replaced into disconnect grace. + + A is attached at generation 0. B's engine resume lands generation 1 and B + is cancelled before it reaches the registry, so A stays the current + attachment. Settling B's generation must hand it to A rather than detach + it: A keeps serving, and A's own disconnect later detaches generation 1, + which is the lease the engine actually holds. + """ + omni = FakeOmni() + handler = _handler(omni) + ws, handle, task = await _open(handler, omni) + token = ws.sent[0]["resume_token"] + registry = handler._attachment_registry + + omni.resume_gate = asyncio.Event() + ws_b, task_b = await _resume(handler, handle.session_id, token) + await asyncio.wait_for(omni.resume_started.wait(), timeout=2.0) + assert handle.lease_generation == 1 + + task_b.cancel() + with suppress(asyncio.CancelledError): + await task_b + omni.resume_gate.set() + await asyncio.gather(*omni.compensations) + + assert omni.detached == [], "A is still serving: its lease must not enter disconnect grace" + assert omni.detach_refused == [] + assert await registry.is_current_attachment(handle.session_id, 1), "A is still the attachment" + handle.deliver(AudioDelta(session_id=handle.session_id, response_id="r1", delta="aGk=")) + await ws.wait_for("response.output_audio.delta") + + # A now owns generation 1: its disconnect detaches the lease the engine holds. + ws.disconnect() + await asyncio.wait_for(task, timeout=2.0) + assert omni.detached == [handle.session_id] + assert omni.detach_refused == [] + + +@pytest.mark.asyncio +async def test_a_send_failure_detaches_the_failed_sockets_own_lease_not_a_pending_resumes() -> None: + """The pump's detach is fenced on the lease of the socket whose send failed, captured with it. + + The pump holds the registry's outbound lock while sending to A. B's + engine resume lands generation 1 meanwhile and B waits for that lock to + activate. When A's send fails, the pump drops A and must detach A's + generation 0, which the engine refuses, not B's generation 1, which would + put B into disconnect grace the moment it activates. + """ + omni = FakeOmni() + handler = _handler(omni) + ws, handle, task = await _open(handler, omni) + token = ws.sent[0]["resume_token"] + registry = handler._attachment_registry + + sending = asyncio.Event() + fail_send = asyncio.Event() + + async def parked_failing_send(payload: dict[str, Any]) -> None: + sending.set() + await fail_send.wait() + raise RuntimeError("socket closed") + + ws.send_json = parked_failing_send # type: ignore[method-assign] + handle.deliver(AudioDelta(session_id=handle.session_id, response_id="r1", delta="aGk=")) + await asyncio.wait_for(sending.wait(), timeout=2.0) + + # B: the engine resume lands; activation waits for the outbound lock the pump holds. + ws_b, task_b = await _resume(handler, handle.session_id, token) + await asyncio.sleep(0.05) + assert omni.resumed == [(handle.session_id, 0)] + assert handle.lease_generation == 1 + assert await registry.is_current_attachment(handle.session_id, 1), "B has not activated yet" + + fail_send.set() + resumed_b = await ws_b.wait_for("session.resumed") + assert resumed_b["attachment_generation"] == 2 + assert omni.detached == [], "A's lease (generation 0) is refused, B's lease is untouched" + assert omni.detach_refused == [(handle.session_id, 0)] + assert await registry.is_current_attachment(handle.session_id, 2) + + # B serves; its own disconnect detaches the lease it resumed. + ws_b.disconnect() + await asyncio.wait_for(task_b, timeout=2.0) + assert omni.detached == [handle.session_id] + ws.disconnect() + with suppress(asyncio.CancelledError, Exception): + await asyncio.wait_for(task, timeout=2.0) + + +@pytest.mark.asyncio +async def test_an_abandoned_resume_goes_to_the_resume_waiting_to_activate_not_the_socket_it_replaces() -> None: + """Three connections: A sending, B resumed and waiting for the outbound lock, C abandoned mid-RPC. + + C's generation 2 must be parked for B's activation. If it were handed to + A instead, A's send failure would detach generation 2 (accepted by the + engine), and B would then activate on a lease already in disconnect grace + and expire despite heartbeating. + """ + omni = FakeOmni() + handler = _handler(omni) + ws, handle, task = await _open(handler, omni) + token = ws.sent[0]["resume_token"] + registry = handler._attachment_registry + + sending = asyncio.Event() + fail_send = asyncio.Event() + + async def parked_failing_send(payload: dict[str, Any]) -> None: + sending.set() + await fail_send.wait() + raise RuntimeError("socket closed") + + ws.send_json = parked_failing_send # type: ignore[method-assign] + handle.deliver(AudioDelta(session_id=handle.session_id, response_id="r1", delta="aGk=")) + await asyncio.wait_for(sending.wait(), timeout=2.0) + + # B: engine resume lands generation 1; activation waits for the lock the pump holds. + ws_b, task_b = await _resume(handler, handle.session_id, token) + await asyncio.sleep(0.05) + assert handle.lease_generation == 1 + + # C: same unrotated token, engine resume lands generation 2, cancelled mid-RPC. + omni.resume_gate = asyncio.Event() + ws_c, task_c = await _resume(handler, handle.session_id, token) + await asyncio.wait_for(omni.resume_started.wait(), timeout=2.0) + assert handle.lease_generation == 2 + task_c.cancel() + with suppress(asyncio.CancelledError): + await task_c + omni.resume_gate.set() + await asyncio.gather(*omni.compensations) + assert omni.detached == [] and omni.detach_refused == [], "generation 2 is parked for B" + + # A's send fails: the pump drops A and detaches A's own generation 0, refused. + fail_send.set() + resumed_b = await ws_b.wait_for("session.resumed") + assert resumed_b["attachment_generation"] == 2 + assert omni.detach_refused == [(handle.session_id, 0)] + assert omni.detached == [], "the lease B serves is untouched" + assert await registry.is_current_attachment(handle.session_id, 2) + + # B serves generation 2 (inherited from C): its disconnect detaches that lease. + ws_b.disconnect() + await asyncio.wait_for(task_b, timeout=2.0) + assert omni.detached == [handle.session_id] + assert omni.detach_refused == [(handle.session_id, 0)] + ws.disconnect() + with suppress(asyncio.CancelledError, Exception): + await asyncio.wait_for(task, timeout=2.0) + + +@pytest.mark.asyncio +async def test_a_resume_cancelled_during_replay_settles_the_newer_lease_it_was_handed() -> None: + """Activation rollback must give back the lease the provisional attachment actually owned. + + A resumed generation 1, delivered ``session.resumed`` and is blocked + replaying events. B, with A's rotated token, resumed generation 2 and was + cancelled mid-RPC; its settlement handed 2 to A. When A is then cancelled + during replay, the registry rolls A back; detaching A's own generation 1 + would be refused and generation 2 would never enter disconnect grace. + """ + omni = FakeOmni() + handler = _handler(omni) + ws, handle, task = await _open(handler, omni) + token = ws.sent[0]["resume_token"] + registry = handler._attachment_registry + ws.disconnect() + await asyncio.wait_for(task, timeout=2.0) + assert omni.detached == [handle.session_id] + # Journaled while detached: A's activation has something to replay. + handle.deliver(AudioDelta(session_id=handle.session_id, response_id="r1", delta="one")) + await asyncio.sleep(0.05) + + # A: session.resumed goes out, the replay entry parks. + ws_a = FakeWebSocket({"duplex": "1", "resume": "1"}) + replaying = asyncio.Event() + sends = 0 + original_send = ws_a.send_json + + async def send_parking_on_replay(payload: dict[str, Any]) -> None: + nonlocal sends + sends += 1 + if sends == 2: + replaying.set() + await asyncio.Event().wait() + await original_send(payload) + + ws_a.send_json = send_parking_on_replay # type: ignore[method-assign] + task_a = asyncio.create_task(handler.handle_realtime_session(ws_a)) + ws_a.feed( + { + "type": "session.resume", + "session_id": handle.session_id, + "resume_token": token, + "last_received_server_event_seq": 0, + } + ) + await asyncio.wait_for(replaying.wait(), timeout=2.0) + resumed_a = await ws_a.wait_for("session.resumed") + assert handle.lease_generation == 1 + + # B: A's rotated token, engine resume lands generation 2, cancelled mid-RPC. + omni.resume_gate = asyncio.Event() + ws_b, task_b = await _resume(handler, handle.session_id, resumed_a["resume_token"]) + await asyncio.wait_for(omni.resume_started.wait(), timeout=2.0) + assert handle.lease_generation == 2 + task_b.cancel() + with suppress(asyncio.CancelledError): + await task_b + omni.resume_gate.set() + await asyncio.gather(*omni.compensations) + assert omni.detached == [handle.session_id], "generation 2 went to A's pending activation" + + # A is cancelled during replay: its rollback settles generation 2, the lease A owned. + task_a.cancel() + with suppress(asyncio.CancelledError): + await task_a + assert not await registry.has_attachment(handle.session_id) + assert omni.detached == [handle.session_id, handle.session_id] + assert omni.detach_refused == [], "generation 1 was never detached: the engine would refuse it" diff --git a/tests/entrypoints/openai/test_duplex_session_attachment.py b/tests/entrypoints/openai/test_duplex_session_attachment.py index 9d7a208b5..10b600407 100644 --- a/tests/entrypoints/openai/test_duplex_session_attachment.py +++ b/tests/entrypoints/openai/test_duplex_session_attachment.py @@ -552,3 +552,90 @@ async def test_registry_resume_cancelled_mid_delivery_rolls_back_like_a_failure( close=close, ) assert recovered.attachment_generation == 3 + + +@pytest.mark.asyncio +async def test_release_hands_back_the_lease_generation_the_dropped_connection_was_serving(mocker) -> None: + """The lease generation travels with the attachment, so a detach is fenced on the right lease.""" + registry = DuplexSessionAttachmentRegistry(replay_ttl_s=60.0, replay_max_bytes_per_session=4096) + send = mocker.AsyncMock() + created = await registry.create("sid-lease", send=send, close=mocker.AsyncMock(), lease_generation=0) + + # An abandoned takeover hands its generation to the connection still attached. + assert await registry.settle_lease_generation("sid-lease", 1) is None + assert await registry.settle_lease_generation("sid-lease", 0) is None, "generations only move forward" + + released = await registry.release_attachment("sid-lease", attachment_generation=created.attachment_generation) + assert released is not None + assert (released.attachment_generation, released.lease_generation) == (1, 1) + assert await registry.release_attachment("sid-lease", attachment_generation=None) is None + assert await registry.settle_lease_generation("sid-lease", 2) == 2, "nobody is attached: the caller detaches it" + + resumed = await registry.resume( + "sid-lease", + resume_token=created.resume_token.plaintext, + last_received_server_event_seq=0, + send=send, + close=mocker.AsyncMock(), + lease_generation=3, + ) + released = await registry.release_attachment("sid-lease", attachment_generation=resumed.attachment_generation) + assert released is not None and released.lease_generation == 3 + + +@pytest.mark.asyncio +async def test_an_orphaned_generation_waits_for_the_resume_that_is_about_to_activate(mocker) -> None: + """While a resume is pending, an orphan goes to its activation, not to the socket it will replace.""" + registry = DuplexSessionAttachmentRegistry(replay_ttl_s=60.0, replay_max_bytes_per_session=4096) + send = mocker.AsyncMock() + created = await registry.create("sid-pending", send=send, close=mocker.AsyncMock(), lease_generation=0) + + await registry.begin_resume("sid-pending") + assert await registry.settle_lease_generation("sid-pending", 2) is None, "parked for the pending activation" + released = await registry.release_attachment("sid-pending", attachment_generation=created.attachment_generation) + assert released is not None and released.lease_generation == 0, "the replaced socket keeps its own lease" + + resumed = await registry.resume( + "sid-pending", + resume_token=created.resume_token.plaintext, + last_received_server_event_seq=0, + send=send, + close=mocker.AsyncMock(), + lease_generation=1, + ) + assert await registry.end_resume("sid-pending") is None, "the activation absorbed the orphan" + released = await registry.release_attachment("sid-pending", attachment_generation=resumed.attachment_generation) + assert released is not None and released.lease_generation == 2 + + # A pending resume that never activates leaves the orphan to whoever ends the claim. + await registry.begin_resume("sid-pending") + assert await registry.settle_lease_generation("sid-pending", 3) is None + assert await registry.end_resume("sid-pending") == 3 + + +@pytest.mark.asyncio +async def test_a_rolled_back_activation_orphans_the_lease_it_was_handed(mocker) -> None: + """A provisional attachment may own a newer generation than its own resume produced.""" + registry = DuplexSessionAttachmentRegistry(replay_ttl_s=60.0, replay_max_bytes_per_session=4096) + created = await registry.create("sid-rollback", send=mocker.AsyncMock(), close=mocker.AsyncMock()) + await registry.detach("sid-rollback", attachment_generation=created.attachment_generation) + + async def failing_activation(payload): + raise RuntimeError("socket closed during replay") + + await registry.begin_resume("sid-rollback") + with pytest.raises(RuntimeError): + await registry.resume( + "sid-rollback", + resume_token=created.resume_token.plaintext, + last_received_server_event_seq=0, + send=failing_activation, + close=mocker.AsyncMock(), + activation_payload_factory=lambda token, generation: {"type": "session.resumed"}, + lease_generation=1, + ) + # Meanwhile a resume abandoned mid-RPC handed generation 2 to the provisional attachment. + # (Order does not matter: both land in the orphan slot while the claim is open.) + assert await registry.settle_lease_generation("sid-rollback", 2) is None + assert await registry.settle_lease_generation("sid-rollback", 1) is None + assert await registry.end_resume("sid-rollback") == 2, "the newest lease is the one to detach" diff --git a/tests/entrypoints/test_duplex_omni.py b/tests/entrypoints/test_duplex_omni.py index f33b3ea4d..ada212224 100644 --- a/tests/entrypoints/test_duplex_omni.py +++ b/tests/entrypoints/test_duplex_omni.py @@ -43,7 +43,20 @@ class FakeDuplexEngine(FakeAsyncOmniEngine): self.opened: list[tuple[str, DuplexSessionConfig]] = [] self.closed: list[tuple[str, str]] = [] self.resumed: list[tuple[str, int]] = [] + #: The control id of every resume RPC, in order (a replay repeats one). + self.resume_control_ids: list[str | None] = [] + #: Per session: the control id that produced the current lease generation, and that generation. + self.lease: dict[str, tuple[str | None, int]] = {} + #: How many resume RPCs, after applying the resume, answer with a + #: timeout: the engine landed it but the waiter never saw the result. + self.resume_timeouts_remaining = 0 self.touched: list[tuple[str, str]] = [] + #: The lease-generation fence of every detach, in order (``None`` = unfenced). + self.detach_generations: list[int | None] = [] + #: When set, ``resume_session_async`` parks on it after signalling ``resume_started``: + #: the engine has applied the resume, its result has not reached the caller. + self.resume_gate: asyncio.Event | None = None + self.resume_started = asyncio.Event() self.commands: list[tuple[str, commands.DuplexCommand]] = [] self.open_error: DuplexSessionError | None = None self.close_error: DuplexSessionError | None = None @@ -71,12 +84,33 @@ class FakeDuplexEngine(FakeAsyncOmniEngine): self.emit(session_id, SessionClosed(reason=reason)) return self._result("close", session_id) - async def resume_session_async(self, session_id, *, expected_lease_generation, timeout=None): + async def resume_session_async(self, session_id, *, expected_lease_generation, control_id=None, timeout=None): self.resumed.append((session_id, expected_lease_generation)) - return self._result("resume", session_id, lease_generation=expected_lease_generation + 1) + self.resume_control_ids.append(control_id) + resumed_by, generation = self.lease.get(session_id, (None, 0)) + if control_id is not None and control_id == resumed_by: + # The engine's replay answer: the generation this resume produced. + # Its answer can be lost to a timeout exactly like the first one. + if self.resume_timeouts_remaining > 0: + self.resume_timeouts_remaining -= 1 + raise DuplexSessionError("duplex resume timed out", code="timeout", retryable=True) + return self._result("resume", session_id, lease_generation=generation) + if expected_lease_generation != generation: + raise DuplexSessionError("duplex lease generation mismatch", code="session_resume_conflict") + generation += 1 + self.lease[session_id] = (control_id, generation) + if self.resume_gate is not None: + self.resume_started.set() + await self.resume_gate.wait() + if self.resume_timeouts_remaining > 0: + self.resume_timeouts_remaining -= 1 + raise DuplexSessionError("duplex resume timed out", code="timeout", retryable=True) + return self._result("resume", session_id, lease_generation=generation) - async def touch_session_async(self, session_id, *, activity, timeout=None): + async def touch_session_async(self, session_id, *, activity, expected_lease_generation=None, timeout=None): self.touched.append((session_id, activity)) + if activity == "detach": + self.detach_generations.append(expected_lease_generation) return self._result("touch", session_id) async def submit_command_async(self, session_id, command): @@ -324,3 +358,141 @@ async def test_shutdown_closes_every_handle(monkeypatch) -> None: assert first.close_reason == "shutdown" assert omni.sessions == {} assert engine.shutdown_called + + +@pytest.mark.asyncio +async def test_detach_session_forwards_the_lease_generation_fence(monkeypatch) -> None: + omni, engine = _make_omni(monkeypatch) + try: + handle = await omni.open_session() + await omni.detach_session(handle.session_id, expected_lease_generation=3) + await omni.detach_session(handle.session_id) + assert engine.detach_generations == [3, None] + finally: + omni.shutdown() + + +@pytest.mark.asyncio +async def test_a_resume_cancelled_mid_rpc_is_adopted_and_detached_once_it_lands(monkeypatch) -> None: + """The resume RPC runs in an executor: cancelling its waiter does not stop the engine. + + Without compensation the engine ends up with the lease resumed (no + disconnect grace) for a connection that is gone, and the handle keeps the + old generation, so a later reconnect resumes against a stale lease. The + outcome has to be observed off the cancelled task: adopt the generation + the resume produced and detach exactly that lease. + """ + omni, engine = _make_omni(monkeypatch) + try: + handle = await omni.open_session() + engine.resume_gate = asyncio.Event() + resume = asyncio.create_task(omni.resume_session(handle.session_id, expected_lease_generation=0)) + await asyncio.wait_for(engine.resume_started.wait(), timeout=2.0) + + resume.cancel() + with pytest.raises(asyncio.CancelledError): + await resume + # Nothing to give back yet: the engine has not answered. + assert handle.lease_generation == 0 + assert engine.detach_generations == [] + + engine.resume_gate.set() + await asyncio.gather(*list(omni._resume_compensations)) + assert handle.lease_generation == 1, "the generation the resume produced is adopted" + assert engine.touched == [(handle.session_id, "detach")] + assert engine.detach_generations == [1], "the detach is fenced on that generation" + finally: + omni.shutdown() + + +@pytest.mark.asyncio +async def test_a_resume_cancelled_then_timed_out_is_settled_by_replaying_it(monkeypatch) -> None: + """A timeout after the cancel leaves the outcome unknown: the resume is replayed under its id. + + The engine may have applied the resume after the waiter gave up, and the + RPC router drops the late answer. Replaying the same control id makes the + engine report the generation that resume produced (or land it now), so + the handle adopts the real generation and the lease is detached again. + """ + omni, engine = _make_omni(monkeypatch) + try: + handle = await omni.open_session() + engine.resume_gate = asyncio.Event() + engine.resume_timeouts_remaining = 1 + resume = asyncio.create_task(omni.resume_session(handle.session_id, expected_lease_generation=0)) + await asyncio.wait_for(engine.resume_started.wait(), timeout=2.0) + + resume.cancel() + with pytest.raises(asyncio.CancelledError): + await resume + engine.resume_gate.set() + await asyncio.gather(*list(omni._resume_compensations)) + + assert len(engine.resume_control_ids) == 2, "the resume was replayed once" + assert engine.resume_control_ids[0] == engine.resume_control_ids[1] is not None + assert engine.lease[handle.session_id][1] == 1, "the replay did not resume a second time" + assert handle.lease_generation == 1 + assert engine.detach_generations == [1] + finally: + omni.shutdown() + + +@pytest.mark.asyncio +async def test_an_abandoned_resume_is_settled_through_the_callers_callback(monkeypatch) -> None: + """A caller that owns attachments decides what the landed generation means; DuplexOmni does not detach.""" + omni, engine = _make_omni(monkeypatch) + try: + handle = await omni.open_session() + settled: list[int] = [] + + async def on_abandoned(lease_generation: int) -> None: + settled.append(lease_generation) + + engine.resume_gate = asyncio.Event() + resume = asyncio.create_task( + omni.resume_session(handle.session_id, expected_lease_generation=0, on_abandoned=on_abandoned) + ) + await asyncio.wait_for(engine.resume_started.wait(), timeout=2.0) + resume.cancel() + with pytest.raises(asyncio.CancelledError): + await resume + engine.resume_gate.set() + await asyncio.gather(*list(omni._resume_compensations)) + + assert settled == [1] + assert handle.lease_generation == 1 + assert engine.detach_generations == [], "the callback owns the decision" + finally: + omni.shutdown() + + +@pytest.mark.asyncio +async def test_an_abandoned_resume_keeps_replaying_across_repeated_timeouts(monkeypatch) -> None: + """A replay can time out too; giving up then would reopen the lost-answer window. + + The engine applied the resume, the abandoned RPC and the first replay both + lose their answers to timeouts, the second replay is answered: the handle + adopts the generation and the lease is settled, with the lease bumped once. + """ + omni, engine = _make_omni(monkeypatch) + try: + omni._resume_replay_backoff_s = (0.01, 0.02) + handle = await omni.open_session() + engine.resume_gate = asyncio.Event() + engine.resume_timeouts_remaining = 2 + resume = asyncio.create_task(omni.resume_session(handle.session_id, expected_lease_generation=0)) + await asyncio.wait_for(engine.resume_started.wait(), timeout=2.0) + + resume.cancel() + with pytest.raises(asyncio.CancelledError): + await resume + engine.resume_gate.set() + await asyncio.wait_for(asyncio.gather(*list(omni._resume_compensations)), timeout=5.0) + + assert len(engine.resume_control_ids) == 3, "the abandoned RPC, a replay that timed out, the replay answered" + assert len(set(engine.resume_control_ids)) == 1 + assert engine.lease[handle.session_id][1] == 1, "the lease was bumped once" + assert handle.lease_generation == 1 + assert engine.detach_generations == [1] + finally: + omni.shutdown() diff --git a/vllm_omni/engine/duplex/messages.py b/vllm_omni/engine/duplex/messages.py index ad264e0c5..7c28b4638 100644 --- a/vllm_omni/engine/duplex/messages.py +++ b/vllm_omni/engine/duplex/messages.py @@ -65,6 +65,11 @@ class TouchDuplexSessionMessage(EngineQueueMessage, kw_only=True): control_id: str session_id: str activity: str + #: For a detach: the lease generation the caller opened or resumed + #: against. The manager refuses to detach a newer lease, so a connection + #: giving up its own resume can never put the lease a later resume owns + #: into disconnect grace. ``None`` detaches whatever generation is current. + expected_lease_generation: int | None = None class DuplexSessionCommandMessage(EngineQueueMessage, kw_only=True): diff --git a/vllm_omni/engine/duplex/session/append_task.py b/vllm_omni/engine/duplex/session/append_task.py index 7425e76a1..7c0723471 100644 --- a/vllm_omni/engine/duplex/session/append_task.py +++ b/vllm_omni/engine/duplex/session/append_task.py @@ -13,6 +13,13 @@ That compensation is the reason this is an object. The five steps of an append all need the same eleven values, so as closures over ``_start_append`` they read as one function with five entry points and no way to test any of them. Named fields make the captured state explicit and the compensation a method. + +The precreated response has two other owners the compensation defers to. When +the runner cancels the append on purpose (``DuplexSessionTasks.cancel_append_tasks``, +for a ``response.cancel``, a barge-in or a close), the runner ends that response +itself right after, with the cancel status the client asked for. When the +session is already closing, the close ends it. Everywhere else, the append is +the last thing that knows the response will never be filled, and it fails it. """ from __future__ import annotations @@ -53,6 +60,10 @@ class AppendAttempt: retained_committed_payload: dict[str, object] | None #: Response reserved before submission, to be failed if submission does not happen. precreated_response_id: str | None + #: Whether ``_start_append`` bound the session's active request to this + #: append (a final or response-bound one), so an append that never happens + #: has to unbind it again. + owns_request: bool = False #: Commits timing state once the runtime accepts the append (before any #: returned output event can clear the continuation chain). on_append_accepted: Callable[[float], None] | None = None @@ -75,11 +86,31 @@ class AppendAttempt: ): self.ctx.session.release_input_bytes(model_state.clear_committed_audio()) - def abandon(self) -> None: - """Give back everything this append reserved but never used.""" + def abandon(self, *, reason: str = "runtime_append_failed", response_taken_over: bool = False) -> None: + """Give back everything this append reserved but never used. + + That includes the response it precreated: once the append is off, + nothing will ever fill it, and a client that saw ``response.created`` + is owed a ``response.done``. It is left alone only when somebody else + ends it: the runner that cancelled this append on purpose + (``response_taken_over``) or the close already under way. + """ if self.pcm_reservation is not None: self.pcm_reservation.rollback() self.discard_retained_audio() + if response_taken_over: + return + session = self.ctx.session + if self.ctx.run.closing or session.state != DuplexSessionState.OPEN: + # The close ends the active response and request itself, with the + # session's one terminal event after them. + return + if self.owns_request and session.epoch == self.epoch: + # Compare-before-clear: only if the session still points at this + # append's request, exactly as the submitted path does. + session.clear_request(self.request_id) + if self._precreated_response_still_active(): + self._fail_precreated_response(reason) def release_on_failure(self, done: asyncio.Task[bool]) -> None: """Done-callback: a cancelled or failed append must not hold the audio.""" @@ -105,7 +136,22 @@ class AppendAttempt: # ------------------------------------------------------------------ # async def run_in_wire_order(self, predecessor: asyncio.Task[bool] | None) -> bool: - """Wait for the previous append, then submit this one if it still applies.""" + """Wait for the previous append, then submit this one if it still applies. + + A cancellation is compensated here, wherever it lands (waiting for the + predecessor or mid-submit): the reservations go back, and the + precreated response is failed unless the runner cancelled the append + itself, in which case it ends the response with the status the + cancellation was for. + """ + try: + return await self._run(predecessor) + except asyncio.CancelledError: + taken_over = self.ctx.tasks.cancelled_by_runner(asyncio.current_task()) + self.abandon(reason="append_cancelled", response_taken_over=taken_over) + raise + + async def _run(self, predecessor: asyncio.Task[bool] | None) -> bool: if predecessor is not None: try: predecessor_ok = await predecessor @@ -127,7 +173,11 @@ class AppendAttempt: # Called off, not failed: the chain behind it still runs. return True if self.pcm_reservation is not None and not self.pcm_reservation.active: - self.discard_retained_audio() + # The client cleared the input buffer while this append waited its + # turn: the audio is gone, so is the response it was going to + # answer, and that is the client's doing, not a runtime failure. + # The reservation rollback itself is already a no-op. + self.abandon(reason="input_cleared") return False return await self._submit() @@ -152,8 +202,6 @@ class AppendAttempt: self.discard_retained_audio() else: self.abandon() - if not append_ok and self._precreated_response_still_active(): - self._fail_precreated_response() if not append_ok and session.state == DuplexSessionState.CLOSED: self.ctx.run.runtime_closed = True return False @@ -171,11 +219,11 @@ class AppendAttempt: self.out.emit_events([session.signal_turn(DuplexTurnEventType.USER_STARTED.value)]) return append_ok except asyncio.CancelledError: - if self.pcm_reservation is not None: - self.pcm_reservation.rollback() raise except Exception as exc: - self.abandon() + # Compensate while the session is still open, so the client gets + # the failed ``response.done`` before the error and the close. + self.abandon(reason="runtime_append_task_failed") logger.exception("Native duplex append task failed: %s", exc) self.model.send_runtime_error("runtime_append_task_failed", exc) if session.state != DuplexSessionState.CLOSED: @@ -188,7 +236,7 @@ class AppendAttempt: and self.ctx.session.active_response_id == self.precreated_response_id ) - def _fail_precreated_response(self) -> None: + def _fail_precreated_response(self, reason: str) -> None: session = self.ctx.session session.end_response(commit_text=False) self.out.emit( @@ -199,7 +247,7 @@ class AppendAttempt: "epoch": session.epoch, "committed": False, "status": "failed", - "status_details": {"type": "failed", "reason": "runtime_append_failed"}, + "status_details": {"type": "failed", "reason": reason}, "playback": session.playback.as_dict(), } ) diff --git a/vllm_omni/engine/duplex/session/context.py b/vllm_omni/engine/duplex/session/context.py index 02853561a..fde9425be 100644 --- a/vllm_omni/engine/duplex/session/context.py +++ b/vllm_omni/engine/duplex/session/context.py @@ -20,7 +20,7 @@ from __future__ import annotations import asyncio from collections.abc import Awaitable, Callable from dataclasses import dataclass, field -from typing import TYPE_CHECKING, Protocol, TypeVar +from typing import TYPE_CHECKING, Any, Protocol, TypeVar if TYPE_CHECKING: from vllm.outputs import RequestOutput @@ -117,6 +117,10 @@ class DuplexSessionTasks: append_tasks: dict[asyncio.Task[bool], DuplexAppendTaskMeta] = field(default_factory=dict) append_tail: asyncio.Task[bool] | None = None active_response_task: asyncio.Task[None] | None = None + #: Appends the runner cancelled on purpose (a cancel, barge-in or close). + #: The runner ends the response such an append precreated, so the append + #: leaves it alone on its way out instead of failing it. + runner_cancelled: set[asyncio.Task[bool]] = field(default_factory=set) def track_append_task( self, @@ -127,16 +131,31 @@ class DuplexSessionTasks: response_bound: bool, ) -> None: self.append_tasks[task] = DuplexAppendTaskMeta(epoch, final, response_bound) - task.add_done_callback(self.append_tasks.pop) + task.add_done_callback(self._forget_append_task) + + def _forget_append_task(self, task: asyncio.Task[bool]) -> None: + self.append_tasks.pop(task, None) + self.runner_cancelled.discard(task) + + def cancelled_by_runner(self, task: asyncio.Task[Any] | None) -> bool: + """Whether ``task`` is an append the runner cancelled through ``cancel_append_tasks``.""" + return task is not None and task in self.runner_cancelled def has_response_bound_append_tasks(self) -> bool: return any(meta.response_bound for meta in self.append_tasks.values()) async def cancel_append_tasks(self, timeout_s: float = 0.25, *, response_bound_only: bool = False) -> bool: + """Cancel the tracked appends; the caller owns their responses from here on. + + Every caller ends the active response itself right after (with the + cancel status, or silently as part of a close), so the cancelled + appends are told not to fail the response they precreated. + """ tasks = [task for task, meta in self.append_tasks.items() if not response_bound_only or meta.response_bound] if not tasks: return False cancelled_tail = self.append_tail if self.append_tail in tasks else None + self.runner_cancelled.update(tasks) for task in tasks: task.cancel() try: diff --git a/vllm_omni/engine/duplex/session/engine_session.py b/vllm_omni/engine/duplex/session/engine_session.py index 37ab18d4b..4e4bfdf6e 100644 --- a/vllm_omni/engine/duplex/session/engine_session.py +++ b/vllm_omni/engine/duplex/session/engine_session.py @@ -280,9 +280,9 @@ class DuplexEngineSession: self._validate_fence(self.fence) self.lease.detach(self._clock()) - def resume_lease(self, *, expected_lease_generation: int) -> int: + def resume_lease(self, *, expected_lease_generation: int, control_id: str | None = None) -> int: self._validate_fence(self.fence) - return self.lease.resume(self._clock(), expected_generation=expected_lease_generation) + return self.lease.resume(self._clock(), expected_generation=expected_lease_generation, control_id=control_id) def begin_lease_operation(self, fence: DuplexFence, operation_id: str) -> None: self._validate_fence(fence) diff --git a/vllm_omni/engine/duplex/session/lease.py b/vllm_omni/engine/duplex/session/lease.py index aea204039..89e75efd3 100644 --- a/vllm_omni/engine/duplex/session/lease.py +++ b/vllm_omni/engine/duplex/session/lease.py @@ -40,6 +40,12 @@ class DuplexLeaseState: detached_at: float | None = None active_operations: set[str] = field(default_factory=set) terminal_reason: str | None = None + #: Control id of the resume that produced ``generation`` (``None`` for the + #: open). A resume replayed with the same id is answered with the current + #: generation instead of bumping it again: the caller lost the first + #: answer (its RPC waiter was cancelled or timed out) and is finding out + #: whether the resume landed. + resumed_by: str | None = None @property def expires_at(self) -> float | None: @@ -61,11 +67,16 @@ class DuplexLeaseState: self.touch(now, DuplexLeaseActivity.DETACH) self.detached_at = now - def resume(self, now: float, *, expected_generation: int) -> int: + def resume(self, now: float, *, expected_generation: int, control_id: str | None = None) -> int: self._require_open() + if control_id is not None and control_id == self.resumed_by: + # The same resume asking again: it already produced this + # generation, so answer without changing anything. + return self.generation if expected_generation != self.generation: raise ValueError(f"duplex lease generation mismatch: expected {self.generation}, got {expected_generation}") self.generation += 1 + self.resumed_by = control_id self.detached_at = None self.last_activity = now return self.generation @@ -104,6 +115,7 @@ class DuplexLeaseState: return False self.terminal_reason = reason self.generation += 1 + self.resumed_by = None self.detached_at = None self.active_operations.clear() return True diff --git a/vllm_omni/engine/duplex/session/manager.py b/vllm_omni/engine/duplex/session/manager.py index 1faea6e22..38eeb7ab6 100644 --- a/vllm_omni/engine/duplex/session/manager.py +++ b/vllm_omni/engine/duplex/session/manager.py @@ -523,32 +523,74 @@ class DuplexSessionManager: # and the stage resource straight back. self._abandoned_opens.discard(session_id) await self._close_runner(runner, kind="close", reason="open_abandoned") + except asyncio.CancelledError: + # A cancel (engine teardown, task cancellation) while awaiting the + # plugin or the result sink is not an ``Exception``: without this + # branch the runner stayed in ``runners`` with its Stage0 + # reservation, and since admission counts runners, every cancelled + # open burned one slot for good. Nobody is answered: the caller's + # RPC waiter is being torn down with us or times out. + logger.info("open_duplex_session cancelled for %s; rolling the admission back", session_id) + await self._rollback_open(session_id, session=session, runner=runner) + raise except Exception as exc: error_code, _, _ = self._control_error(exc) if error_code in {"resource_exhausted", "session_exists"}: logger.info("open_duplex_session rejected: %s", exc) else: logger.exception("open_duplex_session failed: %s", exc) - if runner is not None and self.runners.get(session_id) is runner: - self.runners.pop(session_id, None) - try: - await runner.shutdown() - except Exception: - logger.exception("duplex open rollback: runner shutdown failed for %s", session_id) - if session is not None: - reserved = session.release_all_requests() - if reserved: - try: - await self.stage_port.cleanup(list(reserved)) - except Exception: - logger.warning("duplex open rollback: request cleanup pending for %s", session_id) - self._unregister_session_requests(session_id) + await self._rollback_open(session_id, session=session, runner=runner) await self._put_result(message, operation="open", ok=False, error=exc) finally: if holds_admission_slot: self._admitting.discard(session_id) self._abandoned_opens.discard(session_id) + async def _rollback_open( + self, + session_id: str, + *, + session: DuplexEngineSession | None, + runner: DuplexSessionRunner | None, + ) -> None: + """Undo whatever an open that will not complete already did. + + The bookkeeping that decides admission (``runners`` and the session's + request reservations) is undone synchronously first, so a second + cancellation landing in the awaits below cannot leave the slot burned. + The Stage0 request ids are recorded as a pending request cleanup + before anything is awaited: if the runner shutdown or the cleanup + itself is cancelled or fails, the reaper retries the cleanup instead + of the orchestrator keeping the request state forever. + """ + owned_runner = runner is not None and self.runners.get(session_id) is runner + if owned_runner: + self.runners.pop(session_id, None) + pending: _PendingRequestCleanup | None = None + key: tuple[str, int] | None = None + if session is not None: + reserved = tuple(session.release_all_requests()) + self._unregister_session_requests(session_id) + if reserved: + key = (session_id, session.lease_generation) + pending = _PendingRequestCleanup( + session_id=session_id, + lease_generation=session.lease_generation, + request_ids=reserved, + abort=False, + ) + self._pending_request_cleanups[key] = pending + if owned_runner and runner is not None: + try: + await runner.shutdown() + except Exception: + logger.exception("duplex open rollback: runner shutdown failed for %s", session_id) + if pending is not None and key is not None: + try: + await self._complete_request_cleanup(key, pending) + except Exception: + logger.warning("duplex open rollback: request cleanup pending for %s; the reaper retries", session_id) + def _require_runner(self, session_id: str) -> DuplexSessionRunner: runner = self.runners.get(session_id) if runner is None: @@ -668,7 +710,13 @@ class DuplexSessionManager: runner = self._require_runner(message.session_id) session = runner.session try: - session.resume_lease(expected_lease_generation=message.expected_lease_generation) + # Keyed by the control id, so a caller that lost the answer + # (waiter cancelled or timed out) can replay the same resume + # to learn whether it landed instead of resuming twice. + session.resume_lease( + expected_lease_generation=message.expected_lease_generation, + control_id=message.control_id, + ) except ValueError as exc: raise DuplexSessionError(str(exc), code="session_resume_conflict") from exc await self._put_result(message, operation="resume", ok=True, session=session) @@ -683,6 +731,14 @@ class DuplexSessionManager: session = runner.session activity = DuplexLeaseActivity(message.activity) if activity is DuplexLeaseActivity.DETACH: + expected = message.expected_lease_generation + if expected is not None and expected != session.lease_generation: + # The caller is giving up a lease it no longer holds: a + # later resume owns the current one and must keep it. + raise DuplexSessionError( + f"duplex lease generation mismatch: expected {session.lease_generation}, got {expected}", + code="session_resume_conflict", + ) session.detach_lease() else: session.touch_lease(activity) diff --git a/vllm_omni/engine/duplex/session/runner.py b/vllm_omni/engine/duplex/session/runner.py index 02dc61708..ca9793938 100644 --- a/vllm_omni/engine/duplex/session/runner.py +++ b/vllm_omni/engine/duplex/session/runner.py @@ -1107,6 +1107,7 @@ class DuplexSessionRunner: operation_id=operation_id, retained_committed_payload=retained_committed_payload, precreated_response_id=session.active_response_id if precreate_response else None, + owns_request=final or precreate_response, on_append_accepted=on_append_accepted, before_append=before_append, ) diff --git a/vllm_omni/engine/duplex_omni_engine.py b/vllm_omni/engine/duplex_omni_engine.py index 56ca3e352..c7c3a6328 100644 --- a/vllm_omni/engine/duplex_omni_engine.py +++ b/vllm_omni/engine/duplex_omni_engine.py @@ -214,9 +214,12 @@ class DuplexOmniEngine(AsyncOmniEngine): session_id: str, *, expected_lease_generation: int, + control_id: str | None = None, timeout: float | None = _DEFAULT_CONTROL_TIMEOUT_S, ) -> DuplexControlResultMessage: - control_id = uuid.uuid4().hex + # A caller-supplied id makes the resume replayable: the engine answers + # a repeated id with the generation that resume produced. + control_id = control_id or uuid.uuid4().hex return self._execute_control( ResumeDuplexSessionMessage( control_id=control_id, @@ -234,6 +237,7 @@ class DuplexOmniEngine(AsyncOmniEngine): session_id: str, *, expected_lease_generation: int, + control_id: str | None = None, timeout: float | None = _DEFAULT_CONTROL_TIMEOUT_S, ) -> DuplexControlResultMessage: loop = asyncio.get_running_loop() @@ -242,6 +246,7 @@ class DuplexOmniEngine(AsyncOmniEngine): lambda: self._resume_session( session_id, expected_lease_generation=expected_lease_generation, + control_id=control_id, timeout=timeout, ), ) @@ -251,6 +256,7 @@ class DuplexOmniEngine(AsyncOmniEngine): session_id: str, *, activity: str, + expected_lease_generation: int | None = None, timeout: float | None = _DEFAULT_CONTROL_TIMEOUT_S, ) -> DuplexControlResultMessage: control_id = uuid.uuid4().hex @@ -259,6 +265,7 @@ class DuplexOmniEngine(AsyncOmniEngine): control_id=control_id, session_id=session_id, activity=activity, + expected_lease_generation=expected_lease_generation, ), control_id=control_id, operation="touch", @@ -271,12 +278,18 @@ class DuplexOmniEngine(AsyncOmniEngine): session_id: str, *, activity: str, + expected_lease_generation: int | None = None, timeout: float | None = _DEFAULT_CONTROL_TIMEOUT_S, ) -> DuplexControlResultMessage: loop = asyncio.get_running_loop() return await loop.run_in_executor( None, - lambda: self._touch_session(session_id, activity=activity, timeout=timeout), + lambda: self._touch_session( + session_id, + activity=activity, + expected_lease_generation=expected_lease_generation, + timeout=timeout, + ), ) def _submit_command(self, session_id: str, command: DuplexCommand) -> None: diff --git a/vllm_omni/entrypoints/duplex/serving.py b/vllm_omni/entrypoints/duplex/serving.py index ff1abbe25..5e6f3da70 100644 --- a/vllm_omni/entrypoints/duplex/serving.py +++ b/vllm_omni/entrypoints/duplex/serving.py @@ -36,11 +36,13 @@ from vllm_omni.engine.duplex.events import ( SessionResyncRequired, ) from vllm_omni.engine.duplex.messages import DuplexSessionError -from vllm_omni.entrypoints.duplex.realtime_input import RealtimeEnvelope, parse_resume_request +from vllm_omni.entrypoints.duplex.realtime_input import RealtimeEnvelope, ResumeRequest, parse_resume_request from vllm_omni.entrypoints.duplex.session_attachment import ( + DuplexDetachedAttachment, DuplexJournalGapError, DuplexJournalOverflowError, DuplexSessionAttachmentRegistry, + DuplexSessionResumeResult, InvalidResumeTokenError, ResumeToken, ) @@ -67,6 +69,12 @@ _PUMP_DRAIN_TIMEOUT_S = 5.0 _UNJOURNALED_EVENTS = (SessionCreated,) +def _log_rollback_failure(done: asyncio.Task[None]) -> None: + """A shielded resume rollback outlives its caller; its failure must still be seen.""" + if not done.cancelled() and done.exception() is not None: + logger.error("duplex resume rollback %s failed: %r", done.get_name(), done.exception()) + + @dataclass class _Attachment: """What one websocket connection knows about the session it is attached to.""" @@ -199,10 +207,14 @@ class OmniDuplexSessionHandler: await send_json(envelope.error_payload("invalid_session_config", str(exc))) return None attachment_send, attachment_close = attachment_callbacks(websocket) + # The registry keeps the engine lease generation each attachment is + # serving, so every detach is fenced on the dropped connection's own + # lease (see ``release_attachment``). created = await self._attachment_registry.create( handle.session_id, send=attachment_send, close=attachment_close, + lease_generation=handle.lease_generation, ) resume_supported = bool(handle.capabilities.supports_session_resume) attachment = _Attachment(handle=handle, generation=created.attachment_generation) @@ -258,14 +270,44 @@ class OmniDuplexSessionHandler: except (KeyError, ValueError) as exc: await send_json(envelope.error_payload("session_resume_conflict", str(exc))) return None + # The claim tells the registry a resume is about to activate: a lease + # generation orphaned meanwhile (another resume abandoned mid-RPC) is + # parked for this activation instead of being handed to the socket it + # is about to replace. The claim ends whatever happens below, and an + # orphan nobody else will serve is then ours to put into grace. + await self._attachment_registry.begin_resume(session_id) try: + return await self._resume_claimed(websocket, envelope, request, handle, send_json) + finally: + await self._end_resume_claim(session_id) + + async def _resume_claimed( + self, + websocket: WebSocket, + envelope: RealtimeEnvelope, + request: ResumeRequest, + handle: DuplexSessionHandle, + send_json: SendJson, + ) -> _Attachment | None: + session_id = request.session_id + try: + # If this task is cancelled while the RPC is in flight, DuplexOmni + # observes the outcome and settles a resume that landed through + # ``_settle_abandoned_resume``: the generation goes to the resume + # waiting to activate, else to whoever is attached, else the lease + # goes back into disconnect grace. await self._omni.resume_session( session_id, expected_lease_generation=handle.lease_generation, + on_abandoned=lambda lease_generation: self._settle_abandoned_resume(session_id, lease_generation), ) except DuplexSessionError as exc: await send_json(envelope.error_payload("runtime_resume_failed", str(exc))) return None + # The generation this resume produced, read before anything else can + # run: the handle is shared by every connection of the session, and a + # concurrent resume adopts its own, newer generation into it. + lease_generation = handle.lease_generation attachment_send, attachment_close = attachment_callbacks(websocket) @@ -277,6 +319,12 @@ class OmniDuplexSessionHandler: resume_token=token.plaintext, ).to_realtime() + # From here on the engine lease is resumed (``detached_at`` cleared), so + # every exit that does not hand the attachment to the caller has to + # put the lease back into its disconnect grace, or the session sits + # attached to nothing until idle expiry. That includes cancellation of + # the handler task, which is not an ``Exception``. + resumed: DuplexSessionResumeResult | None = None try: resumed = await self._attachment_registry.resume( session_id, @@ -285,26 +333,27 @@ class OmniDuplexSessionHandler: send=attachment_send, close=attachment_close, activation_payload_factory=activation_payload_factory, + lease_generation=lease_generation, ) + replaced = resumed.replaced_attachment + if replaced is not None: + with suppress(Exception): + await replaced.send( + SessionReplaced(session_id=session_id, attachment_generation=replaced.generation).to_realtime() + ) + with suppress(Exception): + await replaced.close("session_replaced") + except asyncio.CancelledError: + await self._abandon_resume( + session_id, + attachment_generation=resumed.attachment_generation if resumed is not None else None, + lease_generation=lease_generation, + ) + raise except Exception as exc: - # The engine lease was already resumed above, so a failed activation - # would leave the session attached to nothing -- unless another - # connection won the race and is attached right now. Detaching then - # would start the disconnect grace for the *winner*, which ordinary - # heartbeats do not clear. Roll back only what this attempt owns. - if not await self._attachment_registry.has_attachment(session_id): - with suppress(DuplexSessionError): - await self._omni.detach_session(session_id) + await self._abandon_resume(session_id, attachment_generation=None, lease_generation=lease_generation) await send_json(envelope.error_payload("session_resume_conflict", str(exc))) return None - replaced = resumed.replaced_attachment - if replaced is not None: - with suppress(Exception): - await replaced.send( - SessionReplaced(session_id=session_id, attachment_generation=replaced.generation).to_realtime() - ) - with suppress(Exception): - await replaced.close("session_replaced") # A reconnect brings a fresh envelope carrying pcm16/16 kHz wire # defaults. The negotiated input format is a wire default, not part of # the public session object, so it has to be carried over explicitly: @@ -315,6 +364,96 @@ class OmniDuplexSessionHandler: self._start_pump(handle, None) return _Attachment(handle=handle, generation=resumed.attachment_generation) + async def _abandon_resume( + self, + session_id: str, + *, + attachment_generation: int | None, + lease_generation: int, + ) -> None: + """Roll the engine lease of a resume this connection will never serve back into disconnect grace. + + Only what this attempt owns is rolled back. ``attachment_generation`` + is the attachment it activated, or ``None`` when activation did not + happen (or was rolled back by the registry): then the generation this + attempt's engine resume produced is settled like an abandoned resume + (``_settle_abandoned_resume``). The registry knows more than this + method does: a rolled-back activation may have been handed a newer + generation meanwhile, which the registry parks as the orphan, and the + claim ending in ``_resume`` gives it back if nobody else will serve + it. Every engine detach is fenced on the generation being given up, so + a newer lease a concurrent resume produced is never touched. + + The rollback runs shielded: a second cancellation landing between the + registry detach and the engine detach would otherwise strand the + session exactly the way this method compensates for. + """ + rollback = asyncio.ensure_future( + self._roll_back_resume( + session_id, + attachment_generation=attachment_generation, + lease_generation=lease_generation, + ) + ) + rollback.add_done_callback(_log_rollback_failure) + await asyncio.shield(rollback) + + async def _roll_back_resume( + self, + session_id: str, + *, + attachment_generation: int | None, + lease_generation: int, + ) -> None: + if attachment_generation is None: + await self._settle_abandoned_resume(session_id, lease_generation) + return + released = await self._attachment_registry.release_attachment( + session_id, attachment_generation=attachment_generation + ) + if released is not None: + await self._detach_lease(released) + + async def _settle_abandoned_resume(self, session_id: str, lease_generation: int) -> None: + """Settle the engine lease generation of a resume its connection never got to serve. + + The engine already bumped the lease to ``lease_generation`` and + cleared its disconnect grace. The registry picks the owner: a resume + waiting to activate (it will serve the session), else the connection + attached right now (a takeover that never activated leaves the + previous socket serving, and its own disconnect must be able to + detach this lease), else nobody, and then the lease goes back into + disconnect grace here, fenced on the generation so a later resume is + never detached. + """ + orphan = await self._attachment_registry.settle_lease_generation(session_id, lease_generation) + if orphan is not None: + with suppress(DuplexSessionError): + await self._omni.detach_session(session_id, expected_lease_generation=orphan) + + async def _end_resume_claim(self, session_id: str) -> None: + """End this connection's resume claim; detach an orphaned lease nobody else will serve. + + Runs from a ``finally``, so possibly during cancellation: the detach + is shielded like the other rollbacks. + """ + orphan = await self._attachment_registry.end_resume(session_id) + if orphan is None: + return + + async def detach_orphan() -> None: + with suppress(DuplexSessionError): + await self._omni.detach_session(session_id, expected_lease_generation=orphan) + + rollback = asyncio.ensure_future(detach_orphan()) + rollback.add_done_callback(_log_rollback_failure) + await asyncio.shield(rollback) + + async def _detach_lease(self, released: DuplexDetachedAttachment) -> None: + """Engine-owned disconnect grace for the lease a dropped connection was serving.""" + with suppress(DuplexSessionError): + await self._omni.detach_session(released.session_id, expected_lease_generation=released.lease_generation) + # ------------------------------------------------------------------ # # Outbound pump (session-scoped, survives reconnects) # # ------------------------------------------------------------------ # @@ -420,9 +559,13 @@ class OmniDuplexSessionHandler: if handle is None or handle.closed: return if handle.capabilities.supports_session_resume: - if await self._attachment_registry.detach(session_id, attachment_generation=None): - with suppress(DuplexSessionError): - await self._omni.detach_session(session_id) + # The lease generation comes with the dropped attachment, captured + # under the registry lock: the shared handle may already carry the + # generation of a resume that is waiting to activate behind this + # very send, and that lease is not ours to detach. + released = await self._attachment_registry.release_attachment(session_id, attachment_generation=None) + if released is not None: + await self._detach_lease(released) return with suppress(DuplexSessionError): await handle.close(reason="disconnect") @@ -532,10 +675,12 @@ class OmniDuplexSessionHandler: if handle.closed: return if handle.capabilities.supports_session_resume: - if await self._attachment_registry.detach(handle.session_id, attachment_generation=attachment.generation): + released = await self._attachment_registry.release_attachment( + handle.session_id, attachment_generation=attachment.generation + ) + if released is not None: # Engine-owned disconnect grace: expiry arrives as session.expired. - with suppress(DuplexSessionError): - await self._omni.detach_session(handle.session_id) + await self._detach_lease(released) # Otherwise a newer connection already took the session over; the # replaced socket's disconnect must not touch it. return diff --git a/vllm_omni/entrypoints/duplex/session_attachment.py b/vllm_omni/entrypoints/duplex/session_attachment.py index f2dc09f1c..b365995ec 100644 --- a/vllm_omni/entrypoints/duplex/session_attachment.py +++ b/vllm_omni/entrypoints/duplex/session_attachment.py @@ -184,6 +184,15 @@ class DuplexSessionAttachmentCreated: resume_token: ResumeToken = field(repr=False) +@dataclass(frozen=True) +class DuplexDetachedAttachment: + """The transport a detach dropped, with the engine lease generation it was serving.""" + + session_id: str + attachment_generation: int + lease_generation: int + + @dataclass(frozen=True) class DuplexSessionResumeResult: session_id: str @@ -200,6 +209,21 @@ class _DuplexSessionAttachmentState: journal: DuplexEventJournal attachment: DuplexTransportAttachment | None attachment_generation: int + #: The engine lease generation the attached connection is serving: the one + #: its open or resume produced, or one an abandoned resume handed to it. + #: Read together with the attachment under the lock, so a detach always + #: gives back the lease of the connection it drops, never a newer one. + lease_generation: int = 0 + #: Resumes between ``begin_resume`` and ``end_resume``: their engine lease + #: resume is in flight or landed, and they have not activated yet. While + #: one is pending, the next activation is the connection that will serve, + #: so an orphaned generation waits for it instead of going to the socket + #: that is about to be replaced. + pending_resumes: int = 0 + #: An engine lease generation nobody serves yet (its resume was abandoned, + #: or its provisional attachment was rolled back), waiting for the next + #: activation to absorb it. Generations only move forward. + orphaned_lease_generation: int | None = None outbound_lock: asyncio.Lock = field(default_factory=asyncio.Lock, repr=False) recovery_token_digest: bytes | None = field(default=None, repr=False) @@ -231,6 +255,7 @@ class DuplexSessionAttachmentRegistry: *, send: Callable[[dict[str, object]], Awaitable[None]], close: Callable[[str], Awaitable[None]], + lease_generation: int = 0, ) -> DuplexSessionAttachmentCreated: async with self._lock: if session_id in self._sessions: @@ -251,6 +276,7 @@ class DuplexSessionAttachmentRegistry: close=close, ), attachment_generation=generation, + lease_generation=lease_generation, ) return DuplexSessionAttachmentCreated( session_id=session_id, @@ -310,14 +336,91 @@ class DuplexSessionAttachmentRegistry: session answers ``False`` so a second disconnect signal for the same socket cannot restart the engine's disconnect grace window. """ + return await self.release_attachment(session_id, attachment_generation=attachment_generation) is not None + + async def release_attachment( + self, session_id: str, *, attachment_generation: int | None = None + ) -> DuplexDetachedAttachment | None: + """``detach`` that also hands back the lease generation the dropped connection was serving. + + Captured under the same lock as the attachment: the shared session + handle may already carry a newer generation from a resume that has + not activated yet, and the engine detach that follows must be fenced + on the dropped connection's own lease, not on that one. + """ async with self._lock: state = self._sessions.get(session_id) if state is None or state.attachment is None: - return False + return None if attachment_generation is not None and state.attachment_generation != attachment_generation: - return False + return None state.attachment = None - return True + return DuplexDetachedAttachment( + session_id=session_id, + attachment_generation=state.attachment_generation, + lease_generation=state.lease_generation, + ) + + async def begin_resume(self, session_id: str) -> None: + """Claim that a resume of ``session_id`` is about to land and activate. + + Held from before the engine lease resume until ``end_resume``. While + any claim is open, a lease generation that loses its owner is parked + for the next activation (see ``settle_lease_generation``) instead of + being handed to the currently attached connection, which that + activation is about to replace. + """ + async with self._lock: + self._require(session_id).pending_resumes += 1 + + async def end_resume(self, session_id: str) -> int | None: + """Release a ``begin_resume`` claim, whether or not the resume activated. + + Returns a lease generation the caller must detach itself: one that + was parked for an activation that never came, with nobody attached + to take it over. ``None`` when nothing is owed. + """ + async with self._lock: + state = self._sessions.get(session_id) + if state is None: + return None + state.pending_resumes = max(0, state.pending_resumes - 1) + return self._settle_orphan_locked(state) + + async def settle_lease_generation(self, session_id: str, lease_generation: int) -> int | None: + """Find an owner for the engine lease generation of a resume its connection never served. + + The engine already bumped its lease to ``lease_generation``. The owner + is, in order: the resume waiting to activate (it will serve the session + and absorbs the generation when it activates); else the connection + attached right now (it keeps serving, and its own disconnect must be + able to detach this lease); else nobody, in which case the generation + is returned and the caller puts it back into disconnect grace. + Generations only move forward, so an older orphan never displaces a + newer one an owner already holds. + """ + async with self._lock: + state = self._sessions.get(session_id) + if state is None: + # The session is gone from the registry: the engine session + # closed with it, nothing is left to own. + return None + state.orphaned_lease_generation = max(state.orphaned_lease_generation or 0, lease_generation) + return self._settle_orphan_locked(state) + + @staticmethod + def _settle_orphan_locked(state: _DuplexSessionAttachmentState) -> int | None: + orphan = state.orphaned_lease_generation + if orphan is None: + return None + if state.pending_resumes > 0: + # The next activation absorbs it (see ``resume``). + return None + state.orphaned_lease_generation = None + if state.attachment is not None: + state.lease_generation = max(state.lease_generation, orphan) + return None + return orphan async def has_attachment(self, session_id: str) -> bool: """Whether some connection is attached right now, whoever it is. @@ -362,6 +465,7 @@ class DuplexSessionAttachmentRegistry: send: Callable[[dict[str, object]], Awaitable[None]], close: Callable[[str], Awaitable[None]], activation_payload_factory: Callable[[ResumeToken, int], Mapping[str, object]] | None = None, + lease_generation: int | None = None, ) -> DuplexSessionResumeResult: async with self._lock: state = self._require(session_id) @@ -384,6 +488,15 @@ class DuplexSessionAttachmentRegistry: send=send, close=close, ) + # The activating connection serves the newest lease the engine + # holds: its own resume's generation, or a newer one whose + # resume was abandoned while this one was pending. + state.lease_generation = max( + state.lease_generation, + lease_generation if lease_generation is not None else 0, + state.orphaned_lease_generation or 0, + ) + state.orphaned_lease_generation = None if activation_payload_factory is not None: try: await send(dict(activation_payload_factory(rotated_token, attachment_generation))) @@ -418,18 +531,26 @@ class DuplexSessionAttachmentRegistry: *, state: _DuplexSessionAttachmentState, attachment_generation: int, - accepted_token_digest: bytes, + accepted_token_digest: bytes | None, ) -> None: """Undo a resume whose activation never reached the client. Generation-checked: a later resume that already took over must keep its attachment. Restoring ``recovery_token_digest`` is what lets the caller retry with the token it presented, which ``resume`` had rotated past. + + The lease the provisional attachment was serving loses its owner here. + That is not only its own resume's generation: a resume abandoned while + this one was activating may have handed it a newer one. It is parked as + the orphan, so the caller's settlement (which only knows its own + generation) and ``end_resume`` give back the newest lease, never a + stale one the engine would refuse. """ async with self._lock: if self._sessions.get(session_id) is state and state.attachment_generation == attachment_generation: state.attachment = None state.recovery_token_digest = accepted_token_digest + state.orphaned_lease_generation = max(state.orphaned_lease_generation or 0, state.lease_generation) async def close(self, session_id: str) -> DuplexTransportAttachment | None: async with self._lock: diff --git a/vllm_omni/entrypoints/duplex_omni.py b/vllm_omni/entrypoints/duplex_omni.py index 968b69f74..e5d8e5bf6 100644 --- a/vllm_omni/entrypoints/duplex_omni.py +++ b/vllm_omni/entrypoints/duplex_omni.py @@ -26,7 +26,7 @@ Example:: from __future__ import annotations import asyncio -from collections.abc import AsyncIterator, Mapping, Sequence +from collections.abc import AsyncIterator, Awaitable, Callable, Mapping, Sequence from contextlib import suppress from typing import Any from uuid import uuid4 @@ -286,8 +286,13 @@ class DuplexOmni(AsyncOmni): return DuplexOmniEngine(duplex_audio_encoder=encode_audio, **engine_kwargs) def __init__(self, model: str = "", *args: Any, **kwargs: Any) -> None: - super().__init__(*args, model=model, **kwargs) + super().__init__(model, *args, **kwargs) self._handles: dict[str, DuplexSessionHandle] = {} + #: Rollbacks of resumes whose caller was cancelled mid-RPC (see ``resume_session``). + self._resume_compensations: set[asyncio.Task[None]] = set() + #: Initial and maximum delay between replays of an abandoned resume + #: whose answers keep timing out. + self._resume_replay_backoff_s: tuple[float, float] = (0.5, 5.0) # ---- deployment facts ---- @@ -377,32 +382,166 @@ class DuplexOmni(AsyncOmni): *, expected_lease_generation: int, timeout: float | None = _DEFAULT_CONTROL_TIMEOUT_S, + on_abandoned: Callable[[int], Awaitable[None]] | None = None, ) -> DuplexSessionHandle: - """Engine lease resume (CAS on the lease generation); returns the existing handle.""" + """Engine lease resume (CAS on the lease generation); returns the existing handle. + + The RPC runs to completion in an executor thread whatever happens to + this waiter, so a caller cancelled while it is in flight does not stop + the engine from applying the resume (lease generation bumped, + disconnect grace cleared) for a connection that will never serve it. + The outcome is then observed off the cancelled task. If the RPC times + out before answering, the outcome is unknown, so the same resume is + replayed under its control id: the engine answers a replay with the + generation the resume produced instead of resuming again. A resume + that landed is adopted into the handle (a later reconnect resumes + against the real generation) and settled through ``on_abandoned`` with + that generation; without a callback it is detached again, fenced on + the generation, so a resume that came after it is never touched. + + ``on_abandoned`` is for callers that own an attachment concept: they + decide whether the generation now belongs to a connection that is + still serving (a takeover that never activated leaves the previous + socket attached) or whether the lease goes back into disconnect grace. + """ handle = self._require_handle(session_id) - result = await self.engine.resume_session_async( - session_id, - expected_lease_generation=expected_lease_generation, - timeout=timeout, + control_id = uuid4().hex + rpc = asyncio.ensure_future( + self.engine.resume_session_async( + session_id, + expected_lease_generation=expected_lease_generation, + control_id=control_id, + timeout=timeout, + ) ) + try: + result = await asyncio.shield(rpc) + except asyncio.CancelledError: + self._compensate_abandoned_resume( + handle, + rpc, + control_id=control_id, + expected_lease_generation=expected_lease_generation, + timeout=timeout, + on_abandoned=on_abandoned, + ) + raise # The result carries the engine's current public session (state, # epoch, turn ...), which a resumed client must see, not the open-time snapshot. handle._adopt(result) return handle + def _compensate_abandoned_resume( + self, + handle: DuplexSessionHandle, + rpc: asyncio.Future[DuplexControlResultMessage], + *, + control_id: str, + expected_lease_generation: int, + timeout: float | None, + on_abandoned: Callable[[int], Awaitable[None]] | None, + ) -> None: + session_id = handle.session_id + + async def observe() -> DuplexControlResultMessage | None: + # The first attempt is the RPC the caller abandoned. A timeout + # means its answer was dropped and the outcome is unknown: the + # engine may have applied the resume after the waiter gave up. + # Replaying under the same control id makes the engine answer with + # the generation that resume produced, whether it lands now or + # already did, and refuse if a newer resume won. A replay can time + # out too, so this keeps replaying, with backoff, until the engine + # answers one way or the other or the session is gone: giving up + # on a timeout would be exactly the lost-answer window again. + attempt: Awaitable[DuplexControlResultMessage] = rpc + delay, ceiling = self._resume_replay_backoff_s + while True: + try: + return await attempt + except DuplexSessionError as exc: + if exc.code != "timeout": + # Refused (conflict, unknown session, engine gone): + # the resume did not land, nothing to give back. + logger.debug("abandoned duplex resume of %s did not land: %s", session_id, exc) + return None + except Exception as exc: + logger.debug("abandoned duplex resume of %s did not land: %s", session_id, exc) + return None + if handle.closed: + return None + logger.warning( + "abandoned duplex resume of %s lost its answer to a timeout; replaying in %.1fs", session_id, delay + ) + await asyncio.sleep(delay) + delay = min(delay * 2, ceiling) + attempt = self.engine.resume_session_async( + session_id, + expected_lease_generation=expected_lease_generation, + control_id=control_id, + timeout=timeout, + ) + + async def compensate() -> None: + result = await observe() + if result is None: + return + handle._adopt(result) + lease_generation = result.lease_generation + if lease_generation is None: + return + try: + if on_abandoned is not None: + await on_abandoned(lease_generation) + else: + await self.engine.touch_session_async( + session_id, + activity="detach", + expected_lease_generation=lease_generation, + timeout=timeout, + ) + except Exception as exc: + logger.warning("abandoned duplex resume of %s could not be settled: %s", session_id, exc) + + task = asyncio.create_task(compensate(), name=f"duplex-resume-compensation-{session_id}") + self._resume_compensations.add(task) + task.add_done_callback(self._resume_compensations.discard) + async def touch_session( self, session_id: str, *, activity: str = "heartbeat", + expected_lease_generation: int | None = None, timeout: float | None = _DEFAULT_CONTROL_TIMEOUT_S, ) -> None: self._require_handle(session_id) - await self.engine.touch_session_async(session_id, activity=activity, timeout=timeout) + await self.engine.touch_session_async( + session_id, + activity=activity, + expected_lease_generation=expected_lease_generation, + timeout=timeout, + ) - async def detach_session(self, session_id: str, *, timeout: float | None = _DEFAULT_CONTROL_TIMEOUT_S) -> None: - """Start the engine-owned disconnect grace; expiry arrives as ``session.expired``.""" - await self.touch_session(session_id, activity="detach", timeout=timeout) + async def detach_session( + self, + session_id: str, + *, + expected_lease_generation: int | None = None, + timeout: float | None = _DEFAULT_CONTROL_TIMEOUT_S, + ) -> None: + """Start the engine-owned disconnect grace; expiry arrives as ``session.expired``. + + ``expected_lease_generation`` names the lease the caller opened or + resumed against; the engine refuses to detach a newer one, so a + connection giving up its lease cannot start the grace for the lease a + later resume owns. + """ + await self.touch_session( + session_id, + activity="detach", + expected_lease_generation=expected_lease_generation, + timeout=timeout, + ) async def close_session( self, @@ -470,6 +609,8 @@ class DuplexOmni(AsyncOmni): for handle in list(self._handles.values()): handle._mark_closed("shutdown") self._handles.clear() + for compensation in list(self._resume_compensations): + compensation.cancel() super().shutdown(timeout)