rpc interceptors: enforce the deadline across the chain, register by identity

Review follow-ups on the interceptor hook:

- The caller's response_timeout now covers the whole incoming chain
  (_run_incoming_chain wraps interceptors + handler in one wait_for), so time
  an interceptor spends before or after next() counts against it instead of
  handing the handler a fresh full timeout. Deadline expiry cancels the chain
  and maps to RESPONSE_TIMEOUT; cancellation from outside (room disconnect)
  still maps to RECIPIENT_DISCONNECTED. _invoke_rpc_handler no longer wraps
  the handler itself.
- add_rpc_interceptor / remove_rpc_interceptor compare by identity, so two
  distinct interceptors that compare equal coexist and only the exact instance
  is removed.
- mypy: cast the awaited handler result; fix the non-overlapping identity
  check in the test.

Tests: deadline burned before and after next(), outside cancellation,
identity registration; the timeout test now asserts interceptors observe the
cancellation.

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

Author: davidzhao

README.md

@@ -225,7 +225,7 @@ You may find it useful to adjust the `response_timeout` parameter, which indicat
 
 #### Intercepting RPC calls
 
-An `RpcInterceptor` wraps every RPC the local participant performs or handles, which is useful for logging, tracing, or attaching metadata to payloads. Each method receives the call and a `next` continuation; return what `next` returns. Interceptors run in the order they were added, the first being outermost, and errors from the remote side or from your handler flow through them unchanged.
+An `RpcInterceptor` wraps every RPC the local participant performs or handles, which is useful for logging, tracing, or attaching metadata to payloads. Each method receives the call and a `next` continuation; return what `next` returns. Interceptors run in the order they were added, the first being outermost, and errors from the remote side or from your handler flow through them unchanged. On the incoming side the caller's `response_timeout` covers the whole chain, interceptors included.
 
 ```python
 class TimingInterceptor(rtc.RpcInterceptor):

livekit-rtc/livekit/rtc/participant.py

@@ -477,7 +477,8 @@ def add_rpc_interceptor(self, interceptor: RpcInterceptor) -> None:
         Args:
             interceptor (RpcInterceptor): The interceptor to add.
         """
-        if interceptor not in self._rpc_interceptors:
+        # identity, not equality: two distinct interceptors that compare equal must coexist
+        if not any(existing is interceptor for existing in self._rpc_interceptors):
             self._rpc_interceptors.append(interceptor)
 
     def remove_rpc_interceptor(self, interceptor: RpcInterceptor) -> None:
@@ -488,8 +489,9 @@ def remove_rpc_interceptor(self, interceptor: RpcInterceptor) -> None:
         Args:
             interceptor (RpcInterceptor): The interceptor to remove.
         """
-        if interceptor in self._rpc_interceptors:
-            self._rpc_interceptors.remove(interceptor)
+        self._rpc_interceptors = [
+            existing for existing in self._rpc_interceptors if existing is not interceptor
+        ]
 
     def register_rpc_method(
         self,
@@ -598,9 +600,8 @@ async def _handle_rpc_method_invocation(
             request_id, caller_identity, payload, response_timeout, method=method
         )
 
-        handle = _chain_incoming(list(self._rpc_interceptors), self._invoke_rpc_handler)
         try:
-            response_payload = await handle(params)
+            response_payload = await self._run_incoming_chain(params)
         except RpcError as error:
             response_error = error
         except Exception:
@@ -625,26 +626,36 @@ async def _handle_rpc_method_invocation(
             err = res.rpc_method_invocation_response.error
             logger.error(f"error sending rpc method invocation response: {err}")
 
+    async def _run_incoming_chain(self, invocation: RpcInvocationData) -> Optional[str]:
+        """Run the interceptor chain and the handler under the caller's response deadline.
+
+        The deadline covers the whole chain, so time an interceptor spends before or after
+        ``next`` counts against it; when it passes, the chain is cancelled and the caller
+        gets ``RESPONSE_TIMEOUT``. Cancellation from outside (the room disconnecting) maps to
+        ``RECIPIENT_DISCONNECTED``, as before.
+        """
+        handle = _chain_incoming(list(self._rpc_interceptors), self._invoke_rpc_handler)
+        try:
+            return await asyncio.wait_for(handle(invocation), timeout=invocation.response_timeout)
+        except asyncio.TimeoutError:
+            raise RpcError._built_in(RpcError.ErrorCode.RESPONSE_TIMEOUT) from None
+        except asyncio.CancelledError:
+            raise RpcError._built_in(RpcError.ErrorCode.RECIPIENT_DISCONNECTED) from None
+
     async def _invoke_rpc_handler(self, invocation: RpcInvocationData) -> Optional[str]:
         """Run the registered handler for ``invocation`` (the innermost step of the chain).
 
-        Raises ``RpcError`` for an unregistered method, a handler timeout, or a cancelled
-        handler; any other exception from the handler propagates unchanged so interceptors
-        can observe it before the caller's response is built.
+        Raises ``RpcError(UNSUPPORTED_METHOD)`` when nothing is registered; any exception
+        from the handler propagates unchanged so interceptors can observe it before the
+        caller's response is built. The response deadline is enforced by the caller around
+        the whole chain.
         """
         handler = self._rpc_handlers.get(invocation.method)
         if not handler:
             raise RpcError._built_in(RpcError.ErrorCode.UNSUPPORTED_METHOD)
 
         if asyncio.iscoroutinefunction(handler):
-            try:
-                return await asyncio.wait_for(
-                    handler(invocation), timeout=invocation.response_timeout
-                )
-            except asyncio.TimeoutError:
-                raise RpcError._built_in(RpcError.ErrorCode.RESPONSE_TIMEOUT) from None
-            except asyncio.CancelledError:
-                raise RpcError._built_in(RpcError.ErrorCode.RECIPIENT_DISCONNECTED) from None
+            return cast(Optional[str], await handler(invocation))
         return cast(Optional[str], handler(invocation))
 
     async def set_metadata(self, metadata: str) -> None:

livekit-rtc/livekit/rtc/rpc.py

@@ -81,7 +81,9 @@ class RpcInterceptor:
     reaches the caller. On the incoming side, any other exception raised by the handler is
     also visible; the SDK converts it to ``APPLICATION_ERROR`` only after the chain returns,
     and a call for an unregistered method reaches the chain with ``next`` raising
-    ``UNSUPPORTED_METHOD``.
+    ``UNSUPPORTED_METHOD``. The caller's ``response_timeout`` covers the whole incoming
+    chain: when it passes, the chain is cancelled (interceptors see ``CancelledError``) and
+    the caller receives ``RESPONSE_TIMEOUT``.
 
     Example:
         Time every RPC in both directions::

livekit-rtc/tests/test_rpc_interceptors.py

@@ -130,8 +130,9 @@ async def greet(data: RpcInvocationData) -> str:
 
     lp._rpc_handlers["greet"] = greet  # register without the FFI round trip
 
-    handle = _incoming_chain(lp)
-    result = await handle(RpcInvocationData("req-1", "alice", "{}", 5.0, method="greet"))
+    result = await lp._run_incoming_chain(
+        RpcInvocationData("req-1", "alice", "{}", 5.0, method="greet")
+    )
 
     assert result == "hello alice"
     assert log == ["a:in:greet>", "a:in:greet<"]
@@ -152,10 +153,11 @@ async def intercept_incoming(
                 raise
 
     lp.add_rpc_interceptor(Observe())
-    handle = _incoming_chain(lp)
 
     with pytest.raises(rtc.RpcError) as info:
-        await handle(RpcInvocationData("req-1", "alice", "{}", 5.0, method="missing"))
+        await lp._run_incoming_chain(
+            RpcInvocationData("req-1", "alice", "{}", 5.0, method="missing")
+        )
     assert info.value.code == rtc.RpcError.ErrorCode.UNSUPPORTED_METHOD
     assert seen == [info.value]
 
@@ -187,20 +189,77 @@ def sync_ok(data: RpcInvocationData) -> str:
         return "sync"
 
     lp._rpc_handlers.update({"boom": boom, "slow": slow, "sync_ok": sync_ok})
-    handle = _incoming_chain(lp)
 
     # the raw handler exception reaches interceptors; the SDK maps it to APPLICATION_ERROR
     # only when building the response, outside the chain
     with pytest.raises(ValueError):
-        await handle(RpcInvocationData("r1", "alice", "{}", 5.0, method="boom"))
+        await lp._run_incoming_chain(RpcInvocationData("r1", "alice", "{}", 5.0, method="boom"))
     assert isinstance(seen[-1], ValueError)
 
+    # the deadline cancels the chain (interceptors see the cancellation) and the caller
+    # gets RESPONSE_TIMEOUT
+    with pytest.raises(rtc.RpcError) as info:
+        await lp._run_incoming_chain(RpcInvocationData("r2", "alice", "{}", 0.01, method="slow"))
+    assert info.value.code == rtc.RpcError.ErrorCode.RESPONSE_TIMEOUT
+    assert isinstance(seen[-1], asyncio.CancelledError)
+
+    result = await lp._run_incoming_chain(
+        RpcInvocationData("r3", "alice", "{}", 5.0, method="sync_ok")
+    )
+    assert result == "sync"
+
+
+@pytest.mark.parametrize("delay_position", ["before_next", "after_next"])
+async def test_response_deadline_covers_interceptor_time(delay_position: str) -> None:
+    """An interceptor that burns the deadline, on either side of ``next``, times the call out
+    instead of handing the handler a fresh full timeout."""
+    lp = _participant()
+    handler_ran = asyncio.Event()
+
+    class Slow(rtc.RpcInterceptor):
+        async def intercept_incoming(
+            self, invocation: RpcInvocationData, next: IncomingRpcNext
+        ) -> Optional[str]:
+            if delay_position == "before_next":
+                await asyncio.sleep(1.0)
+            result = await next(invocation)
+            if delay_position == "after_next":
+                await asyncio.sleep(1.0)
+            return result
+
+    async def fast(data: RpcInvocationData) -> str:
+        handler_ran.set()
+        return "fast"
+
+    lp.add_rpc_interceptor(Slow())
+    lp._rpc_handlers["fast"] = fast
+
+    start = asyncio.get_running_loop().time()
     with pytest.raises(rtc.RpcError) as info:
-        await handle(RpcInvocationData("r2", "alice", "{}", 0.01, method="slow"))
+        await lp._run_incoming_chain(RpcInvocationData("r1", "alice", "{}", 0.05, method="fast"))
     assert info.value.code == rtc.RpcError.ErrorCode.RESPONSE_TIMEOUT
-    assert seen[-1] is info.value
+    assert asyncio.get_running_loop().time() - start < 0.5
+    assert handler_ran.is_set() == (delay_position == "after_next")
 
-    assert await handle(RpcInvocationData("r3", "alice", "{}", 5.0, method="sync_ok")) == "sync"
+
+async def test_outside_cancellation_maps_to_recipient_disconnected() -> None:
+    lp = _participant()
+    started = asyncio.Event()
+
+    async def hang(data: RpcInvocationData) -> str:
+        started.set()
+        await asyncio.sleep(10)
+        return "never"
+
+    lp._rpc_handlers["hang"] = hang
+    task = asyncio.ensure_future(
+        lp._run_incoming_chain(RpcInvocationData("r1", "alice", "{}", 5.0, method="hang"))
+    )
+    await started.wait()
+    task.cancel()  # what the room does when it disconnects mid-invocation
+    with pytest.raises(rtc.RpcError) as info:
+        await task
+    assert info.value.code == rtc.RpcError.ErrorCode.RECIPIENT_DISCONNECTED
 
 
 async def test_add_and_remove_interceptors() -> None:
@@ -224,13 +283,28 @@ async def fake_ffi(call: RpcCallInfo) -> str:
     assert log == []
 
 
+def test_registration_is_by_identity_not_equality() -> None:
+    class Equalish(rtc.RpcInterceptor):
+        def __eq__(self, other: object) -> bool:
+            return isinstance(other, Equalish)
+
+        def __hash__(self) -> int:
+            return 1
+
+    lp = _participant()
+    first, second = Equalish(), Equalish()
+    assert first == second and first is not second
+
+    lp.add_rpc_interceptor(first)
+    lp.add_rpc_interceptor(second)
+    assert len(lp._rpc_interceptors) == 2, "distinct interceptors that compare equal coexist"
+
+    lp.remove_rpc_interceptor(second)
+    assert len(lp._rpc_interceptors) == 1
+    assert lp._rpc_interceptors[0] is first, "only the exact instance is removed"
+
+
 def test_invocation_data_defaults_keep_positional_construction() -> None:
     # existing code constructs it positionally without `method`
     data = RpcInvocationData("req", "alice", "{}", 2.5)
     assert data.method == ""
-
-
-def _incoming_chain(lp: LocalParticipant) -> IncomingRpcNext:
-    from livekit.rtc.rpc import _chain_incoming
-
-    return _chain_incoming(list(lp._rpc_interceptors), lp._invoke_rpc_handler)