@convex-dev/agent 0.6.0-alpha.0 → 0.6.0-alpha.1

package/src/client/index.ts

@@ -19,7 +19,7 @@ import type {
 } from "ai";
 import { generateObject, generateText, stepCountIs, streamObject } from "ai";
 
-const MIGRATION_URL = "https://github.com/get-convex/agent/blob/main/MIGRATION.md";
+const MIGRATION_URL = "node_modules/@convex-dev/agent/MIGRATION.md";
 const warnedDeprecations = new Set<string>();
 function warnDeprecation(key: string, message: string) {
   if (!warnedDeprecations.has(key)) {
@@ -1016,6 +1016,125 @@ export class Agent<
     );
   }
 
+  /**
+   * Approve a tool call that requires human approval.
+   * Saves a `tool-approval-response` message to the thread.
+   * After calling this, call `agent.streamText` or `agent.generateText`
+   * with `promptMessageId` set to the returned `messageId` to continue
+   * generation — the AI SDK will automatically execute the approved tool.
+   *
+   * The approval response is attached to the same generation order as the
+   * original approval request, preserving tool_call/tool_result adjacency in
+   * the continuation context even if newer thread messages exist.
+   *
+   * @param ctx A ctx object from a mutation.
+   * @param args.threadId The thread containing the tool call.
+   * @param args.approvalId The approval ID from the tool-approval-request part.
+   * @param args.reason Optional reason for approval.
+   * @returns The messageId of the saved approval response message.
+   */
+  async approveToolCall(
+    ctx: MutationCtx,
+    args: { threadId: string; approvalId: string; reason?: string },
+  ): Promise<{ messageId: string }> {
+    return this.respondToToolCallApproval(ctx, { ...args, approved: true });
+  }
+
+  /**
+   * Deny a tool call that requires human approval.
+   * Saves a `tool-approval-response` message to the thread.
+   * After calling this, call `agent.streamText` or `agent.generateText`
+   * with `promptMessageId` set to the returned `messageId` to continue
+   * generation — the AI SDK will automatically create an `execution-denied`
+   * result and let the model respond accordingly.
+   *
+   * @param ctx A ctx object from a mutation.
+   * @param args.threadId The thread containing the tool call.
+   * @param args.approvalId The approval ID from the tool-approval-request part.
+   * @param args.reason Optional reason for denial.
+   * @returns The messageId of the saved denial response message.
+   */
+  async denyToolCall(
+    ctx: MutationCtx,
+    args: { threadId: string; approvalId: string; reason?: string },
+  ): Promise<{ messageId: string }> {
+    return this.respondToToolCallApproval(ctx, { ...args, approved: false });
+  }
+
+  private async respondToToolCallApproval(
+    ctx: MutationCtx,
+    args: {
+      threadId: string;
+      approvalId: string;
+      approved: boolean;
+      reason?: string;
+    },
+  ): Promise<{ messageId: string }> {
+    const promptMessageId = await this.getApprovalRequestMessageId(ctx, {
+      threadId: args.threadId,
+      approvalId: args.approvalId,
+    });
+
+    const { messageId } = await this.saveMessage(ctx, {
+      threadId: args.threadId,
+      promptMessageId,
+      skipEmbeddings: true,
+      message: {
+        role: "tool",
+        content: [
+          {
+            type: "tool-approval-response",
+            approvalId: args.approvalId,
+            approved: args.approved,
+            reason: args.reason,
+          },
+        ],
+      },
+    });
+    return { messageId };
+  }
+
+  private async getApprovalRequestMessageId(
+    ctx: MutationCtx,
+    args: { threadId: string; approvalId: string },
+  ): Promise<string> {
+    // NOTE: This pagination returns messages in descending order (newest first).
+    // The "already handled" check (tool-approval-response) relies on seeing
+    // responses before their corresponding requests. If the pagination order
+    // changes, this logic will need to be updated.
+    let cursor: string | null = null;
+    do {
+      const page = await this.listMessages(ctx, {
+        threadId: args.threadId,
+        paginationOpts: { cursor, numItems: 100 },
+      });
+      for (const message of page.page) {
+        const content = message.message?.content;
+        if (!Array.isArray(content)) continue;
+        for (const part of content) {
+          const typedPart = part as { type?: unknown; approvalId?: unknown };
+          if (
+            typedPart.type === "tool-approval-response" &&
+            typedPart.approvalId === args.approvalId
+          ) {
+            throw new Error(`Approval ${args.approvalId} was already handled`);
+          }
+          if (
+            typedPart.type === "tool-approval-request" &&
+            typedPart.approvalId === args.approvalId
+          ) {
+            return message._id;
+          }
+        }
+      }
+      cursor = page.isDone ? null : page.continueCursor;
+    } while (cursor !== null);
+
+    throw new Error(
+      `Approval request ${args.approvalId} was not found in thread ${args.threadId}`,
+    );
+  }
+
   /**
    * Explicitly save a "step" created by the AI SDK.
    * @param ctx The ctx argument to a mutation or action.

package/src/client/mockModel.ts

@@ -16,8 +16,15 @@ const DEFAULT_USAGE = {
   outputTokens: 10,
   inputTokens: 3,
   totalTokens: 13,
-  inputTokenDetails: undefined,
-  outputTokenDetails: undefined,
+  inputTokenDetails: {
+    noCacheTokens: 3,
+    cacheReadTokens: 0,
+    cacheWriteTokens: 0,
+  },
+  outputTokenDetails: {
+    textTokens: 10,
+    reasoningTokens: 0,
+  },
 };
 
 export type MockModelArgs = {

package/src/client/search.ts

@@ -29,7 +29,11 @@ import type {
   ActionCtx,
 } from "./types.js";
 import { inlineMessagesFiles } from "./files.js";
-import { docsToModelMessages, toModelMessage } from "../mapping.js";
+import {
+  docsToModelMessages,
+  mergeApprovalResponseMessages,
+  toModelMessage,
+} from "../mapping.js";
 
 const DEFAULT_VECTOR_SCORE_THRESHOLD = 0.0;
 // 10k characters should be more than enough for most cases, and stays under
@@ -637,13 +641,13 @@ export async function fetchContextWithPrompt(
   const inputPrompt = promptArray.map(toModelMessage);
   const existingResponses = docsToModelMessages(existingResponseDocs);
 
-  const allMessages = [
+  const allMessages = mergeApprovalResponseMessages([
     ...search,
     ...recent,
     ...inputMessages,
     ...inputPrompt,
     ...existingResponses,
-  ];
+  ]);
   let processedMessages = args.contextHandler
     ? await args.contextHandler(ctx, {
         allMessages,

package/src/client/start.ts

@@ -112,6 +112,7 @@ export async function startGeneration<
       | { step: StepResult<TOOLS> }
       | { object: GenerateObjectResult<unknown> },
     createPendingMessage?: boolean,
+    finishStreamId?: string,
   ) => Promise<void>;
   fail: (reason: string) => Promise<void>;
   getSavedMessages: () => MessageDoc[];
@@ -200,16 +201,18 @@ export async function startGeneration<
     tools?: Tools;
     _internal?: { generateId?: IdGenerator };
   } & CallSettings;
-  if (pendingMessageId) {
-    if (!aiArgs._internal?.generateId) {
-      aiArgs._internal = {
-        ...aiArgs._internal,
-        generateId: pendingMessageId
-          ? () => pendingMessageId ?? crypto.randomUUID()
-          : undefined,
-      };
-    }
-  }
+  // NOTE: We intentionally do NOT override _internal.generateId here.
+  // The AI SDK uses generateId() for many internal IDs (approval IDs,
+  // tool execution IDs, message IDs, etc.) and they must be unique.
+  // The pending message is linked via the explicit `pendingMessageId`
+  // parameter passed to addMessages in the save closure.
+  // Track how many response messages we've already saved across steps.
+  // step.response.messages is cumulative — each step appends to it.
+  // We need to know which messages are new in each step to serialize
+  // only the new ones (important for tool approval flows where the SDK
+  // may add extra messages like approval tool-results).
+  let previousResponseMessageCount = 0;
+
   return {
     args: aiArgs,
     order: order ?? 0,
@@ -228,22 +231,35 @@ export async function startGeneration<
         | { step: StepResult<TOOLS> }
         | { object: GenerateObjectResult<unknown> },
       createPendingMessage?: boolean,
+      /**
+       * If provided, finish this stream atomically with the message save.
+       * This prevents UI flickering from separate mutations (issue #181).
+       */
+      finishStreamId?: string,
     ) => {
       if (threadId && saveMessages !== "none") {
-        const serialized =
-          "object" in toSave
-            ? await serializeObjectResult(
-                ctx,
-                component,
-                toSave.object,
-                activeModel,
-              )
-            : await serializeNewMessagesInStep(
-                ctx,
-                component,
-                toSave.step,
-                activeModel,
-              );
+        let serialized;
+        if ("object" in toSave) {
+          serialized = await serializeObjectResult(
+            ctx,
+            component,
+            toSave.object,
+            activeModel,
+          );
+        } else {
+          const allResponseMessages = toSave.step.response.messages;
+          const newResponseMessages = allResponseMessages.slice(
+            previousResponseMessageCount,
+          );
+          previousResponseMessageCount = allResponseMessages.length;
+          serialized = await serializeNewMessagesInStep(
+            ctx,
+            component,
+            toSave.step,
+            activeModel,
+            newResponseMessages,
+          );
+        }
         const embeddings = await embedMessages(
           ctx,
           { threadId, ...opts, userId },
@@ -265,6 +281,7 @@ export async function startGeneration<
           messages: serialized.messages,
           embeddings,
           failPendingSteps: false,
+          finishStreamId,
         });
         const lastMessage = saved.messages.at(-1)!;
         if (createPendingMessage) {

package/src/client/streamText.ts

@@ -80,6 +80,9 @@ export async function streamText<
 
   const steps: StepResult<TOOLS>[] = [];
 
+  // Track the final step for atomic save with stream finish (issue #181)
+  let pendingFinalStep: StepResult<TOOLS> | undefined;
+
   const streamer =
     threadId && options.saveStreamDeltas
       ? new DeltaStreamer(
@@ -138,7 +141,14 @@ export async function streamText<
     onStepFinish: async (step) => {
       steps.push(step);
       const createPendingMessage = await willContinue(steps, args.stopWhen);
-      await call.save({ step }, createPendingMessage);
+      if (!createPendingMessage && streamer) {
+        // This is the final step with streaming enabled.
+        // Defer saving until stream consumption completes for atomic finish (issue #181).
+        streamer.markFinishedExternally();
+        pendingFinalStep = step;
+      } else {
+        await call.save({ step }, createPendingMessage);
+      }
       return args.onStepFinish?.(step);
     },
   }) as StreamTextResult<TOOLS, OUTPUT>;
@@ -150,8 +160,22 @@ export async function streamText<
       !options.saveStreamDeltas.returnImmediately) ||
     options?.saveStreamDeltas === true
   ) {
-    await stream;
-    await result.consumeStream();
+    try {
+      await stream;
+      await result.consumeStream();
+    } catch (e) {
+      // If the stream errored (e.g. onStepFinish threw), the DeltaStreamer's
+      // 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));
+      throw e;
+    }
+  }
+
+  // If we deferred the final step save, do it now with atomic stream finish.
+  if (pendingFinalStep && streamer) {
+    const finishStreamId = await streamer.getOrCreateStreamId();
+    await call.save({ step: pendingFinalStep }, false, finishStreamId);
   }
   const metadata: GenerationOutputMetadata = {
     promptMessageId,

package/src/client/streaming.ts

@@ -210,6 +210,9 @@ export class DeltaStreamer<T> {
   #ongoingWrite: 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;
 
   constructor(
     public readonly component: AgentComponent,
@@ -259,17 +262,16 @@ export class DeltaStreamer<T> {
   // Avoid race conditions by only creating once
   #creatingStreamIdPromise: Promise<string> | undefined;
   public async getStreamId() {
-    if (this.streamId) {
-      return this.streamId;
-    }
-    if (this.#creatingStreamIdPromise) {
-      return this.#creatingStreamIdPromise;
+    if (!this.streamId) {
+      if (!this.#creatingStreamIdPromise) {
+        this.#creatingStreamIdPromise = this.ctx.runMutation(
+          this.component.streams.create,
+          this.metadata,
+        );
+      }
+      this.streamId = await this.#creatingStreamIdPromise;
     }
-    this.#creatingStreamIdPromise = this.ctx.runMutation(
-      this.component.streams.create,
-      this.metadata,
-    );
-    this.streamId = await this.#creatingStreamIdPromise;
+    return this.streamId;
   }
 
   public async addParts(parts: T[]) {
@@ -290,7 +292,27 @@ export class DeltaStreamer<T> {
     for await (const chunk of stream) {
       await this.addParts([chunk]);
     }
-    await this.finish();
+    // Skip finish if it will be handled externally (atomically with message save)
+    if (!this.#finishedExternally) {
+      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.
+   */
+  public markFinishedExternally(): void {
+    this.#finishedExternally = 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();
   }
 
   async #sendDelta() {

package/src/client/types.ts

@@ -46,6 +46,14 @@ import type {
 import type { StreamingOptions } from "./streaming.js";
 import type { ComponentApi } from "../component/_generated/component.js";
 
+/**
+ * Type-level check that ensures models are from AI SDK v6.
+ * If a v5 model (LanguageModelV2) is passed, TypeScript will show the error message string.
+ */
+type AssertAISDKv6<T> = T extends { specificationVersion: "v3" }
+  ? T
+  : "⚠️ @convex-dev/agent v0.6.0 requires AI SDK v6. Update your dependencies: npm install ai@^6.0.35 @ai-sdk/openai@^3.0.10 (or other provider). See: node_modules/@convex-dev/agent/MIGRATION.md";
+
 export type AgentPrompt = {
   /**
    * System message to include in the prompt. Overwrites Agent instructions.
@@ -91,23 +99,17 @@ export type AgentPrompt = {
 export type Config = {
   /**
    * The LLM model to use for generating / streaming text and objects.
-   * e.g.
+   * Requires AI SDK v6 (@ai-sdk/* packages v3.x).
+   *
+   * @example
    * import { openai } from "@ai-sdk/openai"
    * const myAgent = new Agent(components.agent, {
    *   languageModel: openai.chat("gpt-4o-mini"),
+   * })
    */
-  languageModel?: LanguageModel;
+  languageModel?: AssertAISDKv6<LanguageModel>;
   /**
-   * The model to use for text embeddings. Optional.
-   * If specified, it will use this for generating vector embeddings
-   * of chats, and can opt-in to doing vector search for automatic context
-   * on generateText, etc.
-   * e.g.
-   * import { openai } from "@ai-sdk/openai"
-   * const myAgent = new Agent(components.agent, {
-   *   ...
-   *   textEmbeddingModel: openai.embedding("text-embedding-3-small")
-   * @deprecated — Use embeddingModel instead.
+   * @deprecated Use `embeddingModel` instead.
    */
   textEmbeddingModel?: EmbeddingModel;
   /**

package/src/component/_generated/component.ts

@@ -154,6 +154,7 @@ export type ComponentApi<Name extends string | undefined = string | undefined> =
             vectors: Array<Array<number> | null>;
           };
           failPendingSteps?: boolean;
+          finishStreamId?: string;
           hideFromUserIdSearch?: boolean;
           messages: Array<{
             error?: string;
@@ -1420,22 +1421,19 @@ export type ComponentApi<Name extends string | undefined = string | undefined> =
                           toolName: string;
                           type: "tool-call";
                         }
-                        | {
-                            args: any;
-                            input?: any;
-                            providerExecuted?: boolean;
-                            providerMetadata?: Record<
-                              string,
-                              Record<string, any>
-                            >;
-                            providerOptions?: Record<
-                              string,
-                              Record<string, any>
-                            >;
-                            toolCallId: string;
-                            toolName: string;
-                            type: "tool-call";
-                          }
+                      | {
+                          args: any;
+                          input?: any;
+                          providerExecuted?: boolean;
+                          providerMetadata?: Record<
+                            string,
+                            Record<string, any>
+                          >;
+                          providerOptions?: Record<string, Record<string, any>>;
+                          toolCallId: string;
+                          toolName: string;
+                          type: "tool-call";
+                        }
                       | {
                           args?: any;
                           experimental_content?: Array<
@@ -2533,22 +2531,19 @@ export type ComponentApi<Name extends string | undefined = string | undefined> =
                           toolName: string;
                           type: "tool-call";
                         }
-                        | {
-                            args: any;
-                            input?: any;
-                            providerExecuted?: boolean;
-                            providerMetadata?: Record<
-                              string,
-                              Record<string, any>
-                            >;
-                            providerOptions?: Record<
-                              string,
-                              Record<string, any>
-                            >;
-                            toolCallId: string;
-                            toolName: string;
-                            type: "tool-call";
-                          }
+                      | {
+                          args: any;
+                          input?: any;
+                          providerExecuted?: boolean;
+                          providerMetadata?: Record<
+                            string,
+                            Record<string, any>
+                          >;
+                          providerOptions?: Record<string, Record<string, any>>;
+                          toolCallId: string;
+                          toolName: string;
+                          type: "tool-call";
+                        }
                       | {
                           args?: any;
                           experimental_content?: Array<
@@ -3054,22 +3049,19 @@ export type ComponentApi<Name extends string | undefined = string | undefined> =
                           toolName: string;
                           type: "tool-call";
                         }
-                        | {
-                            args: any;
-                            input?: any;
-                            providerExecuted?: boolean;
-                            providerMetadata?: Record<
-                              string,
-                              Record<string, any>
-                            >;
-                            providerOptions?: Record<
-                              string,
-                              Record<string, any>
-                            >;
-                            toolCallId: string;
-                            toolName: string;
-                            type: "tool-call";
-                          }
+                      | {
+                          args: any;
+                          input?: any;
+                          providerExecuted?: boolean;
+                          providerMetadata?: Record<
+                            string,
+                            Record<string, any>
+                          >;
+                          providerOptions?: Record<string, Record<string, any>>;
+                          toolCallId: string;
+                          toolName: string;
+                          type: "tool-call";
+                        }
                       | {
                           args?: any;
                           experimental_content?: Array<
@@ -4063,22 +4055,19 @@ export type ComponentApi<Name extends string | undefined = string | undefined> =
                           toolName: string;
                           type: "tool-call";
                         }
-                        | {
-                            args: any;
-                            input?: any;
-                            providerExecuted?: boolean;
-                            providerMetadata?: Record<
-                              string,
-                              Record<string, any>
-                            >;
-                            providerOptions?: Record<
-                              string,
-                              Record<string, any>
-                            >;
-                            toolCallId: string;
-                            toolName: string;
-                            type: "tool-call";
-                          }
+                      | {
+                          args: any;
+                          input?: any;
+                          providerExecuted?: boolean;
+                          providerMetadata?: Record<
+                            string,
+                            Record<string, any>
+                          >;
+                          providerOptions?: Record<string, Record<string, any>>;
+                          toolCallId: string;
+                          toolName: string;
+                          type: "tool-call";
+                        }
                       | {
                           args?: any;
                           experimental_content?: Array<

package/src/component/messages.ts

@@ -39,7 +39,7 @@ import {
   vVectorId,
 } from "./vector/tables.js";
 import { changeRefcount } from "./files.js";
-import { getStreamingMessagesWithMetadata } from "./streams.js";
+import { getStreamingMessagesWithMetadata, finishHandler } from "./streams.js";
 import { partial } from "convex-helpers/validators";
 
 function publicMessage(message: Doc<"messages">): MessageDoc {
@@ -141,6 +141,9 @@ const addMessagesArgs = {
   // if set to true, these messages will not show up in text or vector search
   // results for the userId
   hideFromUserIdSearch: v.optional(v.boolean()),
+  // If provided, finish this stream atomically with the message save.
+  // This prevents UI flickering from separate mutations (issue #181).
+  finishStreamId: v.optional(v.id("streamingMessages")),
 };
 export const addMessages = mutation({
   args: addMessagesArgs,
@@ -161,6 +164,8 @@ async function addMessagesHandler(
   const {
     embeddings,
     failPendingSteps,
+    // Destructured separately to exclude from `...rest` (used in addMessages args, not message fields)
+    finishStreamId,
     messages,
     promptMessageId,
     pendingMessageId,
@@ -274,7 +279,7 @@ async function addMessagesHandler(
         order: pendingMessage.order,
         stepOrder: pendingMessage.stepOrder,
       });
-      toReturn.push(pendingMessage);
+      toReturn.push((await ctx.db.get(pendingMessage._id))!);
       continue;
     }
     if (message.message.role === "user") {
@@ -303,6 +308,11 @@ async function addMessagesHandler(
     // TODO: delete the associated stream data for the order/stepOrder
     toReturn.push((await ctx.db.get(messageId))!);
   }
+  // Atomically finish the stream if requested, preventing UI flickering
+  // from separate mutations for message save and stream finish (issue #181).
+  if (finishStreamId) {
+    await finishHandler(ctx, { streamId: finishStreamId });
+  }
   return { messages: toReturn.map(publicMessage) };
 }
 

package/src/deltas.ts

@@ -124,7 +124,6 @@ export async function deriveUIMessagesFromDeltas(
         blankUIMessage(streamMessage, threadId),
         parts,
       );
-      // TODO: this fails on partial tool calls
       messages.push(uiMessage);
     } else {
       const [uiMessages] = deriveUIMessagesFromTextStreamParts(