@convex-dev/agent 0.7.2 → 0.7.3

package/src/vercel/client/streaming.throttle.test.ts

@@ -0,0 +1,144 @@
+import { afterEach, beforeEach, expect, test, vi } from "vitest";
+import { createThread } from "../../client/threads.js";
+import { components, initConvexTest } from "./setup.test.js";
+import { DeltaStreamer } from "./streaming.js";
+
+beforeEach(() => {
+  vi.useFakeTimers();
+  vi.setSystemTime(new Date("2026-01-01T00:00:00Z"));
+});
+
+afterEach(() => {
+  vi.useRealTimers();
+});
+
+test("publishes a quiet tail after preparing an earlier batch takes time", async () => {
+  const t = initConvexTest();
+  let releasePreparation!: () => void;
+  let enteredPreparation!: () => void;
+  const preparing = new Promise<void>((resolve) => {
+    enteredPreparation = resolve;
+  });
+  const release = new Promise<void>((resolve) => {
+    releasePreparation = resolve;
+  });
+  await t.action(async (ctx) => {
+    const threadId = await createThread(ctx, components.agent, {});
+    const streamer = new DeltaStreamer<string>(
+      components.agent,
+      ctx,
+      {
+        throttleMs: 100,
+        compress: null,
+        abortSignal: undefined,
+        onAsyncAbort: async (reason) => {
+          throw new Error(reason);
+        },
+        materialize: async (parts) => {
+          if (parts.includes("first")) {
+            enteredPreparation();
+            await release;
+          }
+          return { parts, fileRefs: [] };
+        },
+      },
+      { threadId, order: 0, stepOrder: 0, format: undefined },
+    );
+    const savedParts = async () => {
+      const deltas = await ctx.runQuery(components.agent.streams.listDeltas, {
+        threadId,
+        cursors: [{ streamId: await streamer.getStreamId(), cursor: 0 }],
+      });
+      return deltas.flatMap((delta) => delta.parts);
+    };
+    try {
+      await streamer.addParts(["first"]);
+      await preparing;
+      await streamer.addParts(["tail"]);
+      await vi.advanceTimersByTimeAsync(250);
+      expect(await savedParts()).toEqual([]);
+      releasePreparation();
+      for (let elapsed = 0; elapsed < 200; elapsed++) {
+        await vi.advanceTimersByTimeAsync(1);
+        if ((await savedParts()).includes("first")) break;
+      }
+      expect(await savedParts()).toEqual(["first"]);
+      await vi.advanceTimersByTimeAsync(50);
+      expect(await savedParts()).toEqual(["first"]);
+      await vi.advanceTimersByTimeAsync(300);
+      expect.soft(await savedParts()).toEqual(["first", "tail"]);
+      const streams = await ctx.runQuery(components.agent.streams.list, {
+        threadId,
+        statuses: ["streaming"],
+      });
+      expect(streams).toHaveLength(1);
+    } finally {
+      releasePreparation();
+      await streamer.finish();
+    }
+    expect(await savedParts()).toEqual(["first", "tail"]);
+  });
+});
+
+test("a wake armed during a write respects the deadline that write set", async () => {
+  const t = initConvexTest();
+  let releasePreparation!: () => void;
+  let enteredPreparation!: () => void;
+  const preparing = new Promise<void>((resolve) => {
+    enteredPreparation = resolve;
+  });
+  const release = new Promise<void>((resolve) => {
+    releasePreparation = resolve;
+  });
+  await t.action(async (ctx) => {
+    const threadId = await createThread(ctx, components.agent, {});
+    const streamer = new DeltaStreamer<string>(
+      components.agent,
+      ctx,
+      {
+        throttleMs: 100,
+        compress: null,
+        abortSignal: undefined,
+        onAsyncAbort: async () => {},
+        materialize: async (parts) => {
+          if (parts.includes("first")) {
+            enteredPreparation();
+            await release;
+          }
+          return { parts, fileRefs: [] };
+        },
+      },
+      { threadId, order: 0, stepOrder: 0, format: undefined },
+    );
+    const streamId = await streamer.getStreamId();
+    const savedParts = async () => {
+      const deltas = await ctx.runQuery(components.agent.streams.listDeltas, {
+        threadId,
+        cursors: [{ streamId, cursor: 0 }],
+      });
+      return deltas.flatMap((delta) => delta.parts);
+    };
+
+    await streamer.addParts(["first"]);
+    await preparing;
+    // The tail is admitted while the first batch is still being prepared, so
+    // the wake it arms is measured against the previous write's deadline.
+    await vi.advanceTimersByTimeAsync(10);
+    await streamer.addParts(["tail"]);
+    releasePreparation();
+    for (let elapsed = 0; elapsed < 200; elapsed++) {
+      await vi.advanceTimersByTimeAsync(1);
+      if ((await savedParts()).includes("first")) break;
+    }
+    expect(await savedParts()).toEqual(["first"]);
+
+    // The first batch has now reset the window, so the tail is not due yet.
+    await vi.advanceTimersByTimeAsync(50);
+    expect(await savedParts()).toEqual(["first"]);
+
+    // It still has to arrive once the window it was re-armed against elapses,
+    // otherwise holding it back would pass this test by never publishing.
+    await vi.advanceTimersByTimeAsync(100);
+    expect(await savedParts()).toEqual(["first", "tail"]);
+  });
+});

package/src/vercel/client/streaming.ts

@@ -133,6 +133,11 @@ export async function listStreams(
 }
 
 export type StreamingOptions = {
+  /**
+   * Whether source parts emitted by the model are included in the persisted
+   * delta stream. Defaults to false.
+   */
+  sendSources?: boolean;
   /**
    * The minimum granularity of deltas to save.
    * Note: this is not a guarantee that every delta will be exactly one line.
@@ -156,6 +161,7 @@ export type StreamingOptions = {
   returnImmediately?: boolean;
 };
 export const DEFAULT_STREAMING_OPTIONS = {
+  sendSources: false,
   // This chunks by sentences / clauses. Punctuation followed by whitespace.
   chunking: /[\p{P}\s]/u,
   throttleMs: 250,
@@ -215,12 +221,22 @@ export class DeltaStreamer<T> {
   #nextParts: T[] = [];
   #latestWrite: number = 0;
   #ongoingWrite: Promise<void> | undefined;
+  #flushTimer: ReturnType<typeof setTimeout> | undefined;
   #abortPromise: Promise<void> | undefined;
   #cursor: number = 0;
   public abortController: AbortController;
-  // When true, the stream will be finished externally (e.g., atomically via addMessages)
-  // and consumeStream should skip calling finish().
-  #finishedExternally: boolean = false;
+  /**
+   * When true, external code finishes the stream row (atomically with the
+   * message save, issue #181) and `consumeStream` must not. Decided once, at
+   * construction.
+   */
+  #finishHandledExternally: boolean;
+  /**
+   * Set only where the row is finished before the source is drained. An
+   * optimization, not a race guard: `addDelta` already returns false for a
+   * non-streaming row.
+   */
+  #stoppedAccepting: boolean = false;
 
   constructor(
     public readonly component: AgentComponent,
@@ -230,6 +246,8 @@ export class DeltaStreamer<T> {
       onAsyncAbort: (reason: string) => Promise<void>;
       abortSignal: AbortSignal | undefined;
       compress: ((parts: T[]) => T[]) | null;
+      /** Defaults to false, meaning `consumeStream` finishes the stream. */
+      finishHandledExternally?: boolean;
       materialize?: (parts: T[]) => Promise<{
         parts: T[];
         fileRefs: Array<{ url: string; fileId: string }>;
@@ -253,6 +271,7 @@ export class DeltaStreamer<T> {
       compress: config.compress,
       materialize: config.materialize ?? null,
     };
+    this.#finishHandledExternally = config.finishHandledExternally ?? false;
     this.#nextParts = [];
     this.abortController = new AbortController();
     if (config.abortSignal) {
@@ -294,23 +313,64 @@ export class DeltaStreamer<T> {
     if (this.abortController.signal.aborted) {
       return;
     }
-    // Once the stream has been finished externally (e.g. by the inline
-    // save in streamText's onStepFinish for the returnImmediately path),
-    // the stream record is already "finished" in the DB. Late deltas
-    // would be silently dropped by streams.addDelta — skip the work.
-    if (this.#finishedExternally) {
+    if (this.#stoppedAccepting) {
       return;
     }
+    // Buffer before awaiting: a part parked in stream creation would be
+    // invisible to #flushPendingParts.
+    this.#nextParts.push(...parts);
     const streamId = await this.getOrCreateStreamId({
       ifAborted: "returnUndefined",
     });
     if (!streamId) return;
-    this.#nextParts.push(...parts);
+    if (this.#stoppedAccepting || this.abortController.signal.aborted) {
+      return;
+    }
     if (
       !this.#ongoingWrite &&
       Date.now() - this.#latestWrite >= this.config.throttleMs
     ) {
       this.#ongoingWrite = this.#sendDelta();
+    } else {
+      this.#scheduleFlush();
+    }
+  }
+
+  // The throttle is only reconsidered when the next part arrives, so a pause in
+  // the stream would otherwise hold whatever is buffered until it resumes.
+  #scheduleFlush() {
+    if (this.#flushTimer) {
+      return;
+    }
+    const wait = Math.max(
+      0,
+      this.config.throttleMs - (Date.now() - this.#latestWrite),
+    );
+    this.#flushTimer = setTimeout(() => {
+      this.#flushTimer = undefined;
+      if (
+        this.#ongoingWrite ||
+        this.#nextParts.length === 0 ||
+        this.#stoppedAccepting ||
+        this.abortController.signal.aborted
+      ) {
+        return;
+      }
+      // A write can start and land while this wake is armed, which moves the
+      // deadline out from under it. Re-arm against the current one rather
+      // than publishing early.
+      if (Date.now() - this.#latestWrite < this.config.throttleMs) {
+        this.#scheduleFlush();
+        return;
+      }
+      this.#ongoingWrite = this.#sendDelta();
+    }, wait);
+  }
+
+  #cancelScheduledFlush() {
+    if (this.#flushTimer) {
+      clearTimeout(this.#flushTimer);
+      this.#flushTimer = undefined;
     }
   }
 
@@ -326,24 +386,71 @@ export class DeltaStreamer<T> {
       await this.#abort(errorToString(error)).catch(() => {});
       throw error;
     }
-    // Skip finish if it will be handled externally (atomically with message save)
-    // or if the stream was aborted (e.g., due to a failed delta write).
     // Abort cleanup owns the terminal component transition, so consumeStream
     // must wait for it instead of also trying to finish the stream.
     if (this.abortController.signal.aborted) {
       await this.#waitForAbortCleanup();
-    } else if (!this.#finishedExternally) {
+      return;
+    }
+    // EOF is the only point where every part has actually been handed over:
+    // parts sit in the AI SDK's transform and tee pipeline until the iterator
+    // yields them, so no earlier callback can observe them.
+    try {
+      await this.#flushPendingParts();
+    } catch (error) {
+      if (this.abortController.signal.aborted) {
+        await this.#waitForAbortCleanup();
+      }
+      throw error;
+    }
+    if (this.abortController.signal.aborted) {
+      await this.#waitForAbortCleanup();
+      return;
+    }
+    if (!this.#finishHandledExternally) {
       await this.finish();
     }
   }
 
   /**
-   * Mark the stream as being finished externally (e.g., atomically via addMessages).
-   * When called, consumeStream() will skip calling finish() since it will be
-   * handled elsewhere in the same mutation as message saving.
+   * Drain everything currently buffered or in flight, so that after this
+   * resolves no delta write is outstanding and #nextParts is empty.
    */
-  public markFinishedExternally(): void {
-    this.#finishedExternally = true;
+  async #flushPendingParts(): Promise<void> {
+    while (!this.abortController.signal.aborted) {
+      const inFlight = this.#ongoingWrite;
+      await inFlight;
+      // #sendDelta reassigns #ongoingWrite from its own tail, so a write can
+      // still be live even though the buffer it drained is now empty.
+      if (this.#ongoingWrite !== inFlight) {
+        continue;
+      }
+      if (this.#nextParts.length === 0) {
+        break;
+      }
+      this.#ongoingWrite = this.#sendDelta();
+    }
+  }
+
+  /**
+   * For the `returnImmediately` path, where nothing awaits consumption and the
+   * save has to happen inline (issue #265).
+   *
+   * Inherent window: parts still inside the AI SDK pipeline at this instant
+   * never reach addParts, so they are never persisted as deltas (the
+   * stream-level `finish` chunk among them). The message saved alongside this
+   * transition is complete, and is authoritative once the stream is finished.
+   *
+   * Narrower race left open: a part admitted after the loop's last emptiness
+   * check is buffered but not drained here, and the write it triggers at EOF
+   * lands on a row the caller has since finished. `#sendDelta` tolerates that
+   * refusal rather than aborting. Closing it means refusing admission for the
+   * duration of the drain, which drops more of the tail than it saves.
+   */
+  public async flushAndStopAccepting(): Promise<void> {
+    this.#cancelScheduledFlush();
+    await this.#flushPendingParts();
+    this.#stoppedAccepting = true;
   }
 
   /**
@@ -373,11 +480,13 @@ export class DeltaStreamer<T> {
   }
 
   async #sendDelta() {
+    this.#cancelScheduledFlush();
     if (this.abortController.signal.aborted) {
       return;
     }
     let success: boolean;
     try {
+      await this.getStreamId();
       const delta = await this.#createDelta();
       if (!delta) {
         return;
@@ -392,11 +501,11 @@ export class DeltaStreamer<T> {
       return;
     }
     if (!success) {
-      // An in-flight #sendDelta started before markFinishedExternally()
+      // A #sendDelta racing the inline save on the returnImmediately path
       // will get `success === false` because the stream row is already
       // "finished". That's a benign late-write miss, not a failure —
       // don't convert it into an abort.
-      if (this.#finishedExternally) {
+      if (this.#stoppedAccepting) {
         return;
       }
       await this.#abortDelta("async abort");
@@ -411,6 +520,11 @@ export class DeltaStreamer<T> {
       this.#ongoingWrite = this.#sendDelta();
     } else {
       this.#ongoingWrite = undefined;
+      // Whatever is still buffered has just lost its owner: this write is
+      // over and no arrival is guaranteed to follow. Hand it to the timer.
+      if (this.#nextParts.length > 0) {
+        this.#scheduleFlush();
+      }
     }
   }
 
@@ -447,6 +561,7 @@ export class DeltaStreamer<T> {
   }
 
   public async finish() {
+    this.#cancelScheduledFlush();
     if (!this.streamId) {
       return;
     }
@@ -491,6 +606,7 @@ export class DeltaStreamer<T> {
 
   #abort(reason: string, waitForOngoingWrite = true): Promise<void> {
     if (!this.#abortPromise) {
+      this.#cancelScheduledFlush();
       this.abortController.abort();
       this.#abortPromise = this.#abortCreatedStream(
         reason,