fix(relay): hand each streamed chunk over before the provider sends the next

With shared metrics on, every provider stream runs through Relay's managed
pipeline. Relay pulls the next provider chunk before it delivers the current
one, and the provider read blocked the stream's private loop, so each chunk
was withheld until the next arrived. On Desktop the reply looked stopped while
the timer kept going and queued follow-ups waited on a turn that was still
running; a steer aborted text the user never saw, stacking user bubbles with
no output between them.

The off-loop read from the previous commit fixes it. Correct its comment to
the real cause, replace the two-loop starvation test (production gives every
stream its own loop, and it bounded wall-clock at 150ms) with a single-stream
contract that deadlocks on the old code, and add a core e2e spec that steers
and queues against a paused provider with Relay managed execution on.

agent/relay_llm.py

@@ -360,9 +360,9 @@ class ManagedLlmStream(Iterator[Any]):
                 run_callback(self._on_stream_created, raw_stream)
             raw_iterator = run_callback(iter, raw_stream)
             while True:
-                # Provider SDK streams are synchronous and their ``next`` commonly blocks on
-                # socket I/O. Do that read off the Relay event loop so concurrent managed
-                # streams can continue consuming their already-arrived responses.
+                # Off the loop: Relay pulls the next provider chunk before it hands over the
+                # current one, so a blocking read here withholds each chunk until the provider
+                # sends the next. Text vanishes for every provider pause and a steer aborts it.
                 chunk, exhausted = await asyncio.to_thread(_next_provider_chunk, run_callback, raw_iterator)
                 if exhausted:
                     break

apps/desktop/e2e/core/relay-managed-steer.spec.ts

@@ -0,0 +1,117 @@
+/**
+ * C2 core: steering and queueing with NeMo Relay managed execution on.
+ *
+ * Opting into shared metrics routes every provider stream through Relay's
+ * managed pipeline. Relay pulls the next provider chunk before it hands over
+ * the current one, so a blocking provider read withheld each chunk until the
+ * provider sent the next: the reply looked stopped while the timer kept going,
+ * and a steer aborted text the user never saw, stacking user bubbles with no
+ * output between them. These steps hold the provider mid-stream and assert the
+ * text is already on screen, then steer and queue on top of it.
+ */
+
+import { expect, type Page, test } from '@playwright/test'
+
+import {
+  coreAppEnv,
+  createCoreSandbox,
+  currentSessionId,
+  launchCoreApp,
+  recordWebSockets,
+  send,
+  waitForInteractive,
+  writeProviderHome
+} from './harness'
+import { assertTranscriptOracle, installDuplicateSampler, type OracleTarget } from './oracle'
+import { gate, startScriptedProvider } from './provider'
+
+const RELAY_MANAGED_EXECUTION = `telemetry:
+  shared_metrics:
+    enabled: true
+`
+
+const nonce = Math.random()
+  .toString(36)
+  .slice(2, 8)
+  .replace(/[^a-z0-9]/g, 'x')
+  .padEnd(4, 'q')
+
+const U = (n: number) => `U${n}-${nonce}`
+const A = (n: number) => `A${n}-${nonce}`
+
+const words = (marker: string, ...rest: string[]) => [
+  `${marker} `,
+  ...rest.map((w, i) => (i === rest.length - 1 ? w : `${w} `))
+]
+
+function viewport(page: Page) {
+  return page.locator('[data-slot="aui_thread-viewport"]').filter({ visible: true }).first()
+}
+
+function stopButton(page: Page) {
+  return page.locator('[data-slot="composer-root"] button[aria-label="Stop"]')
+}
+
+test('steer and queue stay live under Relay managed execution', async () => {
+  const provider = await startScriptedProvider()
+  const sandbox = createCoreSandbox('relay-steer')
+  writeProviderHome(sandbox.hermesHome, provider.url, RELAY_MANAGED_EXECUTION)
+  const { app, page } = await launchCoreApp(coreAppEnv(sandbox))
+  const ws = recordWebSockets(page)
+
+  const finished = (marker: string) =>
+    expect
+      .poll(() => provider.completions.some(c => c.marker === marker && c.finished), {
+        timeout: 120_000,
+        message: `provider finished ${marker}`
+      })
+      .toBe(true)
+
+  const session: OracleTarget = { sessionId: '', expectUserMarkers: [] }
+
+  try {
+    await waitForInteractive(app, page)
+    await installDuplicateSampler(page)
+
+    await test.step('steer lands below the partial reply the user already saw', async () => {
+      const hold = gate()
+      provider.script(U(1), [{ text: words(A(1), 'slow', 'first', 'reply'), holdAfterFirstChunk: hold }])
+      provider.script(U(2), [{ text: words(A(2), 'steered', 'reply') }])
+      await send(page, `${U(1)} slow one`, 'Enter', ws)
+      await provider.streamStarted(U(1))
+      // The provider is paused after its first chunk; that chunk must already be painted.
+      await expect(viewport(page)).toContainText(A(1), { timeout: 10_000 })
+      await send(page, `${U(2)} change course`, 'Enter', ws)
+      await finished(U(2))
+      hold.open()
+      session.sessionId = await currentSessionId(page)
+      session.expectUserMarkers.push(U(1), U(2))
+      await assertTranscriptOracle(page, ws, provider, session, 'relay steer')
+    })
+
+    await test.step('a follow-up queued under a steered reply drains and the turn settles', async () => {
+      const hold = gate()
+      const steerHold = gate()
+      provider.script(U(3), [{ text: words(A(3), 'long', 'running', 'reply'), holdAfterFirstChunk: hold }])
+      provider.script(U(4), [{ text: words(A(4), 'steered', 'again'), holdAfterFirstChunk: steerHold }])
+      provider.script(U(5), [{ text: words(A(5), 'queued', 'reply') }])
+      await send(page, `${U(3)} long one`, 'Enter', ws)
+      await provider.streamStarted(U(3))
+      await expect(viewport(page)).toContainText(A(3), { timeout: 10_000 })
+      await send(page, `${U(4)} steer it`, 'Enter', ws)
+      await provider.streamStarted(U(4))
+      await expect(viewport(page)).toContainText(A(4), { timeout: 10_000 })
+      await send(page, `${U(5)} after that`, 'Control+Enter', ws)
+      hold.open()
+      steerHold.open()
+      await finished(U(5))
+      await expect(stopButton(page)).toHaveCount(0, { timeout: 30_000 })
+      session.expectUserMarkers.push(U(3), U(4), U(5))
+      await assertTranscriptOracle(page, ws, provider, session, 'relay queued follow-up')
+    })
+  } finally {
+    await app.close().catch(() => undefined)
+    await provider.close()
+    sandbox.cleanup()
+  }
+})

tests/agent/test_relay_llm.py

@@ -6,7 +6,6 @@ import asyncio
 import contextvars
 import json
 import threading
-import time
 from types import SimpleNamespace
 
 import pytest
@@ -142,62 +141,6 @@ async def test_async_execution_uses_canonical_relay_operation_name(
     assert observed_names == ["openai.responses"]
 
 
-@pytest.mark.asyncio
-async def test_provider_stream_read_does_not_block_its_event_loop():
-    """A blocked provider read must not prevent a concurrent stream from consuming data."""
-    read_started = threading.Event()
-    release_read = threading.Event()
-
-    class BlockingIterator:
-        def __iter__(self):
-            return self
-
-        def __next__(self):
-            read_started.set()
-            assert release_read.wait(timeout=1)
-            return {"delta": "chunk"}
-
-    class ReadyIterator:
-        def __iter__(self):
-            return self
-
-        def __next__(self):
-            return {"delta": "ready"}
-
-    def provider_stream_for(raw_stream):
-        stream = object.__new__(relay_llm.ManagedLlmStream)
-        stream._stream_factory = lambda _request: raw_stream
-        stream._completed_response_predicate = None
-        stream._on_stream_created = None
-        stream._accept_chunk = None
-        stream._raw_chunks = []
-        stream._provider_completed = False
-        stream._callback_error = None
-        stream._closed = True  # This producer-only test does not own a Relay stream lifecycle.
-        return stream, stream._provider_stream(attempt, {})
-
-    attempt = SimpleNamespace(
-        provider_request=lambda request: request,
-        run_callback=lambda callback, *args: callback(*args),
-    )
-    _slow_stream, slow_provider_stream = provider_stream_for(BlockingIterator())
-    _fast_stream, fast_provider_stream = provider_stream_for(ReadyIterator())
-    release_timer = threading.Timer(0.2, release_read.set)
-    release_timer.start()
-    started_at = time.monotonic()
-    try:
-        slow_task = asyncio.create_task(anext(slow_provider_stream))
-        assert await asyncio.to_thread(read_started.wait, 1)
-        assert await anext(fast_provider_stream) == {"delta": "ready"}
-        assert time.monotonic() - started_at < 0.15
-        assert await slow_task == {"delta": "chunk"}
-    finally:
-        release_read.set()
-        release_timer.cancel()
-        await slow_provider_stream.aclose()
-        await fast_provider_stream.aclose()
-
-
 def test_stream_execution_uses_canonical_relay_operation_name(relay_turn, monkeypatch):
     relay, _turn = relay_turn
     observed_names = []
@@ -434,6 +377,31 @@ def test_managed_stream_does_not_add_sdk_headers_to_strict_callback(relay_turn):
     assert observed == ["provider-native"]
 
 
+def test_managed_stream_delivers_each_chunk_before_the_provider_sends_the_next(relay_turn):
+    """Steering lost streamed text with Relay on: chunk N was withheld until chunk N+1 arrived."""
+    del relay_turn
+    release = threading.Event()
+
+    def paused_provider(_request):
+        yield {"delta": "first"}
+        assert release.wait(10), "first chunk never reached the consumer during the provider pause"
+        yield {"delta": "second"}
+
+    stream = relay_llm.stream(
+        {"payload": "paused"},
+        paused_provider,
+        session_id="session-1",
+        name="paused-native",
+        model_name="paused-model",
+        finalizer=lambda: {"content": "first second"},
+        metadata={"api_mode": "bedrock_converse", "api_request_id": "paused-stream"},
+    )
+
+    assert next(stream) == {"delta": "first"}
+    release.set()
+    assert list(stream) == [{"delta": "second"}]
+
+
 def test_stream_uses_rewritten_request_and_post_intercept_chunks(relay_turn):
     relay, turn = relay_turn
     captured_requests = []