diff --git a/.changeset/bounded-toolkit-progress.md b/.changeset/bounded-toolkit-progress.md new file mode 100644 index 00000000000..288709f016e --- /dev/null +++ b/.changeset/bounded-toolkit-progress.md @@ -0,0 +1,5 @@ +--- +"effect": patch +--- + +Bound each Toolkit handler's result buffer to 16 entries so preliminary results backpressure slow consumers. Interrupt handlers when their result stream closes, including early termination and cancellation. diff --git a/packages/effect/src/unstable/ai/Toolkit.ts b/packages/effect/src/unstable/ai/Toolkit.ts index e8f378d1d6b..47a93b48fb4 100644 --- a/packages/effect/src/unstable/ai/Toolkit.ts +++ b/packages/effect/src/unstable/ai/Toolkit.ts @@ -122,6 +122,8 @@ export interface HandlerContext { * * Preliminary results are streamed to the caller before the handler completes, * enabling real-time progress updates for lengthy operations. + * Each tool call buffers up to 16 results. This effect suspends while that + * buffer is full, until the caller consumes results from the returned stream. */ readonly preliminary: (result: Tool.Success) => Effect.Effect } @@ -336,7 +338,7 @@ const Proto = { readonly result: any readonly isFailure: boolean readonly preliminary: boolean - }, Cause.Done>() + }, Cause.Done>({ capacity: 16 }) const context: HandlerContext = { toolCallId, preliminary: (result) => @@ -391,7 +393,7 @@ const Proto = { const encodedResult = yield* encodeResult(output.result, output.isFailure) return { ...output, encodedResult } })), - Stream.onEnd(Fiber.interrupt(fiber)) + Stream.ensuring(Fiber.interrupt(fiber)) ) satisfies Stream.Stream, any> }) diff --git a/packages/effect/test/unstable/ai/Tool.test.ts b/packages/effect/test/unstable/ai/Tool.test.ts index 77cca7a18f5..72bfe9110ac 100644 --- a/packages/effect/test/unstable/ai/Tool.test.ts +++ b/packages/effect/test/unstable/ai/Tool.test.ts @@ -392,6 +392,93 @@ describe("Tool", () => { })) describe("Preliminary Results", () => { + it.effect.each(["success", "failure"] as const)( + "should backpressure preliminary results and preserve %s settlement", + (settlement) => + Effect.gen(function*() { + const toolkit = Toolkit.make(Tool.make("Progress", { + success: Schema.Number, + failure: Schema.String, + failureMode: "return" + })) + const full = yield* Latch.make() + let accepted = 0 + const handlers = yield* toolkit.pipe(Effect.provide(toolkit.toLayer({ + Progress: Effect.fnUntraced(function*(_, ctx) { + for (let i = 0; i < 64; i++) { + if (i === 16) yield* full.open + yield* ctx.preliminary(i) + accepted++ + } + return settlement === "success" ? 64 : yield* Effect.fail("failed") + }) + }))) + + const stream = yield* handlers.handle("Progress", {}) + yield* full.await + yield* Effect.yieldNow + strictEqual(accepted, 16) + + const results = yield* Stream.runCollect(stream) + strictEqual(accepted, 64) + deepStrictEqual(results, [ + ...Array.from({ length: 64 }, (_, result) => ({ + result, + encodedResult: result, + isFailure: false, + preliminary: true + })), + { + result: settlement === "success" ? 64 : "failed", + encodedResult: settlement === "success" ? 64 : "failed", + isFailure: settlement === "failure", + preliminary: false + } + ]) + }) + ) + + it.effect.each(["early end", "interruption"] as const)( + "should finalize a backpressured handler on consumer %s", + (termination) => + Effect.gen(function*() { + const toolkit = Toolkit.make(Tool.make("Progress", { success: Schema.Number })) + const full = yield* Latch.make() + const finalized = yield* Latch.make() + const handlers = yield* toolkit.pipe(Effect.provide(toolkit.toLayer({ + Progress: Effect.fnUntraced(function*(_, ctx) { + for (let i = 0; i < 64; i++) { + if (i === 16) yield* full.open + yield* ctx.preliminary(i) + } + return 64 + }, Effect.ensuring(finalized.open)) + }))) + + // Keep the handler's parent alive so parent cleanup cannot hide a leak. + const stream = yield* handlers.handle("Progress", {}) + yield* full.await + yield* Effect.yieldNow + assertFalse(Latch.isOpen(finalized)) + + if (termination === "early end") { + const results = yield* stream.pipe(Stream.take(1), Stream.runCollect) + strictEqual(results.length, 1) + } else { + const consuming = yield* Latch.make() + const consumer = yield* stream.pipe( + Stream.tap(() => consuming.open.pipe(Effect.andThen(Effect.never))), + Stream.runDrain, + Effect.forkChild + ) + yield* consuming.await + yield* Fiber.interrupt(consumer) + } + + assertTrue(Latch.isOpen(finalized)) + }) + ) + it.effect("should not have preliminary results when generateText is used", () => Effect.gen(function*() { const toolkit = Toolkit.make(IncrementalTool)