Skip to content

fix(tools): isolate streaming Relay scope - #356

Merged
nvddr merged 4 commits into
mainfrom
agent/fix-async-tool-stream-scope
Aug 12, 2026
Merged

fix(tools): isolate streaming Relay scope#356
nvddr merged 4 commits into
mainfrom
agent/fix-async-tool-stream-scope

Conversation

@nvddr

@nvddr nvddr commented Aug 12, 2026

Copy link
Copy Markdown
Collaborator

Follow-up to #355.

What changed

  • update AGENTS.md so new and migrated native tools are placed in xr-ai-tools, with xr-ai-nat reserved for compatibility surfaces still awaiting migration
  • run each AsyncTool handler in an isolated producer task that owns the Relay scope and handler cleanup
  • preserve streaming backpressure with a one-item queue and propagate producer failures to the consumer
  • add a regression for a swallowed consumer-side error that abandons the stream

Why

AsyncTool.stream previously held a task-local Relay scope open across yield. If a consumer stopped iterating after an error, async-generator finalization could pop that scope from a different context, corrupt the scope stack and prevent streaming handler cleanup from returning the participant to idle.

Impact

Abandoned streams now finalize their handler in the producer Relay context without leaking the tool scope into consumer code. Future streaming tools inherit the safe lifecycle behavior.

Validation

  • tests/.venv/bin/pytest tests/test_native_tools.py tests/test_simple_vlm_example_worker.py (22 passed)
  • uvx --from ruff==0.15.16 ruff check .
  • uvx pyright --pythonpath tests/.venv/bin/python agent-sdk/xr-ai-tools/xr_ai_tools/async_tools.py tests/test_native_tools.py
  • uv build agent-sdk/xr-ai-tools

Signed-off-by: Devdeep Ray <devdeepr@nvidia.com>
@nvddr
nvddr requested a review from blongs-nv August 12, 2026 21:16
@nvddr
nvddr marked this pull request as ready for review August 12, 2026 21:17
@blongs-nv

Copy link
Copy Markdown
Contributor

Reviewed at 848aab8775238a283451e9c9c4d6517b5f45153d by a bot.

This resolves both open blockers from the #355 review: the AGENTS.md sweep is complete (the remaining xr-ai-nat mentions all refer to surfaces genuinely still in NAT), and the producer-task design is the right fix for the scope hazard. The scope isolation, cancellation propagation, aclose-after-cancel safety, and concurrent-stream independence were all verified empirically in this review.

Blockers

  1. A producer that exits without enqueueing a sentinel hangs the consumer forever.

    agent-sdk/xr-ai-tools/xr_ai_tools/async_tools.py:71
    The consumer's only wakeup is await queue.get(), so a BaseException escaping except Exception (:106, e.g. a BaseExceptionGroup from a TaskGroup in a handler) or a producer cancellation from anyone but the consumer's finally (:104 re-raises without a sentinel) deadlocks the stream: the turn never ends and TTS never gets end-of-stream. I reproduced both triggers against this branch. Either race queue.get() against the producer task with asyncio.wait(..., FIRST_COMPLETED), or guarantee a sentinel in an outermost finally with except BaseException.

  2. The new error and cancellation machinery has no tests for its two key paths.

    agent-sdk/xr-ai-tools/xr_ai_tools/async_tools.py:73
    Nothing exercises _StreamFailure propagation (a handler raising mid-stream surfacing via raise item.error) or consumer-task cancellation while awaiting queue.get(), which is the motivating production path (voice barge-in). Both behaviors check out when verified manually, but nothing pins them; add a handler-raises case asserting the exception surfaces from the async for, and a task.cancel() case asserting CancelledError propagates and the handler's finally ran. The dangling streaming fixture at tests/test_simple_vlm_example_worker.py:574 (pre-existing on main, constructed and patched but never iterated) is a cheap place to restore the first.

  3. The changelog is now actively wrong about the AsyncTool lifecycle and this PR adds no entry.

    docs/changelog.md:30
    AGENTS.md:338 requires significant decisions to be recorded, and the existing 2026-08-12 entries describe the consumer-context scope design this PR supersedes. Add an entry stating the producer-task lifecycle as the settled design.

Suggestions

  1. Abandoned streams still depend on GC-timed finalization, one round per generator layer: the response() wrapper at app.py:23 and the abandoned iterator in handler.py:246 are never aclose()d, so on barge-in the superseded turn's deferred set_status("idle") can land after the next turn set "processing", and a held generator reference defers cleanup indefinitely. Deterministic aclose() at those consumer sites is the complement to this fix; worth naming in the follow-ups list.

  2. A routine barge-in now closes the tool scope as an error span, the producer runs one chunk ahead of the consumer (maxsize=1), and the tool span ends before the last chunk is consumed. All fine if deliberate, but they are silent telemetry behavior changes worth a decision or a comment.

  3. await producer in the consumer's finally is unbounded, so a handler whose cleanup blocks (e.g. set_status against a dead endpoint) stalls the consumer's cancellation path; consider a bounded wait.

  4. The stream docstring's "under a tool scope" now overstates the contract (chunks are yielded outside it; the scope is task-local to the producer), and per its own docstring fork_asyncio_context does not transfer scope-local middleware, so a caller registering a sanitizer before calling stream() silently loses it for tool events. One docstring covering both is enough.

  5. In the new test, only the consumer-scope-unchanged assertion actually discriminates the fix (the cleanup and parent-scope assertions pass on the old code too); a one-line comment saying so, and dropping or justifying the incidental context=fork_asyncio_context() on the consumer task, would keep the regression guard legible.

This is the right fix, verified to do what it claims: the producer task is the only design that makes the scope pop context-correct, and the AGENTS.md updates close out the documentation drift. The blockers above should be resolved before merge.

Signed-off-by: Devdeep Ray <devdeepr@nvidia.com>
@nvddr

nvddr commented Aug 12, 2026

Copy link
Copy Markdown
Collaborator Author

Addressed in 7113054.

Blockers

  1. Replaced sentinel-only wakeups with a race between queue.get() and producer completion. The consumer drains a final buffered chunk before awaiting the producer, then propagates every producer exit state directly, including ordinary exceptions, BaseExceptionGroup, and cancellation.
  2. Added regressions for ordinary handler failure after buffered output, BaseExceptionGroup without a hang, consumer-task cancellation with handler cleanup, and producer completion with a final buffered chunk.
  3. Added a changelog entry recording the isolated producer-task lifecycle as the settled design.

Suggestions

  • _VoiceHandlerProcessor now explicitly closes streaming responses, and the simple-VLM response wrapper explicitly closes its nested AsyncTool stream. Barge-in cleanup no longer depends on async-generator GC.
  • The changelog records the deliberate telemetry/backpressure behavior: cancellation closes the tool scope as an error, the producer may be one chunk ahead, and its span may end before the last buffered chunk is consumed.
  • AsyncTool.stream now documents that chunks are yielded in the consumer context and parent scope-local registrations are not transferred.
  • The abandoned-stream regression now explains why the forked consumer and consumer-scope assertion distinguish this fix.
  • I left producer cleanup unbounded. The changelog names this as a follow-up because imposing a timeout requires a policy for detached cleanup; abandoning it can leave Relay state or participant status unfinished.

Validation:

  • 114 focused tests passed (test_native_tools.py, test_simple_vlm_example_worker.py, test_voice_pipeline.py)
  • repository-wide Ruff 0.15.16 passed
  • targeted Pyright passed for all changed implementations and focused native/simple-VLM tests
  • xr-ai-tools and xr-ai-voice source/wheel builds passed

Comment thread agent-sdk/xr-ai-tools/xr_ai_tools/async_tools.py Fixed
Comment thread tests/test_native_tools.py Fixed
Comment thread agent-sdk/xr-ai-tools/xr_ai_tools/async_tools.py Fixed
Signed-off-by: Devdeep Ray <devdeepr@nvidia.com>
Comment thread agent-sdk/xr-ai-tools/xr_ai_tools/async_tools.py Fixed
Comment thread agent-sdk/xr-ai-tools/xr_ai_tools/async_tools.py Fixed
Signed-off-by: Devdeep Ray <devdeepr@nvidia.com>
@nvddr
nvddr merged commit 452ecae into main Aug 12, 2026
12 checks passed
@nvddr
nvddr deleted the agent/fix-async-tool-stream-scope branch August 12, 2026 22:12
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants