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

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,175 @@ 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, existingResponseMessage } =
+      await this.findApprovalContext(ctx, {
+        threadId: args.threadId,
+        approvalId: args.approvalId,
+      });
+
+    const newPart = {
+      type: "tool-approval-response" as const,
+      approvalId: args.approvalId,
+      approved: args.approved,
+      reason: args.reason,
+    };
+
+    // Merge into an existing approval-response message for this step
+    // so the AI SDK sees a single tool message per step.
+    if (existingResponseMessage) {
+      const existingContent = existingResponseMessage.message?.content;
+      const mergedContent = Array.isArray(existingContent)
+        ? [...(existingContent as any[]), newPart]
+        : [newPart];
+      await this.updateMessage(ctx, {
+        messageId: existingResponseMessage._id,
+        patch: {
+          message: { role: "tool", content: mergedContent },
+          status: "success",
+        },
+      });
+      return { messageId: existingResponseMessage._id };
+    }
+
+    const { messageId } = await this.saveMessage(ctx, {
+      threadId: args.threadId,
+      promptMessageId,
+      skipEmbeddings: true,
+      message: {
+        role: "tool",
+        content: [newPart],
+      },
+    });
+    return { messageId };
+  }
+
+  private async findApprovalContext(
+    ctx: MutationCtx,
+    args: { threadId: string; approvalId: string },
+  ): Promise<{
+    promptMessageId: string;
+    existingResponseMessage: MessageDoc | undefined;
+  }> {
+    // 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 existingResponseMessage: MessageDoc | undefined;
+    // Limit the search to the most recent messages. Approvals should always
+    // be near the end of the thread.
+    const page = await this.listMessages(ctx, {
+      threadId: args.threadId,
+      paginationOpts: { cursor: null, numItems: 100 },
+    });
+    {
+      for (const message of page.page) {
+        const content = message.message?.content;
+        if (!Array.isArray(content)) continue;
+        // Check if this assistant message starts a different approval step.
+        // If so, any response message we've seen so far belongs to a newer
+        // step — reset it so we don't merge across step boundaries.
+        // Only reset if the target approval is NOT in this message (i.e.,
+        // this is a genuinely different step, not the same step with
+        // multiple tool calls).
+        if (
+          message.message?.role === "assistant" &&
+          content.some(
+            (p: any) =>
+              p.type === "tool-approval-request" &&
+              p.approvalId !== args.approvalId,
+          ) &&
+          !content.some(
+            (p: any) =>
+              p.type === "tool-approval-request" &&
+              p.approvalId === args.approvalId,
+          )
+        ) {
+          existingResponseMessage = undefined;
+        }
+        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`);
+          }
+          // Track the most recent tool-approval-response message for merging
+          if (
+            typedPart.type === "tool-approval-response" &&
+            !existingResponseMessage
+          ) {
+            existingResponseMessage = message;
+          }
+          if (
+            typedPart.type === "tool-approval-request" &&
+            typedPart.approvalId === args.approvalId
+          ) {
+            return { promptMessageId: message._id, existingResponseMessage };
+          }
+        }
+      }
+    }
+
+    throw new Error(
+      `Approval request ${args.approvalId} was not found in the last 100 messages of 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.test.ts

@@ -292,7 +292,7 @@ describe("search.ts", () => {
       expect(result[1]._id).toBe("2");
     });
 
-    it("should filter out tool calls with approval request but NO approval response", () => {
+    it("should keep tool calls with approval request but NO approval response (auto-deny handles them)", () => {
       const messages: MessageDoc[] = [
         {
           _id: "1",
@@ -321,19 +321,18 @@ describe("search.ts", () => {
 
       const result = filterOutOrphanedToolMessages(messages);
       expect(result).toHaveLength(1);
-      // The assistant message should have the tool-call filtered out
+      // The assistant message should keep the tool-call (auto-deny resolves it downstream)
       const assistantContent = result[0].message?.content;
       expect(Array.isArray(assistantContent)).toBe(true);
       if (Array.isArray(assistantContent)) {
-        // Text and approval-request should remain, but tool-call should be filtered
-        expect(assistantContent).toHaveLength(2);
+        expect(assistantContent).toHaveLength(3);
         expect(assistantContent.find((p) => p.type === "text")).toBeDefined();
         expect(
           assistantContent.find((p) => p.type === "tool-approval-request"),
         ).toBeDefined();
         expect(
           assistantContent.find((p) => p.type === "tool-call"),
-        ).toBeUndefined();
+        ).toBeDefined();
       }
     });
 

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 {
+  autoDenyUnresolvedApprovals,
+  docsToModelMessages,
+  toModelMessage,
+} from "../mapping.js";
 
 const DEFAULT_VECTOR_SCORE_THRESHOLD = 0.0;
 // 10k characters should be more than enough for most cases, and stays under
@@ -285,6 +289,12 @@ export function filterOutOrphanedToolMessages(docs: MessageDoc[]) {
     return approvalId !== undefined && approvalResponseIds.has(approvalId);
   };
 
+  // Helper: check if tool call has a pending approval request
+  // (auto-deny handles these downstream, so they must survive the filter)
+  const hasApprovalRequest = (toolCallId: string) => {
+    return approvalRequestsByToolCallId.has(toolCallId);
+  };
+
   for (const doc of docs) {
     if (
       doc.message?.role === "assistant" &&
@@ -294,7 +304,8 @@ export function filterOutOrphanedToolMessages(docs: MessageDoc[]) {
         (p) =>
           p.type !== "tool-call" ||
           toolResultIds.has(p.toolCallId) ||
-          hasApprovalResponse(p.toolCallId),
+          hasApprovalResponse(p.toolCallId) ||
+          hasApprovalRequest(p.toolCallId),
       );
       if (content.length) {
         result.push({
@@ -657,6 +668,11 @@ export async function fetchContextWithPrompt(
       })
     : allMessages;
 
+  // Post-process: auto-deny unresolved approvals so the AI SDK sees a
+  // complete history. Applied after contextHandler so custom handlers
+  // don't need to handle this.
+  processedMessages = autoDenyUnresolvedApprovals(processedMessages);
+
   // Process messages to inline localhost files (if not, file urls pointing to localhost will be sent to LLM providers)
   if (process.env.CONVEX_CLOUD_URL?.startsWith("http://127.0.0.1")) {
     processedMessages = await inlineMessagesFiles(processedMessages);

package/src/client/start.ts

@@ -10,7 +10,7 @@ import {
   type ToolSet,
 } from "ai";
 import {
-  serializeNewMessagesInStep,
+  serializeResponseMessages,
   serializeObjectResult,
 } from "../mapping.js";
 import { embedMessages, fetchContextWithPrompt } from "./search.js";
@@ -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 serializeResponseMessages(
+            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,31 @@ 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));
+      // 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;
+    }
+  }
+
+  // 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,