fix(reconciler): bound the Flight EOF wait by the render deadline (#204 follow-up)

.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.

packages/rsc-runtime/src/reconciler.ts

@@ -300,20 +300,26 @@ type SettledBoundary = {
   readonly ok: boolean;
 };
 
-const waitPendingOrDeadline = (
-  pending: readonly PendingBoundary[],
+const waitOrDeadline = <A>(
+  wait: Effect.Effect<A, Error>,
   sequence: AgentRenderEventSequence,
-): Effect.Effect<SettledBoundary, Error> => {
+): Effect.Effect<A, Error> => {
   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<SettledBoundary, Error> =>
+  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({

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<void>((resolve) => {
+        resolveCancelled = resolve;
+      });
+      const host: AgentFlightExecutionHost = {
+        execute: async () =>
+          new ReadableStream<Uint8Array>({
+            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: () => {