fix(ai): bound generation event poll sleep by the deadline (#51372)

packages/ai/src/generation.ts

@@ -111,9 +111,10 @@ export class Generation<Response> {
   }
 
   /**
-   * 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<Event, AIError> {
     if (this.terminal) return Stream.make(this.event())
@@ -122,17 +123,26 @@ export class Generation<Response> {
       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()),
           )

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")