[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>

tests/entrypoints/duplex/test_duplex_serving.py

@@ -8,6 +8,7 @@ from __future__ import annotations
 import asyncio
 import base64
 import json
+from collections.abc import Awaitable, Callable
 from contextlib import suppress
 from typing import Any
 
@@ -131,7 +132,15 @@ class FakeOmni:
         self.handles: dict[str, FakeHandle] = {}
         self.resumed: list[tuple[str, int]] = []
         self.detached: list[str] = []
+        #: Detaches the engine refused because the caller's lease generation was stale.
+        self.detach_refused: list[tuple[str, int]] = []
         self.open_error: DuplexSessionError | None = None
+        #: When set, ``resume_session`` parks on it once the engine applied the
+        #: resume: the lease generation is bumped, the result is still in flight.
+        self.resume_gate: asyncio.Event | None = None
+        self.resume_started = asyncio.Event()
+        #: Settlements of resumes whose caller was cancelled mid-RPC (DuplexOmni's contract).
+        self.compensations: list[asyncio.Task[None]] = []
 
     async def open_session(self, config: Any) -> FakeHandle:
         self.opened.append(dict(config))
@@ -146,13 +155,46 @@ class FakeOmni:
     def get_session(self, session_id: str) -> FakeHandle | None:
         return self.handles.get(session_id)
 
-    async def resume_session(self, session_id: str, *, expected_lease_generation: int) -> FakeHandle:
+    async def resume_session(
+        self,
+        session_id: str,
+        *,
+        expected_lease_generation: int,
+        on_abandoned: Callable[[int], Awaitable[None]] | None = None,
+    ) -> FakeHandle:
         self.resumed.append((session_id, expected_lease_generation))
         handle = self.handles[session_id]
-        handle.lease_generation = expected_lease_generation + 1
+        lease_generation = expected_lease_generation + 1
+        handle.lease_generation = lease_generation
+        if self.resume_gate is not None:
+            self.resume_started.set()
+            gate = self.resume_gate
+            try:
+                await gate.wait()
+            except asyncio.CancelledError:
+                # DuplexOmni's contract: the engine applied the resume; once
+                # its answer arrives the landed generation is settled off the
+                # cancelled task, through the caller's callback if it gave one.
+                async def settle() -> None:
+                    await gate.wait()
+                    if on_abandoned is not None:
+                        await on_abandoned(lease_generation)
+                    else:
+                        with suppress(DuplexSessionError):
+                            await self.detach_session(session_id, expected_lease_generation=lease_generation)
+
+                self.compensations.append(asyncio.create_task(settle()))
+                raise
         return handle
 
-    async def detach_session(self, session_id: str) -> None:
+    async def detach_session(self, session_id: str, *, expected_lease_generation: int | None = None) -> None:
+        handle = self.handles[session_id]
+        if expected_lease_generation is not None and expected_lease_generation != handle.lease_generation:
+            # The engine's fence: a detach of a lease the caller no longer holds is refused.
+            self.detach_refused.append((session_id, expected_lease_generation))
+            raise DuplexSessionError(
+                "duplex lease generation mismatch", code="session_resume_conflict", session_id=session_id
+            )
         self.detached.append(session_id)
 
 
@@ -607,3 +649,380 @@ async def test_a_resume_that_fails_to_activate_does_not_detach_the_live_attachme
                 pending.cancel()
                 with suppress(asyncio.CancelledError, Exception):
                     await pending
+
+
+# --------------------------------------------------------------------------- #
+# Cancelled resume and the engine disconnect grace (#7636 Issue 3)            #
+# --------------------------------------------------------------------------- #
+
+
+@pytest.mark.asyncio
+async def test_a_resume_cancelled_during_activation_puts_the_engine_lease_back_into_grace() -> None:
+    """``_resume`` resumes the engine lease first; a cancel mid-activation must detach it again.
+
+    The rollback used to catch ``Exception`` only. A handler task cancelled
+    while sending ``session.resumed`` or a replay entry left the engine with
+    ``detached_at=None`` and the outer handler with no attachment to clean up:
+    the reaper reclaimed nothing after the grace, and with one slot the next
+    open was refused with ``resource_exhausted``.
+    """
+    omni = FakeOmni()
+    handler = _handler(omni)
+    ws, handle, task = await _open(handler, omni)
+    token = ws.sent[0]["resume_token"]
+    ws.disconnect()
+    await asyncio.wait_for(task, timeout=2.0)
+    assert omni.detached == [handle.session_id]
+    registry = handler._attachment_registry
+
+    activating = asyncio.Event()
+
+    async def parked_resume(*args, **kwargs):
+        activating.set()
+        await asyncio.Event().wait()
+
+    original_resume = registry.resume
+    registry.resume = parked_resume  # type: ignore[method-assign]
+    try:
+        ws2, task2 = await _resume(handler, handle.session_id, token)
+        await asyncio.wait_for(activating.wait(), timeout=2.0)
+        assert omni.resumed == [(handle.session_id, 0)], "the engine lease was resumed before activation"
+
+        task2.cancel()
+        with suppress(asyncio.CancelledError):
+            await task2
+    finally:
+        registry.resume = original_resume  # type: ignore[method-assign]
+
+    # The lease is detached again, so the disconnect grace runs for it.
+    assert omni.detached == [handle.session_id, handle.session_id]
+    assert not await registry.has_attachment(handle.session_id)
+    assert handle.close_reasons == []
+    # And the session is still resumable with the token the cancelled attempt presented.
+    ws3, task3 = await _resume(handler, handle.session_id, token)
+    resumed = await ws3.wait_for("session.resumed")
+    assert resumed["session_id"] == handle.session_id
+    ws3.disconnect()
+    await asyncio.wait_for(task3, timeout=2.0)
+
+
+@pytest.mark.asyncio
+async def test_a_resume_cancelled_after_activation_detaches_the_attachment_it_made() -> None:
+    """A cancel after the registry activated the new attachment must undo that attachment too.
+
+    Otherwise the session stays attached to a socket nobody serves, with the
+    engine lease resumed: no grace, no expiry, until the idle TTL.
+    """
+    omni = FakeOmni()
+    handler = _handler(omni)
+    ws, handle, task = await _open(handler, omni)
+    token = ws.sent[0]["resume_token"]
+    registry = handler._attachment_registry
+
+    # The takeover notifies the replaced socket after activation; park there.
+    notifying = asyncio.Event()
+
+    async def parked_send(payload: dict[str, Any]) -> None:
+        notifying.set()
+        await asyncio.Event().wait()
+
+    ws.send_json = parked_send  # type: ignore[method-assign]
+    ws2, task2 = await _resume(handler, handle.session_id, token)
+    await asyncio.wait_for(notifying.wait(), timeout=2.0)
+    assert await registry.is_current_attachment(handle.session_id, 2), "the new socket is attached"
+
+    task2.cancel()
+    with suppress(asyncio.CancelledError):
+        await task2
+
+    assert not await registry.has_attachment(handle.session_id)
+    assert omni.detached == [handle.session_id]
+    assert handle.close_reasons == []
+
+    task.cancel()
+    with suppress(asyncio.CancelledError, Exception):
+        await task
+
+
+@pytest.mark.asyncio
+async def test_an_abandoned_resume_never_detaches_the_lease_a_later_resume_owns() -> None:
+    """Two connections resuming the same session: the loser's rollback must not touch the winner's lease.
+
+    Interleaving: A has activated and is notifying the socket it replaced; B,
+    holding A's rotated token, has completed its engine resume but its RPC
+    result has not reached the handler yet, so nothing is attached for B.
+    Cancelling A detaches A's registry generation, which used to make the
+    engine detach unconditional: it landed on the lease B had just resumed,
+    and the reaper expired B after the grace even though it heartbeats. The
+    engine now refuses a detach fenced on a generation it has moved past.
+    """
+    omni = FakeOmni()
+    handler = _handler(omni)
+    ws, handle, task = await _open(handler, omni)
+    token = ws.sent[0]["resume_token"]
+    registry = handler._attachment_registry
+
+    # A takes over and parks while notifying the replaced socket (after activation).
+    notifying = asyncio.Event()
+
+    async def parked_send(payload: dict[str, Any]) -> None:
+        notifying.set()
+        await asyncio.Event().wait()
+
+    ws.send_json = parked_send  # type: ignore[method-assign]
+    ws_a, task_a = await _resume(handler, handle.session_id, token)
+    await asyncio.wait_for(notifying.wait(), timeout=2.0)
+    resumed_a = await ws_a.wait_for("session.resumed")
+    assert handle.lease_generation == 1
+
+    # B resumes with A's rotated token: the engine applies it, the result is still in flight.
+    omni.resume_gate = asyncio.Event()
+    ws_b, task_b = await _resume(handler, handle.session_id, resumed_a["resume_token"])
+    await asyncio.wait_for(omni.resume_started.wait(), timeout=2.0)
+    assert handle.lease_generation == 2
+
+    task_a.cancel()
+    with suppress(asyncio.CancelledError):
+        await task_a
+    assert omni.detached == [], "A's rollback must not detach the lease B resumed"
+    assert omni.detach_refused == [(handle.session_id, 1)]
+    assert not await registry.has_attachment(handle.session_id), "A's attachment is gone"
+
+    omni.resume_gate.set()
+    resumed_b = await ws_b.wait_for("session.resumed")
+    assert resumed_b["attachment_generation"] == 3
+    assert await registry.is_current_attachment(handle.session_id, 3)
+    assert omni.detached == []
+
+    # B's own disconnect detaches the lease B holds.
+    ws_b.disconnect()
+    await asyncio.wait_for(task_b, timeout=2.0)
+    assert omni.detached == [handle.session_id]
+    assert omni.detach_refused == [(handle.session_id, 1)]
+    task.cancel()
+    with suppress(asyncio.CancelledError, Exception):
+        await task
+
+
+@pytest.mark.asyncio
+async def test_an_abandoned_takeover_leaves_the_live_connection_serving_and_owning_the_lease() -> None:
+    """A takeover cancelled mid-RPC must not put the connection it never replaced into disconnect grace.
+
+    A is attached at generation 0. B's engine resume lands generation 1 and B
+    is cancelled before it reaches the registry, so A stays the current
+    attachment. Settling B's generation must hand it to A rather than detach
+    it: A keeps serving, and A's own disconnect later detaches generation 1,
+    which is the lease the engine actually holds.
+    """
+    omni = FakeOmni()
+    handler = _handler(omni)
+    ws, handle, task = await _open(handler, omni)
+    token = ws.sent[0]["resume_token"]
+    registry = handler._attachment_registry
+
+    omni.resume_gate = asyncio.Event()
+    ws_b, task_b = await _resume(handler, handle.session_id, token)
+    await asyncio.wait_for(omni.resume_started.wait(), timeout=2.0)
+    assert handle.lease_generation == 1
+
+    task_b.cancel()
+    with suppress(asyncio.CancelledError):
+        await task_b
+    omni.resume_gate.set()
+    await asyncio.gather(*omni.compensations)
+
+    assert omni.detached == [], "A is still serving: its lease must not enter disconnect grace"
+    assert omni.detach_refused == []
+    assert await registry.is_current_attachment(handle.session_id, 1), "A is still the attachment"
+    handle.deliver(AudioDelta(session_id=handle.session_id, response_id="r1", delta="aGk="))
+    await ws.wait_for("response.output_audio.delta")
+
+    # A now owns generation 1: its disconnect detaches the lease the engine holds.
+    ws.disconnect()
+    await asyncio.wait_for(task, timeout=2.0)
+    assert omni.detached == [handle.session_id]
+    assert omni.detach_refused == []
+
+
+@pytest.mark.asyncio
+async def test_a_send_failure_detaches_the_failed_sockets_own_lease_not_a_pending_resumes() -> None:
+    """The pump's detach is fenced on the lease of the socket whose send failed, captured with it.
+
+    The pump holds the registry's outbound lock while sending to A. B's
+    engine resume lands generation 1 meanwhile and B waits for that lock to
+    activate. When A's send fails, the pump drops A and must detach A's
+    generation 0, which the engine refuses, not B's generation 1, which would
+    put B into disconnect grace the moment it activates.
+    """
+    omni = FakeOmni()
+    handler = _handler(omni)
+    ws, handle, task = await _open(handler, omni)
+    token = ws.sent[0]["resume_token"]
+    registry = handler._attachment_registry
+
+    sending = asyncio.Event()
+    fail_send = asyncio.Event()
+
+    async def parked_failing_send(payload: dict[str, Any]) -> None:
+        sending.set()
+        await fail_send.wait()
+        raise RuntimeError("socket closed")
+
+    ws.send_json = parked_failing_send  # type: ignore[method-assign]
+    handle.deliver(AudioDelta(session_id=handle.session_id, response_id="r1", delta="aGk="))
+    await asyncio.wait_for(sending.wait(), timeout=2.0)
+
+    # B: the engine resume lands; activation waits for the outbound lock the pump holds.
+    ws_b, task_b = await _resume(handler, handle.session_id, token)
+    await asyncio.sleep(0.05)
+    assert omni.resumed == [(handle.session_id, 0)]
+    assert handle.lease_generation == 1
+    assert await registry.is_current_attachment(handle.session_id, 1), "B has not activated yet"
+
+    fail_send.set()
+    resumed_b = await ws_b.wait_for("session.resumed")
+    assert resumed_b["attachment_generation"] == 2
+    assert omni.detached == [], "A's lease (generation 0) is refused, B's lease is untouched"
+    assert omni.detach_refused == [(handle.session_id, 0)]
+    assert await registry.is_current_attachment(handle.session_id, 2)
+
+    # B serves; its own disconnect detaches the lease it resumed.
+    ws_b.disconnect()
+    await asyncio.wait_for(task_b, timeout=2.0)
+    assert omni.detached == [handle.session_id]
+    ws.disconnect()
+    with suppress(asyncio.CancelledError, Exception):
+        await asyncio.wait_for(task, timeout=2.0)
+
+
+@pytest.mark.asyncio
+async def test_an_abandoned_resume_goes_to_the_resume_waiting_to_activate_not_the_socket_it_replaces() -> None:
+    """Three connections: A sending, B resumed and waiting for the outbound lock, C abandoned mid-RPC.
+
+    C's generation 2 must be parked for B's activation. If it were handed to
+    A instead, A's send failure would detach generation 2 (accepted by the
+    engine), and B would then activate on a lease already in disconnect grace
+    and expire despite heartbeating.
+    """
+    omni = FakeOmni()
+    handler = _handler(omni)
+    ws, handle, task = await _open(handler, omni)
+    token = ws.sent[0]["resume_token"]
+    registry = handler._attachment_registry
+
+    sending = asyncio.Event()
+    fail_send = asyncio.Event()
+
+    async def parked_failing_send(payload: dict[str, Any]) -> None:
+        sending.set()
+        await fail_send.wait()
+        raise RuntimeError("socket closed")
+
+    ws.send_json = parked_failing_send  # type: ignore[method-assign]
+    handle.deliver(AudioDelta(session_id=handle.session_id, response_id="r1", delta="aGk="))
+    await asyncio.wait_for(sending.wait(), timeout=2.0)
+
+    # B: engine resume lands generation 1; activation waits for the lock the pump holds.
+    ws_b, task_b = await _resume(handler, handle.session_id, token)
+    await asyncio.sleep(0.05)
+    assert handle.lease_generation == 1
+
+    # C: same unrotated token, engine resume lands generation 2, cancelled mid-RPC.
+    omni.resume_gate = asyncio.Event()
+    ws_c, task_c = await _resume(handler, handle.session_id, token)
+    await asyncio.wait_for(omni.resume_started.wait(), timeout=2.0)
+    assert handle.lease_generation == 2
+    task_c.cancel()
+    with suppress(asyncio.CancelledError):
+        await task_c
+    omni.resume_gate.set()
+    await asyncio.gather(*omni.compensations)
+    assert omni.detached == [] and omni.detach_refused == [], "generation 2 is parked for B"
+
+    # A's send fails: the pump drops A and detaches A's own generation 0, refused.
+    fail_send.set()
+    resumed_b = await ws_b.wait_for("session.resumed")
+    assert resumed_b["attachment_generation"] == 2
+    assert omni.detach_refused == [(handle.session_id, 0)]
+    assert omni.detached == [], "the lease B serves is untouched"
+    assert await registry.is_current_attachment(handle.session_id, 2)
+
+    # B serves generation 2 (inherited from C): its disconnect detaches that lease.
+    ws_b.disconnect()
+    await asyncio.wait_for(task_b, timeout=2.0)
+    assert omni.detached == [handle.session_id]
+    assert omni.detach_refused == [(handle.session_id, 0)]
+    ws.disconnect()
+    with suppress(asyncio.CancelledError, Exception):
+        await asyncio.wait_for(task, timeout=2.0)
+
+
+@pytest.mark.asyncio
+async def test_a_resume_cancelled_during_replay_settles_the_newer_lease_it_was_handed() -> None:
+    """Activation rollback must give back the lease the provisional attachment actually owned.
+
+    A resumed generation 1, delivered ``session.resumed`` and is blocked
+    replaying events. B, with A's rotated token, resumed generation 2 and was
+    cancelled mid-RPC; its settlement handed 2 to A. When A is then cancelled
+    during replay, the registry rolls A back; detaching A's own generation 1
+    would be refused and generation 2 would never enter disconnect grace.
+    """
+    omni = FakeOmni()
+    handler = _handler(omni)
+    ws, handle, task = await _open(handler, omni)
+    token = ws.sent[0]["resume_token"]
+    registry = handler._attachment_registry
+    ws.disconnect()
+    await asyncio.wait_for(task, timeout=2.0)
+    assert omni.detached == [handle.session_id]
+    # Journaled while detached: A's activation has something to replay.
+    handle.deliver(AudioDelta(session_id=handle.session_id, response_id="r1", delta="one"))
+    await asyncio.sleep(0.05)
+
+    # A: session.resumed goes out, the replay entry parks.
+    ws_a = FakeWebSocket({"duplex": "1", "resume": "1"})
+    replaying = asyncio.Event()
+    sends = 0
+    original_send = ws_a.send_json
+
+    async def send_parking_on_replay(payload: dict[str, Any]) -> None:
+        nonlocal sends
+        sends += 1
+        if sends == 2:
+            replaying.set()
+            await asyncio.Event().wait()
+        await original_send(payload)
+
+    ws_a.send_json = send_parking_on_replay  # type: ignore[method-assign]
+    task_a = asyncio.create_task(handler.handle_realtime_session(ws_a))
+    ws_a.feed(
+        {
+            "type": "session.resume",
+            "session_id": handle.session_id,
+            "resume_token": token,
+            "last_received_server_event_seq": 0,
+        }
+    )
+    await asyncio.wait_for(replaying.wait(), timeout=2.0)
+    resumed_a = await ws_a.wait_for("session.resumed")
+    assert handle.lease_generation == 1
+
+    # B: A's rotated token, engine resume lands generation 2, cancelled mid-RPC.
+    omni.resume_gate = asyncio.Event()
+    ws_b, task_b = await _resume(handler, handle.session_id, resumed_a["resume_token"])
+    await asyncio.wait_for(omni.resume_started.wait(), timeout=2.0)
+    assert handle.lease_generation == 2
+    task_b.cancel()
+    with suppress(asyncio.CancelledError):
+        await task_b
+    omni.resume_gate.set()
+    await asyncio.gather(*omni.compensations)
+    assert omni.detached == [handle.session_id], "generation 2 went to A's pending activation"
+
+    # A is cancelled during replay: its rollback settles generation 2, the lease A owned.
+    task_a.cancel()
+    with suppress(asyncio.CancelledError):
+        await task_a
+    assert not await registry.has_attachment(handle.session_id)
+    assert omni.detached == [handle.session_id, handle.session_id]
+    assert omni.detach_refused == [], "generation 1 was never detached: the engine would refuse it"