fix(tools): isolate streaming Relay scope

AGENTS.md

@@ -125,9 +125,10 @@ mechanically:
 
 **Worker code rules** (apply to every sample worker):
 
-- Import IPC types from `xr_ai_hub`; native agent functions come from
-  `xr_ai_nat`, model clients from `xr_ai_models`, and the native voice runtime
-  from `xr_ai_voice`.
+- Import IPC types from `xr_ai_hub`; new and migrated native agent tools come
+  from `xr_ai_tools`, model clients from `xr_ai_models`, and the native voice
+  runtime from `xr_ai_voice`. Use `xr_ai_nat` only for compatibility surfaces
+  that have not migrated yet.
 - Raw IPC workers keep `_HUB_PUB` / `_HUB_PUSH` as module-level constants,
   wire `SIGINT` and `SIGTERM` to a synchronous `shutdown()`, cancel asyncio
   tasks first, then call `ep.stop()` + `ep.close()`. Voice workers delegate
@@ -162,7 +163,8 @@ Reference implementation: `agent-samples/simple-vlm-example/`.
 Samples must **reuse** the shared building blocks rather than re-implement
 them. They split across SDK packages by what they depend on:
 
-Typed agent functions live in `xr-ai-nat`. `SpatialMathFunctionsConfig`
+New and migrated typed agent tools live in `xr-ai-tools`; `xr-ai-nat` retains
+compatibility function groups while they migrate. `SpatialMathFunctionsConfig`
 registers deterministic coordinate operations that receive an explicit spatial
 frame; tracking and process boundaries remain outside the math functions.
 `TextMemoryFunctionsConfig` provides persistent timestamped text without a
@@ -205,7 +207,8 @@ behavior comes from config alone.
 ### Scope decision and named follow-ups
 
 The function/pipeline boundary is explicit: `xr-ai-pipecat` stays "voice
-pipeline plumbing", while reusable typed agent functions live in `xr-ai-nat`.
+pipeline plumbing", while new and migrated reusable typed agent tools live in
+`xr-ai-tools`.
 Application-specific capabilities stay with their application;
 `xr-ai-pipecat` must not become a catch-all.
 Planned structural follow-ups (own PRs):

agent-samples/simple-vlm-example/worker/simple_vlm_example_worker/app.py

@@ -21,13 +21,18 @@
 def _make_vision_handler(vision: StreamingVisionTool) -> VoiceHandler:
     async def handle(turn):
         async def response():
-            async for chunk in vision.stream(
-                VisionRequest(
-                    participant_id=turn.participant_id,
-                    query=turn.text,
-                )
-            ):
-                yield chunk.text
+            request = VisionRequest(
+                participant_id=turn.participant_id,
+                query=turn.text,
+            )
+            stream = vision.stream(request)
+            try:
+                async for chunk in stream:
+                    yield chunk.text
+            finally:
+                close = getattr(stream, "aclose", None)
+                if close is not None:
+                    await close()
 
         return response()
 

agent-sdk/xr-ai-tools/xr_ai_tools/async_tools.py

@@ -5,7 +5,9 @@
 
 from __future__ import annotations
 
+import asyncio
 from collections.abc import AsyncIterator, Callable, Mapping
+from contextlib import suppress
 from typing import Generic, TypeVar
 
 import nemo_relay
@@ -40,16 +42,69 @@ async def stream(
         self,
         request: RequestT | Mapping[str, object],
     ) -> AsyncIterator[ChunkT]:
-        """Validate and yield one typed result stream under a tool scope."""
+        """Validate and yield chunks produced in an isolated Relay tool scope.
+
+        The producer task owns the scope and handler cleanup; chunks are yielded
+        in the consumer's context. Parent scope-local registrations are not
+        transferred into the isolated context.
+        """
 
         value = self.request_model.model_validate(request)
+        queue: asyncio.Queue[ChunkT] = asyncio.Queue(maxsize=1)
+        producer = asyncio.create_task(
+            self._produce(value, queue),
+            name=f"xr-ai-tool:{self.name}",
+            context=nemo_relay.fork_asyncio_context(),
+        )
+        try:
+            while True:
+                next_chunk = asyncio.create_task(queue.get())
+                try:
+                    done, _ = await asyncio.wait(
+                        (next_chunk, producer),
+                        return_when=asyncio.FIRST_COMPLETED,
+                    )
+                finally:
+                    if not next_chunk.done():
+                        next_chunk.cancel()
+                        with suppress(asyncio.CancelledError):
+                            _ = await next_chunk
+
+                if next_chunk in done:
+                    yield next_chunk.result()
+                    continue
+
+                if not queue.empty():
+                    # Producer completion can win with its final chunk buffered.
+                    yield queue.get_nowait()
+                    continue
+
+                producer.result()
+                return
+        finally:
+            if not producer.done():
+                producer.cancel()
+            with suppress(asyncio.CancelledError):
+                _ = await producer
+
+    async def _produce(
+        self,
+        value: RequestT,
+        queue: asyncio.Queue[ChunkT],
+    ) -> None:
         with nemo_relay.scope.scope(
             self.name,
             nemo_relay.ScopeType.Tool,
             input=value.model_dump(mode="json"),
         ):
-            async for chunk in self.handler(value):
-                yield self.chunk_model.model_validate(chunk)
+            stream = self.handler(value)
+            try:
+                async for chunk in stream:
+                    await queue.put(self.chunk_model.model_validate(chunk))
+            finally:
+                close = getattr(stream, "aclose", None)
+                if close is not None:
+                    await close()
 
 
 __all__ = ["AsyncTool"]

agent-sdk/xr-ai-voice/xr_ai_voice/_processors/handler.py

@@ -243,11 +243,16 @@ async def _run_query(self, frame: GatedQueryFrame, token: object) -> None:
                     accumulated.append(result)
                     await self._push_text(result, pid=pid)
                 return
-            async for chunk in result:
-                if not chunk or self._turn_tokens.get(pid) is not token:
-                    continue
-                accumulated.append(chunk)
-                await self._push_text(chunk, pid=pid)
+            try:
+                async for chunk in result:
+                    if not chunk or self._turn_tokens.get(pid) is not token:
+                        continue
+                    accumulated.append(chunk)
+                    await self._push_text(chunk, pid=pid)
+            finally:
+                close = getattr(result, "aclose", None)
+                if close is not None:
+                    await close()
         except asyncio.CancelledError:
             cancelled = True
             raise

docs/changelog.md

@@ -9,6 +9,26 @@ Significant decisions, in reverse-chronological order. Update this whenever a
 non-trivial architectural or design decision is made so the rationale is
 preserved and not re-litigated.
 
+### 2026-08-12 — Streaming tools isolate Relay scopes in producer tasks
+
+`AsyncTool` runs each handler in a forked producer task that exclusively owns
+its Relay tool scope and closes the handler in that same task and context.
+Chunks cross a one-item queue and are yielded in the consumer's context. The
+consumer races the queue against producer completion so normal completion,
+exceptions (including `BaseExceptionGroup`), and cancellation cannot leave it
+waiting forever. Abandoning or cancelling the consumer cancels and awaits the
+producer. Voice handlers explicitly close streaming responses, and the simple
+VLM adapter also closes its nested tool stream, so cleanup does not depend on
+async-generator garbage collection.
+
+This boundary deliberately means callers do not run inside the tool scope and
+parent scope-local registrations are not copied into the producer context.
+Cancellation records the tool scope as an error; the producer may stay one
+chunk ahead of the consumer, and its span may end before the final buffered
+chunk is consumed. Producer cleanup remains unbounded: a timeout and detached
+cleanup policy requires a separate decision because abandoning cleanup could
+leave Relay state or participant status unfinished.
+
 ### 2026-08-12 — Tool-call handling is not an agent runtime
 
 `agents.py`, `agent_runner.py`, `Agent`, and `AgentRunner` are removed.

tests/test_native_tools.py

@@ -5,7 +5,9 @@
 
 from __future__ import annotations
 
+import asyncio
 import json
+from builtins import BaseExceptionGroup
 from collections.abc import AsyncIterator
 
 import nemo_relay
@@ -52,6 +54,148 @@ async def test_async_tool_validates_and_yields_typed_chunks() -> None:
     assert chunks == [AddResult(total=2), AddResult(total=5)]
 
 
+async def test_async_tool_yields_chunk_buffered_before_producer_completion() -> None:
+    async def single_chunk(request: AddRequest) -> AsyncIterator[AddResult]:
+        yield AddResult(total=request.left)
+
+    tool = AsyncTool(
+        "stream_add",
+        "Stream a running total.",
+        AddRequest,
+        AddResult,
+        single_chunk,
+    )
+
+    chunks = [chunk async for chunk in tool.stream({"left": 2, "right": 3})]
+
+    assert chunks == [AddResult(total=2)]
+
+
+async def test_async_tool_propagates_handler_failure_after_buffered_chunks() -> None:
+    emitted = []
+
+    async def failing_stream(request: AddRequest) -> AsyncIterator[AddResult]:
+        yield AddResult(total=request.left)
+        raise RuntimeError("stream failed")
+
+    tool = AsyncTool(
+        "stream_add",
+        "Stream a running total.",
+        AddRequest,
+        AddResult,
+        failing_stream,
+    )
+
+    with pytest.raises(RuntimeError, match="stream failed"):
+        async for chunk in tool.stream({"left": 2, "right": 3}):
+            emitted.append(chunk)
+
+    assert emitted == [AddResult(total=2)]
+
+
+async def test_async_tool_propagates_base_exception_group_without_hanging() -> None:
+    class FatalStreamError(BaseException):
+        pass
+
+    async def fatal_stream(_request: AddRequest) -> AsyncIterator[AddResult]:
+        if False:
+            yield AddResult(total=0)
+        raise BaseExceptionGroup("stream failed", [FatalStreamError()])
+
+    tool = AsyncTool(
+        "stream_add",
+        "Stream a running total.",
+        AddRequest,
+        AddResult,
+        fatal_stream,
+    )
+
+    async def consume() -> None:
+        async for _chunk in tool.stream({"left": 2, "right": 3}):
+            pass
+
+    with pytest.raises(BaseExceptionGroup, match="stream failed"):
+        await asyncio.wait_for(consume(), timeout=1.0)
+
+
+async def test_async_tool_consumer_cancellation_closes_handler() -> None:
+    started = asyncio.Event()
+    closed = asyncio.Event()
+    blocked = asyncio.Event()
+
+    async def blocking_stream(_request: AddRequest) -> AsyncIterator[AddResult]:
+        started.set()
+        try:
+            await blocked.wait()
+            if False:
+                yield AddResult(total=0)
+        finally:
+            closed.set()
+
+    tool = AsyncTool(
+        "stream_add",
+        "Stream a running total.",
+        AddRequest,
+        AddResult,
+        blocking_stream,
+    )
+
+    async def consume() -> None:
+        async for _chunk in tool.stream({"left": 2, "right": 3}):
+            pass
+
+    consumer = asyncio.create_task(consume())
+    await asyncio.wait_for(started.wait(), timeout=1.0)
+    consumer.cancel()
+
+    with pytest.raises(asyncio.CancelledError):
+        _ = await consumer
+
+    assert closed.is_set()
+
+
+async def test_async_tool_closes_an_abandoned_stream_in_its_relay_context() -> None:
+    closed = asyncio.Event()
+    blocked = asyncio.Event()
+
+    async def blocking_stream(request: AddRequest) -> AsyncIterator[AddResult]:
+        try:
+            yield AddResult(total=request.left)
+            await blocked.wait()
+        finally:
+            closed.set()
+
+    tool = AsyncTool(
+        "stream_add",
+        "Stream a running total.",
+        AddRequest,
+        AddResult,
+        blocking_stream,
+    )
+    consumer_scope_unchanged: list[bool] = []
+
+    async def abandon_stream() -> None:
+        consumer_scope = nemo_relay.scope.get_handle()
+        try:
+            async for chunk in tool.stream({"left": 2, "right": 3}):
+                assert chunk == AddResult(total=2)
+                raise RuntimeError("consumer failed")
+        except RuntimeError:
+            consumer_scope_unchanged.append(
+                nemo_relay.scope.get_handle().uuid == consumer_scope.uuid
+            )
+
+    # A forked consumer reproduces finalization outside the caller's Relay context.
+    await asyncio.create_task(
+        abandon_stream(),
+        context=nemo_relay.fork_asyncio_context(),
+    )
+    await asyncio.wait_for(closed.wait(), timeout=1.0)
+
+    # This is the assertion that fails when the tool scope leaks across a yield.
+    assert consumer_scope_unchanged == [True]
+
+
 def test_tool_definitions_adapt_native_tools_for_model_services() -> None:
     tool = Tool("add", "Add two integers.", AddRequest, AddResult, add)
 

tests/test_simple_vlm_example_worker.py

@@ -5,6 +5,7 @@
 
 from __future__ import annotations
 
+import asyncio
 import sys
 import time
 from pathlib import Path
@@ -285,6 +286,36 @@ def test_config_rejects_a_non_mapping_yaml_document(tmp_path) -> None:
         load_config(config_path)
 
 
+async def test_vision_handler_closes_nested_tool_stream() -> None:
+    closed = asyncio.Event()
+    blocked = asyncio.Event()
+
+    class Vision:
+        async def stream(self, _request):
+            try:
+                yield SimpleNamespace(text="first")
+                await blocked.wait()
+            finally:
+                closed.set()
+
+    handler = app._make_vision_handler(Vision())  # pyright: ignore[reportArgumentType]
+    response = await handler(
+        VoiceQuery(
+            participant_id="alice",
+            text="What is shown?",
+            fresh_match=True,
+            timestamp_us=123,
+        )
+    )
+    assert not isinstance(response, str)
+
+    assert await anext(response) == "first"
+    close = getattr(response, "aclose")
+    await close()
+
+    assert closed.is_set()
+
+
 async def test_app_wires_text_voice_cleanup_readiness_and_shutdown(
     monkeypatch,
     tmp_path,