add support for RPC interceptors

README.md

@@ -223,6 +223,29 @@ except Exception as e:
 
 You may find it useful to adjust the `response_timeout` parameter, which indicates the amount of time you will wait for a response. We recommend keeping this value as low as possible while still satisfying the constraints of your application.
 
+#### 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. On the incoming side the caller's `response_timeout` covers the whole chain, interceptors included.
+
+```python
+class TimingInterceptor(rtc.RpcInterceptor):
+    async def intercept_outgoing(self, call, next):
+        start = time.perf_counter()
+        try:
+            return await next(call)
+        finally:
+            print(f"call {call.method} -> {call.destination_identity}: {time.perf_counter() - start:.3f}s")
+
+    async def intercept_incoming(self, invocation, next):
+        start = time.perf_counter()
+        try:
+            return await next(invocation)
+        finally:
+            print(f"handled {invocation.method} from {invocation.caller_identity}: {time.perf_counter() - start:.3f}s")
+
+room.local_participant.add_rpc_interceptor(TimingInterceptor())
+```
+
 ## Using local media devices
 
 The `MediaDevices` class provides a high-level interface for working with local audio input (microphone) and output (speakers) devices. It's built on top of the `sounddevice` library and integrates seamlessly with LiveKit's audio processing features.  In order to use `MediaDevices`, you must have the `sounddevice` library installed in your local Python environment, if it's not available, `MediaDevices` will not work.

livekit-rtc/livekit/rtc/__init__.py

@@ -115,7 +115,7 @@
 except Exception:  # pragma: no cover - optional dependency (sounddevice)
     _HAS_MEDIA_DEVICES = False
 from .utils import combine_audio_frames
-from .rpc import RpcError, RpcInvocationData
+from .rpc import RpcCallInfo, RpcError, RpcInterceptor, RpcInvocationData
 from .synchronizer import AVSynchronizer
 from .data_stream import (
     TextStreamInfo,
@@ -211,6 +211,8 @@
     "AudioResamplerQuality",
     "RpcError",
     "RpcInvocationData",
+    "RpcCallInfo",
+    "RpcInterceptor",
     "EventEmitter",
     "combine_audio_frames",
     "AVSynchronizer",

livekit-rtc/livekit/rtc/participant.py

@@ -48,11 +48,17 @@
     TrackPublication,
 )
 from .transcription import Transcription
-from .rpc import RpcError
+from .rpc import (
+    RpcCallInfo,
+    RpcError,
+    RpcInterceptor,
+    RpcInvocationData,
+    _chain_incoming,
+    _chain_outgoing,
+)
 from ._proto.rpc_pb2 import RpcMethodInvocationResponseRequest
 from .log import logger
 
-from .rpc import RpcInvocationData
 from .data_stream import (
     TextStreamWriter,
     TextStreamInfo,
@@ -230,6 +236,8 @@ def disconnect_reason(
 
 
 RpcHandler = Callable[["RpcInvocationData"], Union[Awaitable[Optional[str]], Optional[str]]]
+
+
 F = TypeVar(
     "F", bound=Callable[[RpcInvocationData], Union[Awaitable[Optional[str]], Optional[str]]]
 )
@@ -247,6 +255,7 @@ def __init__(
         self._room_queue = room_queue
         self._track_publications: dict[str, LocalTrackPublication] = {}
         self._rpc_handlers: Dict[str, RpcHandler] = {}
+        self._rpc_interceptors: List[RpcInterceptor] = []
         # Handles of data stream writers that have been opened but not yet
         # closed, so the room can drop them at disconnect. The FFI close
         # request consumes the handle (take_handle), so an entry is removed as
@@ -427,15 +436,27 @@ async def perform_rpc(
         Raises:
             RpcError: On failure. Details in `message`.
         """
+        call = RpcCallInfo(
+            destination_identity=destination_identity,
+            method=method,
+            payload=payload,
+            response_timeout=response_timeout,
+            max_round_trip_latency=max_round_trip_latency,
+        )
+        # snapshot the interceptor list so add/remove during a call is well defined
+        perform = _chain_outgoing(list(self._rpc_interceptors), self._perform_rpc_ffi)
+        return await perform(call)
+
+    async def _perform_rpc_ffi(self, call: RpcCallInfo) -> str:
         req = proto_ffi.FfiRequest()
         req.perform_rpc.local_participant_handle = self._ffi_handle.handle
-        req.perform_rpc.destination_identity = destination_identity
-        req.perform_rpc.method = method
-        req.perform_rpc.payload = payload
-        if response_timeout is not None:
-            req.perform_rpc.response_timeout_ms = int(response_timeout * 1000)
-        if max_round_trip_latency is not None:
-            req.perform_rpc.max_round_trip_latency_ms = int(max_round_trip_latency * 1000)
+        req.perform_rpc.destination_identity = call.destination_identity
+        req.perform_rpc.method = call.method
+        req.perform_rpc.payload = call.payload
+        if call.response_timeout is not None:
+            req.perform_rpc.response_timeout_ms = int(call.response_timeout * 1000)
+        if call.max_round_trip_latency is not None:
+            req.perform_rpc.max_round_trip_latency_ms = int(call.max_round_trip_latency * 1000)
 
         queue = FfiClient.instance.queue.subscribe()
         try:
@@ -449,6 +470,31 @@ async def perform_rpc(
 
         return cast(str, cb.perform_rpc.payload)
 
+    def add_rpc_interceptor(self, interceptor: RpcInterceptor) -> None:
+        """
+        Add an :class:`RpcInterceptor` that wraps every RPC this participant performs or
+        handles. Interceptors run in the order they were added, the first being outermost.
+        Adding the same instance twice is a no-op.
+
+        Args:
+            interceptor (RpcInterceptor): The interceptor to add.
+        """
+        # 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:
+        """
+        Remove a previously added :class:`RpcInterceptor`. Calls already in flight keep the
+        chain they started with.
+
+        Args:
+            interceptor (RpcInterceptor): The interceptor to remove.
+        """
+        self._rpc_interceptors = [
+            existing for existing in self._rpc_interceptors if existing is not interceptor
+        ]
+
     def register_rpc_method(
         self,
         method_name: str,
@@ -552,33 +598,20 @@ async def _handle_rpc_method_invocation(
         response_error: Optional[RpcError] = None
         response_payload: Optional[str] = None
 
-        params = RpcInvocationData(request_id, caller_identity, payload, response_timeout)
-
-        handler = self._rpc_handlers.get(method)
+        params = RpcInvocationData(
+            request_id, caller_identity, payload, response_timeout, method=method
+        )
 
-        if not handler:
-            response_error = RpcError._built_in(RpcError.ErrorCode.UNSUPPORTED_METHOD)
-        else:
-            try:
-                if asyncio.iscoroutinefunction(handler):
-                    try:
-                        response_payload = await asyncio.wait_for(
-                            handler(params), timeout=response_timeout
-                        )
-                    except asyncio.TimeoutError:
-                        raise RpcError._built_in(RpcError.ErrorCode.RESPONSE_TIMEOUT)
-                    except asyncio.CancelledError:
-                        raise RpcError._built_in(RpcError.ErrorCode.RECIPIENT_DISCONNECTED)
-                else:
-                    response_payload = cast(Optional[str], handler(params))
-            except RpcError as error:
-                response_error = error
-            except Exception:
-                logger.exception(
-                    f"Uncaught error returned by RPC handler for {method}. "
-                    "Returning APPLICATION_ERROR instead. "
-                )
-                response_error = RpcError._built_in(RpcError.ErrorCode.APPLICATION_ERROR)
+        try:
+            response_payload = await self._run_incoming_chain(params)
+        except RpcError as error:
+            response_error = error
+        except Exception:
+            logger.exception(
+                f"Uncaught error returned by RPC handler for {method}. "
+                "Returning APPLICATION_ERROR instead. "
+            )
+            response_error = RpcError._built_in(RpcError.ErrorCode.APPLICATION_ERROR)
 
         req = proto_ffi.FfiRequest(
             rpc_method_invocation_response=RpcMethodInvocationResponseRequest(
@@ -595,6 +628,68 @@ 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``, whatever the chain raises while unwinding. Cancellation
+        from outside (the room disconnecting) maps to ``RECIPIENT_DISCONNECTED``, as before.
+
+        A ``TimeoutError`` raised by the handler or an interceptor itself before the
+        deadline (an HTTP client timing out, say) is not the response deadline: it
+        propagates as an application error rather than being reported as
+        ``RESPONSE_TIMEOUT``.
+        """
+        handle = _chain_incoming(list(self._rpc_interceptors), self._invoke_rpc_handler)
+        loop = asyncio.get_running_loop()
+        # ensure_future: `next` continuations are typed as Awaitable, not Coroutine
+        chain_task: asyncio.Future[Optional[str]] = asyncio.ensure_future(handle(invocation))
+
+        # the deadline is recorded independently of whatever the chain raises while it
+        # unwinds, so a TimeoutError (or anything else) from cancellation cleanup cannot be
+        # mistaken for an application failure, nor a genuine in-chain timeout for the deadline
+        deadline_fired = False
+
+        def _on_deadline() -> None:
+            nonlocal deadline_fired
+            deadline_fired = True
+            chain_task.cancel()
+
+        deadline = loop.call_later(invocation.response_timeout, _on_deadline)
+        try:
+            return await 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])
+            raise RpcError._built_in(RpcError.ErrorCode.RECIPIENT_DISCONNECTED) from None
+        except Exception:
+            if deadline_fired:
+                raise RpcError._built_in(RpcError.ErrorCode.RESPONSE_TIMEOUT) from None
+            raise
+        finally:
+            deadline.cancel()
+
+    async def _invoke_rpc_handler(self, invocation: RpcInvocationData) -> Optional[str]:
+        """Run the registered handler for ``invocation`` (the innermost step of the chain).
+
+        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):
+            return cast(Optional[str], await handler(invocation))
+        return cast(Optional[str], handler(invocation))
+
     async def set_metadata(self, metadata: str) -> None:
         """
         Set the metadata for the local participant.