diff --git a/agent/relay_llm.py b/agent/relay_llm.py index b53bf28dd3..526829194a 100644 --- a/agent/relay_llm.py +++ b/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 diff --git a/apps/desktop/e2e/core/relay-managed-steer.spec.ts b/apps/desktop/e2e/core/relay-managed-steer.spec.ts new file mode 100644 index 0000000000..5899bae28c --- /dev/null +++ b/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() + } +}) diff --git a/tests/agent/test_relay_llm.py b/tests/agent/test_relay_llm.py index 1d393adce8..ac6e6b3aca 100644 --- a/tests/agent/test_relay_llm.py +++ b/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 = []