fix(agent): close two run-teardown races in ReusableReadableStream and async drain

.changeset/quiet-streams-drain.md

@@ -0,0 +1,5 @@
+---
+'@openrouter/agent': patch
+---
+
+Fix two run-teardown races. `ReusableReadableStream.createConsumer()` called while `cancel()` is awaiting the source reader now yields an already-done consumer instead of one that reads a cleared buffer slot and throws. `onRunEnd: 'drain'` now persists and broadcasts (`delivery: 'dropped'`) a background task settlement that lands during the final drain turn, where it was previously left unharvested with the persisted task still `working`.

packages/agent/src/lib/model-result.ts

@@ -4983,11 +4983,10 @@ export class ModelResult<
       }
     }
 
-    // Drain budget exhausted with work still in flight: cut it loose.
     if (registry.hasInFlight()) {
       registry.abortAll('Async tool drain budget exhausted at run end');
-      await this.dropSettledTasks();
     }
+    await this.dropSettledTasks();
 
     return response;
   }

packages/agent/src/lib/reusable-stream.ts

@@ -28,6 +28,7 @@ export class ReusableReadableStream<T> {
   private sourceComplete = false;
   private sourceError: Error | null = null;
   private pumpStarted = false;
+  private cancelled = false;
   private sourceCancelPromise: Promise<void> | null = null;
   private readonly streamReplay: StreamReplay;
   private readonly onValue: ((value: T) => void) | undefined;
@@ -73,21 +74,12 @@ export class ReusableReadableStream<T> {
    * Create a new consumer that can independently iterate over the stream.
    * Full-replay consumers start at position 0. Active-consumer replay starts
    * at the current trim watermark. Multiple attached consumers advance
-   * independently in either mode.
+   * independently in either mode. Consumers created after `cancel()` are
+   * already done.
    */
   createConsumer(): AsyncIterableIterator<T> {
     const consumerId = this.nextConsumerId++;
-    const state: ConsumerState = {
-      position: this.trimOffset,
-      waitingPromise: null,
-      cancelled: false,
-    };
-    this.consumers.set(consumerId, state);
-
-    // Start pumping the source stream if not already started
-    if (!this.pumpStarted) {
-      this.startPump();
-    }
+    this.registerConsumer(consumerId);
 
     // eslint-disable-next-line @typescript-eslint/no-this-alias
     const self = this;
@@ -263,6 +255,22 @@ export class ReusableReadableStream<T> {
     this.bufferHead = 0;
   }
 
+  private registerConsumer(consumerId: number): void {
+    if (this.cancelled) {
+      return;
+    }
+    this.consumers.set(consumerId, {
+      position: this.trimOffset,
+      waitingPromise: null,
+      cancelled: false,
+    });
+
+    // Start pumping the source stream if not already started
+    if (!this.pumpStarted) {
+      this.startPump();
+    }
+  }
+
   /**
    * Start pumping data from the source stream into the buffer
    */
@@ -323,6 +331,13 @@ export class ReusableReadableStream<T> {
     return this.sourceCancelPromise;
   }
 
+  private cancelUnstartedSource(): Promise<void> {
+    if (!this.sourceCancelPromise) {
+      this.sourceCancelPromise = this.sourceStream.cancel();
+    }
+    return this.sourceCancelPromise;
+  }
+
   /**
    * Notify all waiting consumers that new data is available
    */
@@ -343,6 +358,7 @@ export class ReusableReadableStream<T> {
    * Cancel the source stream and all consumers
    */
   async cancel(): Promise<void> {
+    this.cancelled = true;
     // Cancel all consumers
     for (const consumer of this.consumers.values()) {
       consumer.cancelled = true;
@@ -358,6 +374,8 @@ export class ReusableReadableStream<T> {
     // Cancel the source stream
     if (this.sourceReader) {
       await this.cancelSourceReader(this.sourceReader);
+    } else if (!this.pumpStarted) {
+      await this.cancelUnstartedSource();
     }
     /*
      * The pump may have landed one in-flight chunk between the synchronous

packages/agent/tests/unit/async-tool-background.test.ts

@@ -480,6 +480,89 @@ describe('tool.background — placeholder & delivery', () => {
     expect(sawAbort).toBe(true);
   });
 
+  it("onRunEnd: 'drain' reports a task that settles during the last drain turn as dropped", async () => {
+    const first = makeControlledBackgroundTool('render_a');
+    const second = makeControlledBackgroundTool('render_b');
+
+    mockBetaResponsesSend
+      .mockResolvedValueOnce({
+        ok: true,
+        value: makeResponse('resp_1', [
+          functionCallItem('call_a', 'render_a', '{"script":"a"}'),
+          functionCallItem('call_b', 'render_b', '{"script":"b"}'),
+        ]),
+      })
+      .mockImplementationOnce(async () => {
+        first.release({
+          url: 'https://cdn/a.mp4',
+        });
+        return {
+          ok: true,
+          value: makeResponse('resp_2', [
+            messageItem('msg_1', 'both started'),
+          ]),
+        };
+      })
+      // The only drain turn (maxDrainTurns: 1) carries call_a; call_b
+      // settles while the model is still answering, after the loop has
+      // spent its last turn.
+      .mockImplementationOnce(async () => {
+        second.release({
+          url: 'https://cdn/b.mp4',
+        });
+        await new Promise((resolve) => setTimeout(resolve, 10));
+        return {
+          ok: true,
+          value: makeResponse('resp_3', [
+            messageItem('msg_2', 'a is done'),
+          ]),
+        };
+      });
+
+    const result = callModel(client, {
+      model: 'test-model',
+      input: 'render',
+      tools: [
+        first.tool,
+        second.tool,
+      ] as const,
+      asyncTools: {
+        onRunEnd: 'drain',
+        drainTimeoutMs: 5_000,
+        maxDrainTurns: 1,
+      },
+    });
+
+    const settled: Array<{
+      toolCallId: string;
+      delivery: string;
+    }> = [];
+    for await (const event of result.getFullResponsesStream()) {
+      if (isToolAsyncSettledEvent(event)) {
+        settled.push({
+          toolCallId: event.toolCallId,
+          delivery: event.delivery,
+        });
+      }
+    }
+
+    expect(mockBetaResponsesSend).toHaveBeenCalledTimes(3);
+    expect(settled).toEqual([
+      {
+        toolCallId: 'call_a',
+        delivery: 'injected',
+      },
+      {
+        toolCallId: 'call_b',
+        delivery: 'dropped',
+      },
+    ]);
+    expect(result.getAsyncTasks().map((t) => t.status)).toEqual([
+      'completed',
+      'completed',
+    ]);
+  });
+
   it('cancelTask(taskId) cancels a working background task', async () => {
     const controlled = makeControlledBackgroundTool('render_video');
 

packages/agent/tests/unit/reusable-stream.test.ts

@@ -209,4 +209,54 @@ describe('ReusableReadableStream', () => {
     expect((await fresh.next()).done).toBe(true);
     // No source.close(): cancel() already terminated the source stream.
   });
+
+  it('cancel() before the first consumer cancels the source and later consumers are done', async () => {
+    let sourceCancelled = false;
+    const source = new ReadableStream<number>({
+      cancel(): void {
+        sourceCancelled = true;
+      },
+    });
+    const stream = new ReusableReadableStream<number>(source);
+
+    await stream.cancel();
+
+    expect(sourceCancelled).toBe(true);
+    expect(source.locked).toBe(false);
+    await expect(source.getReader().read()).resolves.toEqual({
+      done: true,
+      value: undefined,
+    });
+    expect((await stream.createConsumer().next()).done).toBe(true);
+    await stream.cancel();
+  });
+
+  it('active-consumers: a consumer created while cancel() awaits the source reader is done', async () => {
+    const source = controlledStream<number>();
+    const stream = new ReusableReadableStream<number>(source.stream, {
+      streamReplay: 'active-consumers',
+    });
+    const first = stream.createConsumer();
+    await Promise.resolve();
+    source.push(1);
+
+    /*
+     * The pump has already read a chunk that lands in the buffer while
+     * cancel() awaits sourceReader.cancel(), so the late consumer is created
+     * between the two backlog sweeps.
+     */
+    const cancelPromise = stream.cancel();
+    const late = stream.createConsumer();
+    await cancelPromise;
+
+    expect(await first.next()).toEqual({
+      done: true,
+      value: undefined,
+    });
+    expect(await late.next()).toEqual({
+      done: true,
+      value: undefined,
+    });
+    expect(stream.findLastBuffered(() => true)).toBeUndefined();
+  });
 });