[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:
19 files changed
+1697
-96
No files matched your search
@@ -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
@@ -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
|
||||
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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"
|
||||
@@ -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)
|
||||
@@ -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"
|
||||
@@ -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()
|
||||
@@ -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):
|
||||
|
||||
@@ -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(),
|
||||
}
|
||||
)
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
|
||||
Reference in new issue
Block a user