diff --git a/packages/ai/src/generation.ts b/packages/ai/src/generation.ts index ed52314f8b4..fd110c536cf 100644 --- a/packages/ai/src/generation.ts +++ b/packages/ai/src/generation.ts @@ -111,9 +111,10 @@ export class Generation { } /** - * Status observations as a stream, ending after the first terminal observation. Each poll is bounded by the time - * remaining until `poll.timeout`, so a hung status request fails the stream instead of stalling it. (`Stream.interruptWhen` - * would express this directly but deadlocks under `TestClock` when the source completes while the timer sleeps.) + * Status observations as a stream, ending after the first terminal observation. Each poll and each sleep between polls + * is bounded by the time remaining until `poll.timeout`, so a hung status request or a long interval fails the stream at + * the deadline instead of stalling it. (`Stream.interruptWhen` would express this directly but deadlocks under + * `TestClock` when the source completes while the timer sleeps.) */ events(options?: AwaitOptions): Stream.Stream { if (this.terminal) return Stream.make(this.event()) @@ -122,17 +123,26 @@ export class Generation { Clock.currentTimeMillis.pipe( Effect.map((start) => { const deadline = start + Duration.toMillis(timeout) + // Fail before polling once the deadline has passed: a fast status request could otherwise win the zero-budget + // race and schedule another zero-delay poll. const refresh = Clock.currentTimeMillis.pipe( Effect.flatMap((now) => - this.refresh().pipe( - Effect.timeoutOrElse({ - duration: Duration.millis(Math.max(0, deadline - now)), - orElse: () => this.timeoutError(timeout), - }), - ), + now >= deadline + ? this.timeoutError(timeout) + : this.refresh().pipe( + Effect.timeoutOrElse({ + duration: Duration.millis(deadline - now), + orElse: () => this.timeoutError(timeout), + }), + ), ), ) - return Stream.fromEffectSchedule(refresh, this.schedule(options?.poll)).pipe( + const schedule = this.schedule(options?.poll).pipe( + Schedule.modifyDelay((meta) => + Effect.succeed(Duration.min(meta.duration, Duration.millis(Math.max(0, deadline - meta.now)))), + ), + ) + return Stream.fromEffectSchedule(refresh, schedule).pipe( Stream.takeUntil((generation) => generation.terminal), Stream.map((generation) => generation.event()), ) diff --git a/packages/ai/test/generation.test.ts b/packages/ai/test/generation.test.ts index 5b49558210d..b3ea1be7491 100644 --- a/packages/ai/test/generation.test.ts +++ b/packages/ai/test/generation.test.ts @@ -97,6 +97,26 @@ describe("Generation", () => { }), ) + it.effect("fails an event stream at the deadline when the poll interval is longer than the timeout", () => + Effect.gen(function* () { + const scripted = yield* scriptedRoute(["running"], "never") + const generation = new Generation(scripted.route, "t", { id: "gen_1", status: "queued" }) + + const fiber = yield* Effect.forkChild( + generation + .events({ poll: { interval: "30 seconds", timeout: "10 seconds" } }) + .pipe(Stream.runCollect, Effect.flip), + ) + yield* TestClock.adjust("9 seconds") + expect(fiber.pollUnsafe()).toBeUndefined() + yield* TestClock.adjust("1 second") + const error = yield* Fiber.join(fiber) + + expect(error.reason._tag).toBe("Timeout") + expect(yield* Ref.get(scripted.polls)).toBe(1) + }), + ) + it.effect("surfaces the route failure body for failed generations", () => Effect.gen(function* () { const scripted = yield* scriptedRoute(["running", "failed"], "unused")