add support for RPC interceptors

allowing telemetry hooks so we can provide automatic traces when used in
agents.

Author: davidzhao

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.
+
+```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,
@@ -247,6 +253,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 +434,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 +468,29 @@ 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.
+        """
+        if interceptor not 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.
+        """
+        if interceptor in self._rpc_interceptors:
+            self._rpc_interceptors.remove(interceptor)
+
     def register_rpc_method(
         self,
         method_name: str,
@@ -552,33 +594,21 @@ 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)
+        handle = _chain_incoming(list(self._rpc_interceptors), self._invoke_rpc_handler)
+        try:
+            response_payload = await handle(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 +625,28 @@ 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 _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.
+        """
+        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], handler(invocation))
+
     async def set_metadata(self, metadata: str) -> None:
         """
         Set the metadata for the local participant.

livekit-rtc/livekit/rtc/rpc.py

@@ -12,7 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Optional, Dict, Union, ClassVar
+from __future__ import annotations
+
+from typing import Awaitable, Callable, ClassVar, Dict, Optional, Sequence, Union
 from enum import IntEnum
 from ._proto import rpc_pb2 as proto_rpc
 from dataclasses import dataclass
@@ -27,12 +29,123 @@ class RpcInvocationData:
         caller_identity (str): The unique participant identity of the caller.
         payload (str): The payload of the request. User-definable format, typically JSON.
         response_timeout (float): The maximum time the caller will wait for a response.
+        method (str): The name of the invoked RPC method.
     """
 
     request_id: str
     caller_identity: str
     payload: str
     response_timeout: float
+    method: str = ""
+
+
+@dataclass
+class RpcCallInfo:
+    """An outgoing RPC call, as passed to :meth:`RpcInterceptor.intercept_outgoing`.
+
+    Mirrors the arguments of :meth:`LocalParticipant.perform_rpc`. An interceptor may pass a
+    modified copy to ``next`` (for example to add a header to a JSON payload).
+
+    Attributes:
+        destination_identity (str): The identity of the participant being called.
+        method (str): The method name.
+        payload (str): The request payload.
+        response_timeout (Optional[float]): Seconds to wait for a response, or ``None`` for the default.
+        max_round_trip_latency (Optional[float]): See :meth:`LocalParticipant.perform_rpc`.
+    """
+
+    destination_identity: str
+    method: str
+    payload: str
+    response_timeout: Optional[float] = None
+    max_round_trip_latency: Optional[float] = None
+
+
+OutgoingRpcNext = Callable[["RpcCallInfo"], Awaitable[str]]
+"""Continuation handed to :meth:`RpcInterceptor.intercept_outgoing`: performs the call."""
+IncomingRpcNext = Callable[["RpcInvocationData"], Awaitable[Optional[str]]]
+"""Continuation handed to :meth:`RpcInterceptor.intercept_incoming`: runs the handler."""
+
+
+class RpcInterceptor:
+    """Observe or wrap RPC calls made and handled by a :class:`LocalParticipant`.
+
+    Register with :meth:`LocalParticipant.add_rpc_interceptor`. Each method receives the call
+    and a ``next`` continuation and must return (or raise) what ``next`` returns (or raises),
+    unless it deliberately short-circuits the call. Interceptors run in registration order:
+    the first one added is the outermost. Both methods default to pass-through, so override
+    only the direction you care about.
+
+    Errors flow through the chain unchanged: a :class:`RpcError` raised by the remote side
+    (outgoing) or by the handler (incoming) is visible to every interceptor before it
+    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``.
+
+    Example:
+        Time every RPC in both directions::
+
+            class TimingInterceptor(rtc.RpcInterceptor):
+                async def intercept_outgoing(self, call, next):
+                    start = time.perf_counter()
+                    try:
+                        return await next(call)
+                    finally:
+                        log("rpc call", call.method, time.perf_counter() - start)
+
+                async def intercept_incoming(self, invocation, next):
+                    start = time.perf_counter()
+                    try:
+                        return await next(invocation)
+                    finally:
+                        log("rpc handled", invocation.method, time.perf_counter() - start)
+
+            room.local_participant.add_rpc_interceptor(TimingInterceptor())
+    """
+
+    async def intercept_outgoing(self, call: RpcCallInfo, next: OutgoingRpcNext) -> str:
+        """Wrap an outgoing :meth:`LocalParticipant.perform_rpc`. Return the response payload."""
+        return await next(call)
+
+    async def intercept_incoming(
+        self, invocation: RpcInvocationData, next: IncomingRpcNext
+    ) -> Optional[str]:
+        """Wrap the handling of an incoming invocation. Return the response payload."""
+        return await next(invocation)
+
+
+def _chain_outgoing(
+    interceptors: Sequence[RpcInterceptor], terminal: OutgoingRpcNext
+) -> OutgoingRpcNext:
+    """Compose ``interceptors`` around ``terminal``; the first interceptor is outermost."""
+    call_next = terminal
+    for interceptor in reversed(interceptors):
+        call_next = _bind_outgoing(interceptor, call_next)
+    return call_next
+
+
+def _chain_incoming(
+    interceptors: Sequence[RpcInterceptor], terminal: IncomingRpcNext
+) -> IncomingRpcNext:
+    call_next = terminal
+    for interceptor in reversed(interceptors):
+        call_next = _bind_incoming(interceptor, call_next)
+    return call_next
+
+
+def _bind_outgoing(interceptor: RpcInterceptor, call_next: OutgoingRpcNext) -> OutgoingRpcNext:
+    async def bound(call: RpcCallInfo) -> str:
+        return await interceptor.intercept_outgoing(call, call_next)
+
+    return bound
+
+
+def _bind_incoming(interceptor: RpcInterceptor, call_next: IncomingRpcNext) -> IncomingRpcNext:
+    async def bound(invocation: RpcInvocationData) -> Optional[str]:
+        return await interceptor.intercept_incoming(invocation, call_next)
+
+    return bound
 
 
 class RpcError(Exception):