@convex-dev/agent 0.7.1 → 0.7.3

package/src/vercel/client/streaming.ts

@@ -8,6 +8,7 @@ import {
   type UIMessageChunk,
 } from "ai";
 import { v } from "convex/values";
+import { errorToString } from "../../errors.js";
 import {
   vMessageDoc,
   vPaginationResult,
@@ -132,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.
@@ -155,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,
@@ -214,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,
@@ -229,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 }>;
@@ -252,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) {
@@ -293,20 +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;
     }
-    await this.getStreamId();
+    // 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;
+    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;
     }
   }
 
@@ -319,45 +383,110 @@ export class DeltaStreamer<T> {
       // A provider can throw while responding to an abort. Join the durable
       // abort transition here, outside the active delta writer, before
       // preserving the provider error for the caller.
-      await this.#abort(
-        error instanceof Error ? error.message : "stream consumption failed",
-      ).catch(() => {});
+      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;
   }
 
   /**
    * Get the stream ID, waiting for it to be created if necessary.
    * Useful for passing to addMessages for atomic finish.
    */
-  public async getOrCreateStreamId(): Promise<string> {
-    return this.getStreamId();
+  public async getOrCreateStreamId(): Promise<string>;
+  public async getOrCreateStreamId(options: {
+    ifAborted: "returnUndefined";
+  }): Promise<string | undefined>;
+  public async getOrCreateStreamId(options?: {
+    ifAborted?: "returnUndefined";
+  }): Promise<string | undefined> {
+    if (options?.ifAborted !== "returnUndefined") {
+      return this.getStreamId();
+    }
+    if (this.abortController.signal.aborted) {
+      await this.#abortPromise;
+      return undefined;
+    }
+    const streamId = await this.getStreamId();
+    if (this.abortController.signal.aborted) {
+      await this.#abortPromise;
+      return undefined;
+    }
+    return streamId;
   }
 
   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;
@@ -368,15 +497,15 @@ export class DeltaStreamer<T> {
         delta,
       );
     } catch (e) {
-      await this.#abortDelta(e instanceof Error ? e.message : "unknown error");
+      await this.#abortDelta(errorToString(e));
       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");
@@ -391,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();
+      }
     }
   }
 
@@ -427,6 +561,7 @@ export class DeltaStreamer<T> {
   }
 
   public async finish() {
+    this.#cancelScheduledFlush();
     if (!this.streamId) {
       return;
     }
@@ -471,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,

package/src/vercel/client/types.ts

@@ -203,6 +203,8 @@ export type ContextOptions = {
   /**
    * How many recent messages to include. These are added after the search
    * messages, and do not count against the search limit.
+   * If the limit splits an order, additional messages are included to complete
+   * that order. If it cannot be completed, the incomplete order is omitted.
    * Default: 100
    */
   recentMessages?: number;

package/src/vercel/client/utils.test.ts

@@ -0,0 +1,97 @@
+import type { StepResult } from "ai";
+import { describe, expect, test } from "vitest";
+import { hasSuccessfulToolCall, willContinue } from "./utils.js";
+
+// Minimal StepResult builder — only the fields willContinue and
+// hasSuccessfulToolCall actually read. Loosely typed on purpose so test
+// fixtures can be terse; cast at the boundary.
+type StepFixture = {
+  finishReason?: string;
+  content?: Array<{ type: string; toolName?: string }>;
+  toolCalls?: Array<{ toolCallId: string; toolName: string }>;
+  toolResults?: Array<{ toolCallId: string; toolName: string }>;
+};
+
+function makeStep(partial: StepFixture): StepResult<any> {
+  return {
+    finishReason: "tool-calls",
+    content: [],
+    toolCalls: [],
+    toolResults: [],
+    ...partial,
+  } as unknown as StepResult<any>;
+}
+
+describe("hasSuccessfulToolCall", () => {
+  test("returns true when last step has a tool-result for the named tool", () => {
+    const step = makeStep({
+      content: [{ type: "tool-result", toolName: "search" }],
+    });
+    expect(hasSuccessfulToolCall("search")({ steps: [step] })).toBe(true);
+  });
+
+  test("returns false when only a tool-error is present for the named tool", () => {
+    const step = makeStep({
+      content: [{ type: "tool-error", toolName: "search" }],
+    });
+    expect(hasSuccessfulToolCall("search")({ steps: [step] })).toBe(false);
+  });
+
+  test("returns false when the matching tool name is missing", () => {
+    const step = makeStep({
+      content: [{ type: "tool-result", toolName: "other" }],
+    });
+    expect(hasSuccessfulToolCall("search")({ steps: [step] })).toBe(false);
+  });
+
+  test("only inspects the last step", () => {
+    const earlier = makeStep({
+      content: [{ type: "tool-result", toolName: "search" }],
+    });
+    const last = makeStep({
+      content: [{ type: "tool-error", toolName: "search" }],
+    });
+    expect(hasSuccessfulToolCall("search")({ steps: [earlier, last] })).toBe(
+      false,
+    );
+  });
+
+  test("returns false when steps is empty", () => {
+    expect(hasSuccessfulToolCall("search")({ steps: [] })).toBe(false);
+  });
+});
+
+describe("willContinue", () => {
+  test("does not stop when a tool-error fills in for a missing tool-result", async () => {
+    // Two tool calls; one returns a result, the other errors.
+    const step = makeStep({
+      toolCalls: [
+        { toolCallId: "1", toolName: "a" },
+        { toolCallId: "2", toolName: "b" },
+      ],
+      toolResults: [{ toolCallId: "1", toolName: "a" }],
+      content: [
+        { type: "tool-result", toolName: "a" },
+        { type: "tool-error", toolName: "b" },
+      ],
+    });
+    // No stopWhen → returns false (no further stop conditions). The point
+    // is the function progresses past the early `toolCalls > completed`
+    // bail; pre-fix it returned early because tool-error wasn't counted.
+    expect(await willContinue([step], undefined)).toBe(false);
+  });
+
+  test("stops when a tool call has neither a result nor an error yet", async () => {
+    const step = makeStep({
+      toolCalls: [{ toolCallId: "1", toolName: "a" }],
+      toolResults: [],
+      content: [],
+    });
+    expect(await willContinue([step], undefined)).toBe(false);
+  });
+
+  test("stops when finishReason is not tool-calls", async () => {
+    const step = makeStep({ finishReason: "stop" });
+    expect(await willContinue([step], undefined)).toBe(false);
+  });
+});

package/src/vercel/client/utils.ts

@@ -1,17 +1,23 @@
 import type { Context } from "@ai-sdk/provider-utils";
 import type { StepResult, StopCondition, ToolSet } from "ai";
 
+export { errorToString } from "../../errors.js";
+
 /**
- * A stop condition that only matches tool calls which completed
- * successfully (i.e. produced a `tool-result`, not a `tool-error`).
+ * A stop condition that only matches tool calls of the given name which
+ * completed successfully — i.e. produced a `tool-result` content part.
+ * Failed tool calls (which surface as `tool-error` parts under AI SDK v6)
+ * do not match.
  *
- * Use this instead of the AI SDK's `hasToolCall` when you want the
- * agent to retry on argument validation failures rather than stopping.
+ * Use this instead of the AI SDK's `hasToolCall` when you want the agent
+ * to retry on argument-validation or runtime tool failures rather than
+ * stopping. Evaluated only against the last step (consistent with how
+ * `stopWhen` is applied after each step).
  */
 export function hasSuccessfulToolCall(toolName: string): StopCondition<any> {
   return ({ steps }) =>
-    steps[steps.length - 1]?.toolResults?.some(
-      (result) => result.toolName === toolName,
+    steps[steps.length - 1]?.content?.some(
+      (p) => p.type === "tool-result" && p.toolName === toolName,
     ) ?? false;
 }
 
@@ -31,12 +37,13 @@ export async function willContinue<
   if (step.finishReason !== "tool-calls") return false;
   // Count both successful results and errors as completed outputs.
   // Failed tool calls are represented as tool-error content parts, so only
-  // checking toolResults misses them.
+  // checking toolResults misses them. The fallback to step.toolResults.length
+  // is for callers whose steps lack content (mocks); the optional chain is
+  // defensive rather than load-bearing.
   const completedOutputs =
     step.content?.filter(
       (p) => p.type === "tool-result" || p.type === "tool-error",
     ).length ?? step.toolResults.length;
-  // we don't have a tool result, so we'll wait for more
   if (step.toolCalls.length > completedOutputs) return false;
   if (Array.isArray(stopWhen)) {
     return (await Promise.all(stopWhen.map(async (s) => s({ steps })))).every(
@@ -45,10 +52,3 @@ export async function willContinue<
   }
   return !!stopWhen && !(await stopWhen({ steps }));
 }
-
-export function errorToString(error: unknown): string {
-  if (error instanceof Error) {
-    return error.message;
-  }
-  return String(error);
-}

package/src/vercel/index.ts

@@ -818,6 +818,7 @@ export class Agent<
       metadata: args.metadata ? [args.metadata] : undefined,
       skipEmbeddings: args.skipEmbeddings,
       promptMessageId: args.promptMessageId,
+      order: args.order,
       pendingMessageId: args.pendingMessageId,
     });
     const message = messages.at(-1)!;
@@ -1765,6 +1766,7 @@ export class Agent<
         threadId: v.string(),
         userId: v.optional(v.string()),
         promptMessageId: v.optional(v.string()),
+        order: v.optional(v.union(v.number(), v.literal("next"))),
         messages: v.array(vMessageWithMetadata),
         failPendingSteps: v.optional(v.boolean()),
         embeddings: v.optional(vMessageEmbeddings),