@convex-dev/agent 0.6.0 → 0.6.2

package/src/client/streamText.ts

@@ -80,9 +80,21 @@ export async function streamText<
 
   const steps: StepResult<TOOLS>[] = [];
 
-  // Track the final step for atomic save with stream finish (issue #181)
+  // Track the final step for atomic save with stream finish (issue #181).
+  // Only used when streamText awaits stream consumption itself; the
+  // `returnImmediately` path saves inline instead (see onStepFinish below).
   let pendingFinalStep: StepResult<TOOLS> | undefined;
 
+  // Whether streamText will await stream consumption before returning.
+  // When false (saveStreamDeltas.returnImmediately === true), we cannot
+  // defer the final-step save to a post-await block — the function has
+  // already returned by the time onStepFinish fires. See issue #265.
+  const willAwaitStream =
+    Boolean(threadId) &&
+    (options.saveStreamDeltas === true ||
+      (typeof options.saveStreamDeltas === "object" &&
+        !options.saveStreamDeltas.returnImmediately));
+
   const streamer =
     threadId && options.saveStreamDeltas
       ? new DeltaStreamer(
@@ -142,10 +154,19 @@ export async function streamText<
       steps.push(step);
       const createPendingMessage = await willContinue(steps, args.stopWhen);
       if (!createPendingMessage && streamer) {
-        // This is the final step with streaming enabled.
-        // Defer saving until stream consumption completes for atomic finish (issue #181).
+        // Final step with streaming enabled.
         streamer.markFinishedExternally();
-        pendingFinalStep = step;
+        if (willAwaitStream) {
+          // We're about to `await stream` below — defer the save so it
+          // happens atomically with stream finish (issue #181).
+          pendingFinalStep = step;
+        } else {
+          // returnImmediately path: streamText is about to return without
+          // awaiting consumption, so the deferred-save block below won't
+          // see this step. Save inline now (issue #265).
+          const finishStreamId = await streamer.getOrCreateStreamId();
+          await call.save({ step }, false, finishStreamId);
+        }
       } else {
         await call.save({ step }, createPendingMessage);
       }
@@ -155,11 +176,7 @@ export async function streamText<
   const stream = streamer?.consumeStream(
     result.toUIMessageStream<AIUIMessage<TOOLS>>(),
   );
-  if (
-    (typeof options?.saveStreamDeltas === "object" &&
-      !options.saveStreamDeltas.returnImmediately) ||
-    options?.saveStreamDeltas === true
-  ) {
+  if (willAwaitStream) {
     try {
       await stream;
       await result.consumeStream();
@@ -168,6 +185,15 @@ export async function streamText<
       // finish() was never called, leaving the streaming message stuck in
       // "streaming" state. Clean it up by marking it as aborted.
       await streamer?.fail(e instanceof Error ? e.message : String(e));
+      // Save the deferred final step if it was already generated but not yet persisted
+      if (pendingFinalStep) {
+        try {
+          await call.save({ step: pendingFinalStep }, false);
+        } catch (saveError) {
+          console.error("Failed to save deferred final step:", saveError);
+        }
+        pendingFinalStep = undefined;
+      }
       throw e;
     }
   }

package/src/client/streaming.ts

@@ -244,16 +244,25 @@ export class DeltaStreamer<T> {
     this.abortController = new AbortController();
     if (config.abortSignal) {
       config.abortSignal.addEventListener("abort", async () => {
-        if (this.abortController.signal.aborted) {
-          return;
-        }
-        if (this.streamId) {
+        try {
+          if (this.abortController.signal.aborted) {
+            return;
+          }
           this.abortController.abort();
-          await this.#ongoingWrite;
-          await this.ctx.runMutation(this.component.streams.abort, {
-            streamId: this.streamId,
-            reason: "abortSignal",
-          });
+          // Wait for in-flight stream creation before trying to abort it
+          if (this.#creatingStreamIdPromise) {
+            await this.#creatingStreamIdPromise;
+          }
+          if (this.streamId) {
+            await this.#ongoingWrite;
+            await this.ctx.runMutation(this.component.streams.abort, {
+              streamId: this.streamId,
+              reason: "abortSignal",
+            });
+          }
+        } catch {
+          // Best-effort cleanup — the stream will be garbage-collected
+          // by the 10-minute timeout if this fails.
         }
       });
     }
@@ -278,6 +287,13 @@ 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) {
+      return;
+    }
     await this.getStreamId();
     this.#nextParts.push(...parts);
     if (
@@ -293,7 +309,9 @@ export class DeltaStreamer<T> {
       await this.addParts([chunk]);
     }
     // 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)
+    // or if the stream was aborted (e.g., due to a failed delta write).
+    // Aborted streams are cleaned up via streams.abort (called by the abort
+    // signal handler), so we don't need to call finish() for them.
     if (!this.#finishedExternally && !this.abortController.signal.aborted) {
       await this.finish();
     }
@@ -331,6 +349,13 @@ export class DeltaStreamer<T> {
         delta,
       );
       if (!success) {
+        // An in-flight #sendDelta started before markFinishedExternally()
+        // 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) {
+          return;
+        }
         await this.config.onAsyncAbort("async abort");
         this.abortController.abort();
         return;
@@ -376,10 +401,7 @@ export class DeltaStreamer<T> {
       return;
     }
     await this.#ongoingWrite;
-    if (this.abortController.signal.aborted) {
-      return;
-    }
-    await this.#sendDelta();
+    await this.#sendDelta(); // #sendDelta checks aborted internally
     if (this.abortController.signal.aborted) {
       return;
     }
@@ -439,16 +461,15 @@ export function compressTextStreamParts(
       } else {
         compressed.push(part);
       }
+    } else if (part.type === "file") {
+      compressed.push({
+        type: "file",
+        file: {
+          ...part.file,
+          uint8Array: undefined as unknown as Uint8Array,
+        },
+      });
     } else {
-      if (part.type === "file") {
-        compressed.push({
-          type: "file",
-          file: {
-            ...part.file,
-            uint8Array: undefined as unknown as Uint8Array,
-          },
-        });
-      }
       compressed.push(part);
     }
   }

package/src/client/utils.ts

@@ -1,5 +1,19 @@
 import type { StepResult, StopCondition } from "ai";
 
+/**
+ * A stop condition that only matches tool calls which completed
+ * successfully (i.e. produced a `tool-result`, not a `tool-error`).
+ *
+ * Use this instead of the AI SDK's `hasToolCall` when you want the
+ * agent to retry on argument validation failures rather than stopping.
+ */
+export function hasSuccessfulToolCall(toolName: string): StopCondition<any> {
+  return ({ steps }) =>
+    steps[steps.length - 1]?.toolResults?.some(
+      (result) => result.toolName === toolName,
+    ) ?? false;
+}
+
 export async function willContinue(
   steps: StepResult<any>[],
 
@@ -9,8 +23,15 @@ export async function willContinue(
   // we aren't doing another round after a tool result
   // TODO: whether to handle continuing after too much context used..
   if (step.finishReason !== "tool-calls") return false;
+  // Count both successful results and errors as completed outputs.
+  // In AI SDK v6, failed tool calls produce tool-error content parts
+  // instead of tool-result, so only checking toolResults misses them.
+  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 > step.toolResults.length) return false;
+  if (step.toolCalls.length > completedOutputs) return false;
   if (Array.isArray(stopWhen)) {
     return (await Promise.all(stopWhen.map(async (s) => s({ steps })))).every(
       (stop) => !stop,

package/src/component/_generated/dataModel.ts

@@ -38,7 +38,7 @@ export type Doc<TableName extends TableNames> = DocumentByName<
  * Convex documents are uniquely identified by their `Id`, which is accessible
  * on the `_id` field. To learn more, see [Document IDs](https://docs.convex.dev/using/document-ids).
  *
- * Documents can be loaded using `db.get(id)` in query and mutation functions.
+ * Documents can be loaded using `db.get(tableName, id)` in query and mutation functions.
  *
  * IDs are just strings at runtime, but this type can be used to distinguish them from other
  * strings when type checking.

package/src/component/_generated/server.ts

@@ -107,11 +107,6 @@ export const internalAction: ActionBuilder<DataModel, "internal"> =
  */
 export const httpAction: HttpActionBuilder = httpActionGeneric;
 
-type GenericCtx =
-  | GenericActionCtx<DataModel>
-  | GenericMutationCtx<DataModel>
-  | GenericQueryCtx<DataModel>;
-
 /**
  * A set of services for use within Convex query functions.
  *

package/src/mapping.ts

@@ -297,21 +297,49 @@ export function toModelMessageWarnings(
   }) as any;
 }
 
+/**
+ * Serialize explicitly provided response messages for a step.
+ * Used by the streaming/generation loop where the caller tracks which
+ * messages are new via slicing.
+ */
+export async function serializeResponseMessages<TOOLS extends ToolSet>(
+  ctx: ActionCtx,
+  component: AgentComponent,
+  step: StepResult<TOOLS>,
+  model: ModelOrMetadata | undefined,
+  responseMessages: ModelMessage[],
+): Promise<{ messages: MessageWithMetadata[] }> {
+  return serializeStepMessages(ctx, component, step, model, responseMessages);
+}
+
+/**
+ * Serialize the new messages from a step using a heuristic to determine
+ * which response messages are new (last 1-2 messages).
+ */
 export async function serializeNewMessagesInStep<TOOLS extends ToolSet>(
   ctx: ActionCtx,
   component: AgentComponent,
   step: StepResult<TOOLS>,
   model: ModelOrMetadata | undefined,
-  /**
-   * If provided, these are the new response messages for this step
-   * (pre-sliced by the caller). When not provided, falls back to the
-   * existing heuristic of slicing the last 1-2 messages.
-   *
-   * This is needed for tool approval flows where the SDK adds extra
-   * messages (e.g. approval tool-results) at the beginning of
-   * responseMessages that the old slice(-1/-2) logic would miss.
-   */
-  newResponseMessages?: ModelMessage[],
+): Promise<{ messages: MessageWithMetadata[] }> {
+  const hasToolMessage = step.response.messages.at(-1)?.role === "tool";
+  let messagesToSerialize: ModelMessage[];
+  if (hasToolMessage) {
+    messagesToSerialize = step.response.messages.slice(-2);
+  } else if (step.content.length) {
+    messagesToSerialize = step.response.messages.slice(-1);
+  } else {
+    messagesToSerialize = [{ role: "assistant" as const, content: [] }];
+  }
+  return serializeStepMessages(ctx, component, step, model, messagesToSerialize);
+}
+
+async function serializeStepMessages<TOOLS extends ToolSet>(
+  ctx: ActionCtx,
+  component: AgentComponent,
+  step: StepResult<TOOLS>,
+  model: ModelOrMetadata | undefined,
+  messagesToSerialize: ModelMessage[],
 ): Promise<{ messages: MessageWithMetadata[] }> {
   // If there are tool results, there's another message with the tool results
   // ref: https://github.com/vercel/ai/blob/main/packages/ai/src/generate-text/to-response-messages.ts#L120
@@ -330,18 +358,6 @@ export async function serializeNewMessagesInStep<TOOLS extends ToolSet>(
   } satisfies Omit<MessageWithMetadata, "message" | "text" | "fileIds">;
   const toolFields = { sources: step.sources };
 
-  // Determine which messages to serialize for this step
-  let messagesToSerialize: ModelMessage[];
-  if (newResponseMessages) {
-    messagesToSerialize = newResponseMessages;
-  } else if (hasToolMessage) {
-    messagesToSerialize = step.response.messages.slice(-2);
-  } else if (step.content.length) {
-    messagesToSerialize = step.response.messages.slice(-1);
-  } else {
-    messagesToSerialize = [{ role: "assistant" as const, content: [] }];
-  }
-
   const messages: MessageWithMetadata[] = await Promise.all(
     messagesToSerialize.map(async (msg): Promise<MessageWithMetadata> => {
       const { message, fileIds } = await serializeMessage(ctx, component, msg);

package/src/test.ts

@@ -9,10 +9,9 @@ const modules = import.meta.glob("./component/**/*.ts");
  * @param t - The test convex instance, e.g. from calling `convexTest`.
  * @param name - The name of the component, as registered in convex.config.ts.
  */
-export function register(
-  t: TestConvex<SchemaDefinition<GenericSchema, boolean>>,
-  name: string = "agent",
-) {
+export function register<
+  Schema extends SchemaDefinition<GenericSchema, boolean>,
+>(t: TestConvex<Schema>, name: string = "agent") {
   t.registerComponent(name, schema, modules);
 }
 export default { register, schema, modules };