[Bugfix] Compensate cancelled duplex open, append and resume; fix serving README drift (#7636 Issues 1, 2, 3, 21) (#7785)

Signed-off-by: chickeyton <ngton2014@gmail.com>
Signed-off-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com>
Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>
Co-authored-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com>
This commit is contained in:
authored and GitHub committed 2026-09-24 19:23:08 -07:00
1 parent 551a03102f
commit 9e7e3a802c
19 files changed
+1697 -96

No files matched your search

+7 -2
View File
@@ -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
+16 -11
View File
@@ -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
+26
View File
@@ -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)
+194 -4
View File
@@ -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"
+138 -2
View File
@@ -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)
+422 -3
View File
@@ -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"
@@ -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"
+175 -3
View File
@@ -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()
+5
View File
@@ -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):
+59 -11
View File
@@ -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(),
}
)
+21 -2
View File
@@ -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:
@@ -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)
+13 -1
View File
@@ -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
+71 -15
View File
@@ -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)
@@ -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,
)
+15 -2
View File
@@ -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:
+168 -23
View File
@@ -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
@@ -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:
+152 -11
View File
@@ -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)