diff --git a/.changeset/bound-flight-eof-wait.md b/.changeset/bound-flight-eof-wait.md new file mode 100644 index 000000000..38656cfef --- /dev/null +++ b/.changeset/bound-flight-eof-wait.md @@ -0,0 +1,6 @@ +--- +"@agent-bundle/runtime": patch +--- + +Bound the Flight EOF wait by the render deadline and cancel stalled Flight +sources when the elapsed-time limit is exceeded. diff --git a/packages/rsc-runtime/src/reconciler.ts b/packages/rsc-runtime/src/reconciler.ts index 150a5ce8e..534218b29 100644 --- a/packages/rsc-runtime/src/reconciler.ts +++ b/packages/rsc-runtime/src/reconciler.ts @@ -300,20 +300,26 @@ type SettledBoundary = { readonly ok: boolean; }; -const waitPendingOrDeadline = ( - pending: readonly PendingBoundary[], +const waitOrDeadline = ( + wait: Effect.Effect, sequence: AgentRenderEventSequence, -): Effect.Effect => { +): Effect.Effect => { const remaining = sequence.remainingMs; if (remaining <= 0) return Effect.fail(elapsedTimeExceeded(sequence.maxElapsedMs)); return Effect.raceFirst( - waitSettledBoundary(pending), + wait, Effect.sleep(Duration.millis(remaining)).pipe( Effect.flatMap(() => Effect.fail(elapsedTimeExceeded(sequence.maxElapsedMs))), ), ); }; +const waitPendingOrDeadline = ( + pending: readonly PendingBoundary[], + sequence: AgentRenderEventSequence, +): Effect.Effect => + waitOrDeadline(waitSettledBoundary(pending), sequence); + const settledBoundaryInputs = ( previous: TreeSnapshot, next: TreeSnapshot, @@ -365,7 +371,7 @@ const reconcileLoopStream = ( Error > => { if (snapshot.pending.length === 0) { - return Deferred.await(flightDone).pipe( + return waitOrDeadline(Deferred.await(flightDone), sequence).pipe( Effect.andThen(Queue.clear(progressInputs)), Effect.flatMap((queued) => Effect.try({ diff --git a/packages/rsc-runtime/tests/dispatcher.test.ts b/packages/rsc-runtime/tests/dispatcher.test.ts index 4c96b2c8b..3ffaaa87e 100644 --- a/packages/rsc-runtime/tests/dispatcher.test.ts +++ b/packages/rsc-runtime/tests/dispatcher.test.ts @@ -559,6 +559,59 @@ describe('AgentRenderDispatcher streaming', () => { await expect(reader.read()).rejects.toMatchObject({ code: 'elapsed-time-exceeded' }); }); + it('bounds Flight EOF for stream and dispatch and cancels the source', { retry: 2, timeout: 5_000 }, async () => { + const finiteReader = ( + await createWorkerHost('ready').execute({ + invocation, + signal: new AbortController().signal, + }) + ).getReader(); + const chunks: Uint8Array[] = []; + while (true) { + const next = await finiteReader.read(); + if (next.done) break; + chunks.push(next.value); + } + + for (const mode of ['stream', 'dispatch'] as const) { + let cancelCalls = 0; + let resolveCancelled: (() => void) | undefined; + const cancelled = new Promise((resolve) => { + resolveCancelled = resolve; + }); + const host: AgentFlightExecutionHost = { + execute: async () => + new ReadableStream({ + cancel() { + cancelCalls += 1; + resolveCancelled?.(); + }, + start(controller) { + for (const chunk of chunks) controller.enqueue(chunk); + }, + }), + }; + const dispatcher = createAgentRenderDispatcher(host, { limits: { maxElapsedMs: 150 } }); + const signal = new AbortController().signal; + const startedAt = Date.now(); + + if (mode === 'stream') { + const reader = dispatcher.stream({ invocation, signal }).getReader(); + const shell = await reader.read(); + if (shell.value?.type !== 'shell') throw new Error('expected a shell event'); + await expect(reader.read()).rejects.toMatchObject({ code: 'elapsed-time-exceeded' }); + } else { + await expect(dispatcher.dispatch({ invocation, signal })).rejects.toMatchObject({ + code: 'elapsed-time-exceeded', + }); + } + + expect(Date.now() - startedAt).toBeLessThan(1_000); + await cancelled; + expect(cancelCalls).toBe(1); + } + }); + it('converts a synchronous host throw into a stream failure', async () => { const host: AgentFlightExecutionHost = { execute: () => {