add support for RPC interceptors - #806
Conversation
allowing telemetry hooks so we can provide automatic traces when used in agents.
![devin-ai-integration[bot]](https://avatars.githubusercontent.com/in/811515?s=60&v=4){kind=small}
This comment was marked as resolved.
This comment was marked as resolved.
Sorry, something went wrong.
…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>
This comment was marked as resolved.
This comment was marked as resolved.
Sorry, something went wrong.
asyncio.wait_for raises the same TimeoutError a handler or interceptor might raise on its own (an HTTP client timing out, say), so _run_incoming_chain reported both as RESPONSE_TIMEOUT. The pre-existing handler wrapper had the same conflation, but now that the deadline spans the chain it is cheap to get right: a TimeoutError raised inside the chain is tagged (_ChainTimeoutError) and re-raised as the original exception, which _handle_rpc_method_invocation maps to APPLICATION_ERROR like any other handler failure; only wait_for's own expiry becomes RESPONSE_TIMEOUT. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
This comment was marked as resolved.
This comment was marked as resolved.
Sorry, something went wrong.
…ng error wait_for reported whatever the chain raised while being cancelled; an interceptor that catches the cancellation and raises during cleanup could turn a deadline expiry into APPLICATION_ERROR (or, via the TimeoutError tag, back into a mislabelled application timeout). The chain now runs as its own future under a call_later deadline that sets a flag before cancelling it. Once the flag is set, anything the chain raises while unwinding is reported as RESPONSE_TIMEOUT; a CancelledError without the flag is external cancellation and stays RECIPIENT_DISCONNECTED; any exception before the deadline, TimeoutError included, propagates as an application error. The _ChainTimeoutError tag is no longer needed. Tests: cleanup raising TimeoutError and RuntimeError after the deadline both map to RESPONSE_TIMEOUT; the existing deadline, in-chain timeout and outside cancellation cases still hold. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
xianshijing-lk
left a comment
xianshijing-lk
left a comment
There was a problem hiding this comment.
some comments from a quick look, look good to me assuming you are addressing them.
| # 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]) |
There was a problem hiding this comment.
I know the code is also doing the unbound wait, but wonder if it makes sense to bound the wait here, like
if not chain_task.done():
await asyncio.wait([chain_task], timeout=_CANCEL_TIMEOUT)
There was a problem hiding this comment.
Good call, done in 91a23d8, though it needed more than the timeout: awaiting the chain directly forwards the outer cancel into the chain and keeps this task parked until the chain finishes, so the wait here would only have started once a stubborn handler had already stopped (at the caller's deadline). The chain is now awaited through asyncio.shield, so an outside cancel is raised here at once; we then cancel the chain explicitly, wait up to 2s for it to unwind, and log a warning naming the method if it does not stop. Test added with a handler that swallows cancellation.
|
|
||
| if asyncio.iscoroutinefunction(handler): | ||
| return cast(Optional[str], await handler(invocation)) | ||
| return cast(Optional[str], handler(invocation)) |
There was a problem hiding this comment.
I wonder if it is slightly cleaner to do
result = handler(invocation)
if inspect.isawaitable(result):
result = await result
return cast(Optional[str], result)
| 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)) |
There was a problem hiding this comment.
any chance that some incoming RPCs are no-interceptor + sync handling ? that just run inline with zero scheduling ? If that is a valid use case, probably we don't want to allocate a task and timer for them ?
There was a problem hiding this comment.
we would want to observe everything. so I wouldn't create an exception around observability.
rpc handlers are always async AFAIK, and I think that's the right thing to ensure no sync operations could take place here
…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>
This comment was marked as resolved.
This comment was marked as resolved.
Sorry, something went wrong.
Once the caller has been answered nobody awaits the chain, and the shield stops watching it the moment its own future is cancelled, so an exception from the handler's cancellation cleanup surfaced as a never-retrieved task exception at garbage collection. It is now retrieved and logged, whether the handler stops within the unwind bound or after it. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
…leanup handler Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
This comment was marked as resolved.
This comment was marked as resolved.
Sorry, something went wrong.
2 participants
allowing telemetry hooks so we can provide automatic traces when used in agents.