diff --git a/.changeset/quiet-streams-drain.md b/.changeset/quiet-streams-drain.md new file mode 100644 index 0000000..93f292e --- /dev/null +++ b/.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`. diff --git a/packages/agent/src/lib/model-result.ts b/packages/agent/src/lib/model-result.ts index bbc4f04..a9a510f 100644 --- a/packages/agent/src/lib/model-result.ts +++ b/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; } diff --git a/packages/agent/src/lib/reusable-stream.ts b/packages/agent/src/lib/reusable-stream.ts index fdd4ba4..1e4908f 100644 --- a/packages/agent/src/lib/reusable-stream.ts +++ b/packages/agent/src/lib/reusable-stream.ts @@ -28,6 +28,7 @@ export class ReusableReadableStream { private sourceComplete = false; private sourceError: Error | null = null; private pumpStarted = false; + private cancelled = false; private sourceCancelPromise: Promise | null = null; private readonly streamReplay: StreamReplay; private readonly onValue: ((value: T) => void) | undefined; @@ -73,21 +74,12 @@ export class ReusableReadableStream { * 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 { 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 { 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 { return this.sourceCancelPromise; } + private cancelUnstartedSource(): Promise { + 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 { * Cancel the source stream and all consumers */ async cancel(): Promise { + this.cancelled = true; // Cancel all consumers for (const consumer of this.consumers.values()) { consumer.cancelled = true; @@ -358,6 +374,8 @@ export class ReusableReadableStream { // 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 diff --git a/packages/agent/tests/unit/async-tool-background.test.ts b/packages/agent/tests/unit/async-tool-background.test.ts index 8a95b4e..d989add 100644 --- a/packages/agent/tests/unit/async-tool-background.test.ts +++ b/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'); diff --git a/packages/agent/tests/unit/reusable-stream.test.ts b/packages/agent/tests/unit/reusable-stream.test.ts index bd27dd4..144328c 100644 --- a/packages/agent/tests/unit/reusable-stream.test.ts +++ b/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({ + cancel(): void { + sourceCancelled = true; + }, + }); + const stream = new ReusableReadableStream(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(); + const stream = new ReusableReadableStream(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(); + }); });