rpc interceptors: bound the unwind after outside cancellation; await any awaitable a handler returns

Awaiting the chain directly forwarded an outside cancel (the room disconnecting)
into the chain and kept the invocation task parked until the chain finished, so
a handler that swallowed cancellation held room.disconnect() up to the caller's
deadline. The chain is now awaited through a shield: the cancel is raised here
at once, the chain is cancelled explicitly and given a bounded time to unwind,
and a handler that does not stop is named in a warning.

Handlers are typed as any callable returning a payload or an awaitable of one;
dispatch now checks the result with inspect.isawaitable instead of
iscoroutinefunction, which missed callable objects with an async __call__ and
sync wrappers returning a coroutine.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

Author: davidzhao

livekit-rtc/livekit/rtc/participant.py

@@ -16,6 +16,7 @@
 
 import ctypes
 import asyncio
+import inspect
 import datetime
 import enum
 import os
@@ -237,6 +238,10 @@ def disconnect_reason(
 
 RpcHandler = Callable[["RpcInvocationData"], Union[Awaitable[Optional[str]], Optional[str]]]
 
+# how long a cancelled incoming RPC chain gets to unwind before the caller is answered
+# without it (the room's disconnect waits on these invocations)
+_RPC_CANCEL_UNWIND_TIMEOUT = 2.0
+
 
 F = TypeVar(
     "F", bound=Callable[[RpcInvocationData], Union[Awaitable[Optional[str]], Optional[str]]]
@@ -658,14 +663,25 @@ def _on_deadline() -> None:
 
         deadline = loop.call_later(invocation.response_timeout, _on_deadline)
         try:
-            return await chain_task
+            # shielded: a cancel from outside (the room disconnecting) is raised here at
+            # once. Awaiting the chain directly would instead forward the cancel to it and
+            # keep this task parked until the chain finished, so a handler that ignored
+            # cancellation held up room.disconnect() for as long as the caller's deadline.
+            return await asyncio.shield(chain_task)
         except asyncio.CancelledError:
             if deadline_fired:
                 raise RpcError._built_in(RpcError.ErrorCode.RESPONSE_TIMEOUT) from None
-            # cancelled from outside: awaiting propagated the cancel into the chain; let it
-            # finish unwinding before answering the caller
-            if not chain_task.done():
-                await asyncio.wait([chain_task])
+            # cancelled from outside: stop the chain and let it unwind before answering the
+            # caller, but not for long; this is the path room.disconnect() waits on
+            chain_task.cancel()
+            _, pending = await asyncio.wait([chain_task], timeout=_RPC_CANCEL_UNWIND_TIMEOUT)
+            if pending:
+                logger.warning(
+                    "RPC handler for %s did not stop within %.1fs of being cancelled; "
+                    "answering the caller without it",
+                    invocation.method,
+                    _RPC_CANCEL_UNWIND_TIMEOUT,
+                )
             raise RpcError._built_in(RpcError.ErrorCode.RECIPIENT_DISCONNECTED) from None
         except Exception:
             if deadline_fired:
@@ -686,9 +702,13 @@ async def _invoke_rpc_handler(self, invocation: RpcInvocationData) -> Optional[s
         if not handler:
             raise RpcError._built_in(RpcError.ErrorCode.UNSUPPORTED_METHOD)
 
-        if asyncio.iscoroutinefunction(handler):
-            return cast(Optional[str], await handler(invocation))
-        return cast(Optional[str], handler(invocation))
+        # RpcHandler admits any callable returning a payload or an awaitable of one: a
+        # coroutine function, but also a callable object with an async __call__ or a sync
+        # wrapper handing back a coroutine, which iscoroutinefunction would not recognize
+        result = handler(invocation)
+        if inspect.isawaitable(result):
+            result = await result
+        return result
 
     async def set_metadata(self, metadata: str) -> None:
         """

livekit-rtc/tests/test_rpc_interceptors.py

@@ -321,6 +321,73 @@ async def hang(data: RpcInvocationData) -> str:
     assert info.value.code == rtc.RpcError.ErrorCode.RECIPIENT_DISCONNECTED
 
 
+async def test_unwind_after_outside_cancellation_is_bounded(
+    monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture
+) -> None:
+    """A handler that ignores cancellation must not hang the disconnect that cancelled it:
+    the caller is answered once the unwind bound passes, and the offender is named.
+
+    Awaiting the chain directly could never bound this: asyncio forwards the cancel to the
+    chain and keeps the awaiting task parked until the chain finishes, so the wait would
+    only begin once the stubborn handler had already stopped (at the caller's deadline)."""
+    from livekit.rtc import participant as participant_mod
+
+    monkeypatch.setattr(participant_mod, "_RPC_CANCEL_UNWIND_TIMEOUT", 0.05)
+    lp = _participant()
+    started = asyncio.Event()
+    release = asyncio.Event()
+
+    async def ignores_cancel(data: RpcInvocationData) -> str:
+        started.set()
+        try:
+            await asyncio.sleep(10)
+        except asyncio.CancelledError:
+            pass
+        await release.wait()  # keeps running long after it was told to stop
+        return "late"
+
+    lp._rpc_handlers["stubborn"] = ignores_cancel
+    task = asyncio.ensure_future(
+        lp._run_incoming_chain(RpcInvocationData("r1", "alice", "{}", 5.0, method="stubborn"))
+    )
+    await started.wait()
+    task.cancel()
+    loop = asyncio.get_running_loop()
+    t0 = loop.time()
+    with pytest.raises(rtc.RpcError) as info:
+        await task
+    assert info.value.code == rtc.RpcError.ErrorCode.RECIPIENT_DISCONNECTED
+    assert loop.time() - t0 < 1.0
+    assert any("stubborn" in r.getMessage() for r in caplog.records)
+
+    release.set()  # let the orphaned handler finish so the loop closes clean
+    await asyncio.sleep(0)
+
+
+async def test_handlers_returning_an_awaitable_are_awaited() -> None:
+    """RpcHandler admits any callable returning a payload or an awaitable of one, not only
+    coroutine functions: a callable object with an async __call__, a sync wrapper handing
+    back a coroutine, and a plain sync handler all work."""
+    lp = _participant()
+
+    class Handler:
+        async def __call__(self, data: RpcInvocationData) -> str:
+            return f"obj:{data.payload}"
+
+    async def _inner(data: RpcInvocationData) -> str:
+        return f"wrapped:{data.payload}"
+
+    lp._rpc_handlers["obj"] = Handler()
+    lp._rpc_handlers["wrapped"] = lambda data: _inner(data)  # sync, returns a coroutine
+    lp._rpc_handlers["sync"] = lambda data: f"sync:{data.payload}"
+
+    for method, expected in (("obj", "obj:x"), ("wrapped", "wrapped:x"), ("sync", "sync:x")):
+        result = await lp._run_incoming_chain(
+            RpcInvocationData("r1", "alice", "x", 1.0, method=method)
+        )
+        assert result == expected, method
+
+
 async def test_add_and_remove_interceptors() -> None:
     lp = _participant()
     log: list[str] = []