@convex-dev/agent 0.6.2 → 0.6.4

package/src/client/definePlaygroundAPI.ts

@@ -6,6 +6,7 @@ import {
   type ApiFromModules,
   type GenericActionCtx,
   type GenericDataModel,
+  type GenericMutationCtx,
   type GenericQueryCtx,
 } from "convex/server";
 import { v } from "convex/values";
@@ -64,7 +65,10 @@ export function definePlaygroundAPI<DataModel extends GenericDataModel>(
     }
   }
 
-  async function validateApiKey(ctx: RunQueryCtx, apiKey: string) {
+  async function validateApiKey(
+    ctx: QueryCtx | MutationCtx | ActionCtx,
+    apiKey: string,
+  ) {
     await ctx.runQuery(component.apiKeys.validate, { apiKey });
   }
 
@@ -276,18 +280,22 @@ export function definePlaygroundAPI<DataModel extends GenericDataModel>(
         },
         { contextOptions, storageOptions, saveStreamDeltas: true },
       );
-      const outputMessages = await Promise.all(
-        (await steps).map(async (step) => {
-          const { messages } = await serializeNewMessagesInStep(
-            ctx,
-            component,
-            step,
-            {
-              model: getModelName(agent.options.languageModel),
-              provider: getProviderName(agent.options.languageModel),
-            },
-          );
-          return messages.map((messageWithMetadata, i) => {
+      const outputMessages: MessageDoc[][] = [];
+      let previousResponseMessageCount = 0;
+      for (const step of await steps) {
+        const { messages } = await serializeNewMessagesInStep(
+          ctx,
+          component,
+          step,
+          {
+            model: getModelName(agent.options.languageModel),
+            provider: getProviderName(agent.options.languageModel),
+          },
+          previousResponseMessageCount,
+        );
+        previousResponseMessageCount = step.response.messages.length;
+        outputMessages.push(
+          messages.map((messageWithMetadata, i) => {
             return {
               ...messageWithMetadata,
               tool: isTool(messageWithMetadata.message),
@@ -300,9 +308,9 @@ export function definePlaygroundAPI<DataModel extends GenericDataModel>(
               order: 0,
               stepOrder: i + 1,
             } satisfies MessageDoc;
-          });
-        }),
-      );
+          }),
+        );
+      }
       return { text: await text, messages: outputMessages.flat() };
     },
     returns: v.object({ text: v.string(), messages: v.array(vMessageDoc) }),
@@ -367,4 +375,12 @@ export function definePlaygroundAPI<DataModel extends GenericDataModel>(
   };
 }
 
-type RunQueryCtx = { runQuery: GenericQueryCtx<GenericDataModel>["runQuery"] };
+type QueryCtx = Pick<GenericQueryCtx<GenericDataModel>, "runQuery">;
+type MutationCtx = Pick<
+  GenericMutationCtx<GenericDataModel>,
+  "runQuery" | "runMutation"
+>;
+type ActionCtx = Pick<
+  GenericActionCtx<GenericDataModel>,
+  "runQuery" | "runMutation" | "runAction"
+>;

package/src/client/index.test.ts

@@ -2,6 +2,7 @@ import { describe, expect, test } from "vitest";
 import {
   Agent,
   createThread,
+  createTool,
   filterOutOrphanedToolMessages,
   type MessageDoc,
 } from "./index.js";
@@ -62,6 +63,74 @@ export const createThreadManually = mutation({
   },
 });
 
+const saveStepAgent = new Agent(components.agent, {
+  name: "save-step-test",
+  instructions: "test",
+  tools: {
+    echo: createTool({
+      description: "Echo a value",
+      inputSchema: z.object({ value: z.string() }),
+      execute: async (_ctx, input) => `echo:${input.value}`,
+    }),
+  },
+  languageModel: mockModel({
+    contentSteps: [
+      [
+        {
+          type: "tool-call",
+          toolCallId: "ss-1",
+          toolName: "echo",
+          input: JSON.stringify({ value: "hi" }),
+        },
+      ],
+      [{ type: "text", text: "done" }],
+    ],
+  }),
+  stopWhen: stepCountIs(5),
+});
+
+export const replayStepsViaSaveStep = action({
+  args: { withWatermark: v.boolean() },
+  handler: async (ctx, args) => {
+    const { thread } = await saveStepAgent.createThread(ctx, {
+      userId: "ss-gen",
+    });
+    const genResult = await thread.generateText({ prompt: "echo hi" });
+    const steps = genResult.steps;
+
+    const { threadId } = await saveStepAgent.createThread(ctx, {
+      userId: "ss-replay",
+    });
+    const { messageId: promptMessageId } = await saveStepAgent.saveMessage(ctx, {
+      threadId,
+      message: { role: "user", content: "echo hi" },
+      skipEmbeddings: true,
+    });
+    let previousStep: (typeof steps)[number] | undefined;
+    for (const step of steps) {
+      await saveStepAgent.saveStep(ctx, {
+        threadId,
+        promptMessageId,
+        step,
+        previousStep: args.withWatermark ? previousStep : undefined,
+      });
+      previousStep = step;
+    }
+
+    const replayed = await saveStepAgent.listMessages(ctx, {
+      threadId,
+      paginationOpts: { cursor: null, numItems: 50 },
+      statuses: ["success", "pending", "failed"],
+    });
+    const contentTypes = replayed.page.flatMap((m) =>
+      Array.isArray(m.message?.content)
+        ? m.message!.content.map((c: { type?: string }) => c.type ?? "text")
+        : ["text"],
+    );
+    return { stepCount: steps.length, contentTypes };
+  },
+});
+
 export const createThreadMutation = agent.createThreadMutation();
 export const generateObjectAction = agent.asObjectAction({
   schema: z.object({ hello: z.string().describe("A string for testing") }),
@@ -162,6 +231,7 @@ const testApi: ApiFromModules<{
     generateTextAction: typeof generateTextAction;
     generateObjectAction: typeof generateObjectAction;
     saveMessageMutation: typeof saveMessageMutation;
+    replayStepsViaSaveStep: typeof replayStepsViaSaveStep;
   };
 }>["fns"] = anyApi["index.test"] as any;
 
@@ -177,6 +247,27 @@ describe("Agent thick client", () => {
     expect(result).toBeDefined();
     expect(result).toMatch(TEST_TEXT);
   });
+  test("saveStep with previousStep saves each step's new messages exactly once", async () => {
+    const t = initConvexTest(schema);
+    const res = await t.action(testApi.replayStepsViaSaveStep, {
+      withWatermark: true,
+    });
+    expect(res.stepCount).toBe(2);
+    const toolCalls = res.contentTypes.filter((t) => t === "tool-call").length;
+    const toolResults = res.contentTypes.filter(
+      (t) => t === "tool-result",
+    ).length;
+    expect(toolCalls).toBe(1);
+    expect(toolResults).toBe(1);
+  });
+  test("saveStep without previousStep duplicates prior messages", async () => {
+    const t = initConvexTest(schema);
+    const res = await t.action(testApi.replayStepsViaSaveStep, {
+      withWatermark: false,
+    });
+    const toolCalls = res.contentTypes.filter((t) => t === "tool-call").length;
+    expect(toolCalls).toBeGreaterThan(1);
+  });
 });
 
 describe("filterOutOrphanedToolMessages", () => {

package/src/client/index.ts

@@ -1182,7 +1182,9 @@ export class Agent<
   }
 
   /**
-   * Explicitly save a "step" created by the AI SDK.
+   * Explicitly save a "step" created by the AI SDK. For multi-step generation
+   * loops, pass `previousStep` so we save only the new response messages —
+   * see the arg JSDoc for why.
    * @param ctx The ctx argument to a mutation or action.
    * @param args The Step generated by the AI SDK.
    */
@@ -1199,6 +1201,15 @@ export class Agent<
        * The step to save, possibly including multiple tool calls.
        */
       step: StepResult<TOOLS>;
+      /**
+       * The previous step in the same generation loop, if any. Pass it so we
+       * can compute how many of `step.response.messages` are already saved.
+       * Omit for the first step. AI SDK v6's `step.response.messages` is
+       * cumulative across steps; without this, multi-step callers duplicate
+       * every prior message on every save — the exact failure mode this fix
+       * addresses, just at the public-API layer.
+       */
+      previousStep?: StepResult<TOOLS>;
       /**
        * The model used to generate the step.
        * Defaults to the chat model for the Agent.
@@ -1211,6 +1222,18 @@ export class Agent<
       provider?: string;
     },
   ): Promise<{ messages: MessageDoc[] }> {
+    const previousResponseMessageCount =
+      args.previousStep?.response.messages.length ?? 0;
+    if (
+      args.previousStep !== undefined &&
+      args.step.response.messages.length < previousResponseMessageCount
+    ) {
+      throw new Error(
+        `saveStep: step.response.messages length (${args.step.response.messages.length}) is less than ` +
+          `previousStep.response.messages length (${previousResponseMessageCount}). ` +
+          `Ensure previousStep is from the immediately preceding step in the same generation loop.`,
+      );
+    }
     const { messages } = await serializeNewMessagesInStep(
       ctx,
       this.component,
@@ -1219,6 +1242,7 @@ export class Agent<
         provider: args.provider ?? getProviderName(this.options.languageModel),
         model: args.model ?? getModelName(this.options.languageModel),
       },
+      previousResponseMessageCount,
     );
     const embeddings = await this.generateEmbeddings(
       ctx,

package/src/client/saveInputMessages.ts

@@ -31,7 +31,10 @@ export async function saveInputMessages(
     storageOptions?: {
       saveMessages?: "all" | "promptAndOutput";
     };
-  } & Pick<Config, "usageHandler" | "textEmbeddingModel" | "embeddingModel" | "callSettings">,
+  } & Pick<
+    Config,
+    "usageHandler" | "textEmbeddingModel" | "embeddingModel" | "callSettings"
+  >,
 ): Promise<{
   promptMessageId: string | undefined;
   pendingMessage: MessageDoc;

package/src/client/streaming.integration.test.ts

@@ -9,11 +9,7 @@ import {
   DeltaStreamer,
   mergeTransforms,
 } from "./streaming.js";
-import {
-  getParts,
-  deriveUIMessagesFromDeltas,
-  deriveUIMessagesFromTextStreamParts,
-} from "../deltas.js";
+import { getParts, deriveUIMessagesFromDeltas } from "../deltas.js";
 import type { TestConvex } from "convex-test";
 import type { StreamDelta, StreamMessage } from "../validators.js";
 import { dedupeMessages } from "../react/useUIMessages.js";
@@ -152,9 +148,7 @@ describe("HTTP Streaming Initiation", () => {
 
       // Verify we can reconstruct the text from deltas
       const { parts } = getParts(deltas);
-      const textParts = parts.filter(
-        (p: any) => p.type === "text-delta",
-      );
+      const textParts = parts.filter((p: any) => p.type === "text-delta");
       expect(textParts.length).toBeGreaterThan(0);
     });
   });
@@ -265,10 +259,9 @@ describe("Stream Exclusion Logic", () => {
       await streamer2.addParts([{ type: "start" }]);
 
       // Default list: only streaming
-      const defaultStreams = await ctx.runQuery(
-        components.agent.streams.list,
-        { threadId },
-      );
+      const defaultStreams = await ctx.runQuery(components.agent.streams.list, {
+        threadId,
+      });
       expect(defaultStreams).toHaveLength(1);
       expect(defaultStreams[0].status).toBe("streaming");
       expect(defaultStreams[0].order).toBe(1);
@@ -325,10 +318,10 @@ describe("Stream Exclusion Logic", () => {
       expect(finishedStreams[0].status).toBe("finished");
 
       // Query for only aborted
-      const abortedStreams = await ctx.runQuery(
-        components.agent.streams.list,
-        { threadId, statuses: ["aborted"] },
-      );
+      const abortedStreams = await ctx.runQuery(components.agent.streams.list, {
+        threadId,
+        statuses: ["aborted"],
+      });
       expect(abortedStreams).toHaveLength(1);
       expect(abortedStreams[0].status).toBe("aborted");
 
@@ -487,10 +480,10 @@ describe("Delta Stream Consumption", () => {
       expect(laterParts.length).toBeLessThanOrEqual(allParts.length);
 
       // Fetching from the end cursor should yield nothing
-      const noDeltas = await ctx.runQuery(
-        components.agent.streams.listDeltas,
-        { threadId, cursors: [{ cursor: endCursor, streamId }] },
-      );
+      const noDeltas = await ctx.runQuery(components.agent.streams.listDeltas, {
+        threadId,
+        cursors: [{ cursor: endCursor, streamId }],
+      });
       expect(noDeltas).toHaveLength(0);
     });
   });
@@ -526,16 +519,13 @@ describe("Delta Stream Consumption", () => {
       const id2 = streamer2.streamId!;
 
       // Fetch deltas for both streams simultaneously
-      const deltas = await ctx.runQuery(
-        components.agent.streams.listDeltas,
-        {
-          threadId,
-          cursors: [
-            { cursor: 0, streamId: id1 },
-            { cursor: 0, streamId: id2 },
-          ],
-        },
-      );
+      const deltas = await ctx.runQuery(components.agent.streams.listDeltas, {
+        threadId,
+        cursors: [
+          { cursor: 0, streamId: id1 },
+          { cursor: 0, streamId: id2 },
+        ],
+      });
 
       // Should have deltas for both streams
       const s1Deltas = deltas.filter((d) => d.streamId === id1);
@@ -568,10 +558,10 @@ describe("Delta Stream Consumption", () => {
         threadId,
         statuses: ["finished"],
       });
-      const deltas = await ctx.runQuery(
-        components.agent.streams.listDeltas,
-        { threadId, cursors: [{ cursor: 0, streamId }] },
-      );
+      const deltas = await ctx.runQuery(components.agent.streams.listDeltas, {
+        threadId,
+        cursors: [{ cursor: 0, streamId }],
+      });
 
       // Derive UI messages
       const uiMessages = await deriveUIMessagesFromDeltas(
@@ -615,10 +605,10 @@ describe("Delta Stream Consumption", () => {
       await streamer.consumeStream(result.toUIMessageStream());
       const streamId = streamer.streamId!;
 
-      const deltas = await ctx.runQuery(
-        components.agent.streams.listDeltas,
-        { threadId, cursors: [{ cursor: 0, streamId }] },
-      );
+      const deltas = await ctx.runQuery(components.agent.streams.listDeltas, {
+        threadId,
+        cursors: [{ cursor: 0, streamId }],
+      });
       const { parts } = getParts(deltas);
 
       // Compressed: all text-deltas for one text section should be merged
@@ -677,77 +667,6 @@ describe("Delta Stream Consumption", () => {
     expect((parts[0] as { type: string }).type).toBe("new");
     expect(cursor).toBe(6);
   });
-
-  test("TextStreamPart format delta reconstruction with tool calls", () => {
-    const streamId = "s1";
-    const streamMessage: StreamMessage = {
-      streamId,
-      order: 1,
-      stepOrder: 0,
-      status: "streaming",
-    };
-    const deltas: StreamDelta[] = [
-      {
-        streamId,
-        start: 0,
-        end: 1,
-        parts: [{ type: "text-delta", id: "txt-0", text: "Let me call a tool. " }],
-      },
-      {
-        streamId,
-        start: 1,
-        end: 2,
-        parts: [
-          {
-            type: "tool-call",
-            toolCallId: "tc1",
-            toolName: "search",
-            input: { query: "hello" },
-          },
-        ],
-      },
-      {
-        streamId,
-        start: 2,
-        end: 3,
-        parts: [
-          {
-            type: "tool-result",
-            toolCallId: "tc1",
-            toolName: "search",
-            output: "Found 3 results",
-          },
-        ],
-      },
-      {
-        streamId,
-        start: 3,
-        end: 4,
-        parts: [
-          { type: "text-delta", id: "txt-1", text: "Here are the results." },
-        ],
-      },
-    ];
-
-    const [messages, , changed] = deriveUIMessagesFromTextStreamParts(
-      "thread1",
-      [streamMessage],
-      [],
-      deltas,
-    );
-
-    expect(messages).toHaveLength(1);
-    expect(changed).toBe(true);
-
-    const msg = messages[0];
-    expect(msg.text).toContain("Let me call a tool.");
-    expect(msg.text).toContain("Here are the results.");
-
-    const toolParts = msg.parts.filter((p: any) =>
-      p.type.startsWith("tool-"),
-    );
-    expect(toolParts.length).toBeGreaterThan(0);
-  });
 });
 
 // ============================================================================
@@ -879,18 +798,21 @@ describe("Fallback Behavior", () => {
       order: 0,
       stepOrder: 0,
       status: "streaming",
+      format: "UIMessageChunk",
     };
     const finishedMsg: StreamMessage = {
       streamId: "s2",
       order: 1,
       stepOrder: 0,
       status: "finished",
+      format: "UIMessageChunk",
     };
     const abortedMsg: StreamMessage = {
       streamId: "s3",
       order: 2,
       stepOrder: 0,
       status: "aborted",
+      format: "UIMessageChunk",
     };
 
     const msgs = await deriveUIMessagesFromDeltas(
@@ -1001,10 +923,10 @@ describe("Stream Lifecycle Integration", () => {
       expect(finished).toHaveLength(1);
 
       // 4. Derive UI messages from stored deltas
-      const deltas = await ctx.runQuery(
-        components.agent.streams.listDeltas,
-        { threadId, cursors: [{ cursor: 0, streamId }] },
-      );
+      const deltas = await ctx.runQuery(components.agent.streams.listDeltas, {
+        threadId,
+        cursors: [{ cursor: 0, streamId }],
+      });
       const uiMessages = await deriveUIMessagesFromDeltas(
         threadId,
         finished,
@@ -1056,10 +978,10 @@ describe("Stream Lifecycle Integration", () => {
       expect(aborted[0].status).toBe("aborted");
 
       // Even aborted streams have their deltas stored
-      const deltas = await ctx.runQuery(
-        components.agent.streams.listDeltas,
-        { threadId, cursors: [{ cursor: 0, streamId }] },
-      );
+      const deltas = await ctx.runQuery(components.agent.streams.listDeltas, {
+        threadId,
+        cursors: [{ cursor: 0, streamId }],
+      });
       expect(deltas.length).toBeGreaterThan(0);
     });
   });

package/src/client/types.ts

@@ -32,8 +32,6 @@ export interface Output<_T = any, _P = any, _E = any> {
 import type {
   GenericActionCtx,
   GenericDataModel,
-  GenericMutationCtx,
-  GenericQueryCtx,
   WithoutSystemFields,
 } from "convex/server";
 import type {
@@ -356,10 +354,7 @@ export type TextArgs<
   OUTPUT extends Output<any, any, any> = never,
 > = Omit<
   Parameters<
-    typeof generateText<
-      TOOLS extends undefined ? AgentTools : TOOLS,
-      OUTPUT
-    >
+    typeof generateText<TOOLS extends undefined ? AgentTools : TOOLS, OUTPUT>
   >[0],
   "model" | "prompt" | "messages"
 > & {
@@ -376,10 +371,7 @@ export type StreamingTextArgs<
   OUTPUT extends Output<any, any, any> = never,
 > = Omit<
   Parameters<
-    typeof streamText<
-      TOOLS extends undefined ? AgentTools : TOOLS,
-      OUTPUT
-    >
+    typeof streamText<TOOLS extends undefined ? AgentTools : TOOLS, OUTPUT>
   >[0],
   "model" | "prompt" | "messages"
 > & {
@@ -496,11 +488,7 @@ export interface Thread<DefaultTools extends ToolSet> {
     OUTPUT extends Output<any, any, any> = never,
   >(
     generateTextArgs: AgentPrompt &
-      TextArgs<
-        TOOLS extends undefined ? DefaultTools : TOOLS,
-        TOOLS,
-        OUTPUT
-      >,
+      TextArgs<TOOLS extends undefined ? DefaultTools : TOOLS, TOOLS, OUTPUT>,
     options?: Options,
   ): Promise<
     GenerateTextResult<TOOLS extends undefined ? DefaultTools : TOOLS, OUTPUT> &
@@ -541,10 +529,7 @@ export interface Thread<DefaultTools extends ToolSet> {
       saveStreamDeltas?: boolean | StreamingOptions;
     },
   ): Promise<
-    StreamTextResult<
-      TOOLS extends undefined ? DefaultTools : TOOLS,
-      OUTPUT
-    > &
+    StreamTextResult<TOOLS extends undefined ? DefaultTools : TOOLS, OUTPUT> &
       ThreadOutputMetadata
   >;
   /**
@@ -629,9 +614,9 @@ export type SyncStreamsReturnValue =
   | undefined;
 
 /* Type utils follow */
-export type QueryCtx = Pick<GenericQueryCtx<GenericDataModel>, "runQuery">;
+export type QueryCtx = Pick<GenericActionCtx<GenericDataModel>, "runQuery">;
 export type MutationCtx = Pick<
-  GenericMutationCtx<GenericDataModel>,
+  GenericActionCtx<GenericDataModel>,
   "runQuery" | "runMutation"
 >;
 export type ActionCtx = Pick<

package/src/component/apiKeys.ts

@@ -13,7 +13,7 @@ export const issue = mutation({
         .first();
       if (existingApiKey) {
         console.warn(`API key ${args.name} already exists, deleting...`);
-        await ctx.db.delete(existingApiKey._id);
+        await ctx.db.delete("apiKeys", existingApiKey._id);
       }
     }
     const apiKey = await ctx.db.insert("apiKeys", args);
@@ -27,7 +27,7 @@ export const validate = query({
     apiKey: v.id("apiKeys"),
   },
   handler: async (ctx, args) => {
-    const apiKey = await ctx.db.get(args.apiKey);
+    const apiKey = await ctx.db.get("apiKeys", args.apiKey);
     if (!apiKey) {
       throw new Error("Invalid API key");
     }
@@ -43,14 +43,14 @@ export const destroy = mutation({
   }),
   handler: async (ctx, args) => {
     if (args.apiKey) {
-      const apiKey = await ctx.db.get(args.apiKey);
+      const apiKey = await ctx.db.get("apiKeys", args.apiKey);
       if (!apiKey) {
         return "missing";
       }
       if (apiKey.name !== args.name) {
         return "name mismatch";
       }
-      await ctx.db.delete(args.apiKey);
+      await ctx.db.delete("apiKeys", args.apiKey);
     } else if (args.name) {
       const apiKey = await ctx.db
         .query("apiKeys")
@@ -59,7 +59,7 @@ export const destroy = mutation({
       if (!apiKey) {
         return "missing";
       }
-      await ctx.db.delete(apiKey._id);
+      await ctx.db.delete("apiKeys", apiKey._id);
     } else {
       return "must provide either apiKey or name";
     }

package/src/component/files.test.ts

@@ -96,7 +96,7 @@ describe("files", () => {
       });
       // Manually set refcount to 0
       await t.run(async (ctx) => {
-        await ctx.db.patch(fileId, { refcount: 0 });
+        await ctx.db.patch("files", fileId, { refcount: 0 });
       });
       files.push(fileId);
     }