@convex-dev/agent 0.6.3 → 0.6.4

package/src/mapping.test.ts

@@ -4,6 +4,7 @@ import {
   serializeDataOrUrl,
   toModelMessageDataOrUrl,
   serializeMessage,
+  serializeNewMessagesInStep,
   toModelMessage,
   serializeContent,
   toModelMessageContent,
@@ -16,7 +17,7 @@ import fs from "fs";
 import path from "path";
 import type { SerializedContent } from "./mapping.js";
 import { validate } from "convex-helpers/validators";
-import type { ToolResultPart } from "ai";
+import type { ModelMessage, StepResult, ToolResultPart, ToolSet } from "ai";
 import type { Infer } from "convex/values";
 
 const testAssetsDir = path.join(__dirname, "../test-assets");
@@ -259,6 +260,245 @@ describe("mapping", () => {
     expect((content as unknown[])[0]).toMatchObject(approvalResponse);
   });
 
+  describe("serializeNewMessagesInStep", () => {
+    const ctx = {
+      runAction: async () => undefined,
+      runMutation: async () => undefined,
+      storage: {
+        store: async () => "storageId",
+        getUrl: async () => "https://example.com/file",
+        delete: async () => undefined,
+      },
+    } as unknown as ActionCtx;
+    const component = api as unknown as AgentComponent;
+
+    const step0Messages: ModelMessage[] = [
+      {
+        role: "assistant",
+        content: [
+          { type: "tool-call", toolCallId: "c1", toolName: "search", input: {} },
+        ],
+      },
+      {
+        role: "tool",
+        content: [
+          {
+            type: "tool-result",
+            toolCallId: "c1",
+            toolName: "search",
+            output: { type: "text", value: "ok" },
+          },
+        ],
+      },
+    ];
+    const step1Messages: ModelMessage[] = [
+      ...step0Messages,
+      { role: "assistant", content: [{ type: "text", text: "thinking" }] },
+    ];
+    const step2Messages: ModelMessage[] = [
+      ...step1Messages,
+      {
+        role: "assistant",
+        content: [
+          { type: "tool-call", toolCallId: "c2", toolName: "search", input: {} },
+        ],
+      },
+      {
+        role: "tool",
+        content: [
+          {
+            type: "tool-result",
+            toolCallId: "c2",
+            toolName: "search",
+            output: { type: "text", value: "done" },
+          },
+        ],
+      },
+    ];
+
+    const makeStep = (messages: ModelMessage[]): StepResult<ToolSet> =>
+      ({
+        content: [],
+        text: "",
+        reasoning: [],
+        reasoningText: undefined,
+        files: [],
+        sources: [],
+        toolCalls: [],
+        staticToolCalls: [],
+        dynamicToolCalls: [],
+        toolResults: [],
+        staticToolResults: [],
+        dynamicToolResults: [],
+        finishReason: "stop",
+        rawFinishReason: undefined,
+        usage: { inputTokens: 1, outputTokens: 1, totalTokens: 2 },
+        warnings: undefined,
+        request: {},
+        response: {
+          id: "resp",
+          timestamp: new Date(),
+          modelId: "test",
+          messages,
+        },
+        providerMetadata: undefined,
+      }) as unknown as StepResult<ToolSet>;
+
+    const contentTypes = (msg: { content: unknown }): string[] => {
+      const c = msg.content;
+      if (!Array.isArray(c)) return ["text"];
+      return c.map((p: { type?: string }) => p.type ?? "?");
+    };
+
+    test("first step (count=0) serializes all response messages", async () => {
+      const res = await serializeNewMessagesInStep(
+        ctx,
+        component,
+        makeStep(step0Messages),
+        undefined,
+        0,
+      );
+      expect(res.messages).toHaveLength(2);
+      expect(res.messages[0].message.role).toBe("assistant");
+      expect(contentTypes(res.messages[0].message)).toEqual(["tool-call"]);
+      expect(res.messages[1].message.role).toBe("tool");
+      expect(contentTypes(res.messages[1].message)).toEqual(["tool-result"]);
+    });
+
+    test("middle step (count=2) serializes only the new text message", async () => {
+      const res = await serializeNewMessagesInStep(
+        ctx,
+        component,
+        makeStep(step1Messages),
+        undefined,
+        2,
+      );
+      expect(res.messages).toHaveLength(1);
+      expect(res.messages[0].message.role).toBe("assistant");
+      expect(contentTypes(res.messages[0].message)).toEqual(["text"]);
+    });
+
+    test("multi-message step (count=3) serializes the new tool-call + tool-result pair", async () => {
+      const res = await serializeNewMessagesInStep(
+        ctx,
+        component,
+        makeStep(step2Messages),
+        undefined,
+        3,
+      );
+      expect(res.messages).toHaveLength(2);
+      expect(res.messages[0].message.role).toBe("assistant");
+      expect(contentTypes(res.messages[0].message)).toEqual(["tool-call"]);
+      expect(res.messages[1].message.role).toBe("tool");
+      expect(contentTypes(res.messages[1].message)).toEqual(["tool-result"]);
+    });
+
+    // Regression test for the actually-broken shape: a single step appended
+    // assistant(text) + assistant(tool-call) + tool(tool-result), so the new
+    // tail has length 3 and the last message is a tool message. The old
+    // heuristic took `slice(-2)` whenever the last role was "tool" and would
+    // have dropped the leading text. The watermark returns all three.
+    test("returns all three messages when a step adds text + tool-call + tool-result", async () => {
+      const stepMessages: ModelMessage[] = [
+        ...step0Messages, // length 2
+        { role: "assistant", content: [{ type: "text", text: "Let me check..." }] },
+        {
+          role: "assistant",
+          content: [
+            { type: "tool-call", toolCallId: "c3", toolName: "search", input: {} },
+          ],
+        },
+        {
+          role: "tool",
+          content: [
+            {
+              type: "tool-result",
+              toolCallId: "c3",
+              toolName: "search",
+              output: { type: "text", value: "done" },
+            },
+          ],
+        },
+      ];
+      const res = await serializeNewMessagesInStep(
+        ctx,
+        component,
+        makeStep(stepMessages),
+        undefined,
+        step0Messages.length,
+      );
+      expect(res.messages).toHaveLength(3);
+      expect(res.messages[0].message.role).toBe("assistant");
+      expect(contentTypes(res.messages[0].message)).toEqual(["text"]);
+      expect(res.messages[1].message.role).toBe("assistant");
+      expect(contentTypes(res.messages[1].message)).toEqual(["tool-call"]);
+      expect(res.messages[2].message.role).toBe("tool");
+      expect(contentTypes(res.messages[2].message)).toEqual(["tool-result"]);
+    });
+
+    test("empty response messages slice falls back to synthetic empty assistant", async () => {
+      const res = await serializeNewMessagesInStep(
+        ctx,
+        component,
+        makeStep(step1Messages),
+        undefined,
+        step1Messages.length,
+      );
+      expect(res.messages).toHaveLength(1);
+      expect(res.messages[0].message.role).toBe("assistant");
+      expect(res.messages[0].message.content).toEqual([]);
+    });
+
+    // Pin the caller-drift behavior: if the watermark is past the end of
+    // response.messages (e.g. the caller mistracked), the slice is empty and
+    // we fall through to the synthetic anchor. Future "fixes" should not
+    // accidentally change this without intent.
+    test("watermark beyond response.messages.length returns the synthetic fallback", async () => {
+      const res = await serializeNewMessagesInStep(
+        ctx,
+        component,
+        makeStep(step1Messages),
+        undefined,
+        step1Messages.length + 5,
+      );
+      expect(res.messages).toHaveLength(1);
+      expect(res.messages[0].message.role).toBe("assistant");
+      expect(res.messages[0].message.content).toEqual([]);
+    });
+
+    // AI SDK v6 makes step.response.messages cumulative across steps:
+    // step N's array contains all messages from steps 0..N. Without the
+    // previousResponseMessageCount watermark, every multi-step save duplicates
+    // all prior messages. These tests demonstrate the bug and the fix.
+    describe("multi-step loop — previousStep watermark", () => {
+      test("without watermark, step 2 re-saves all cumulative messages (demonstrates the bug)", async () => {
+        // step2Messages = step0 (2 msgs) + step1 (1 msg) + step2 new (2 msgs) = 5 total
+        const res = await serializeNewMessagesInStep(
+          ctx,
+          component,
+          makeStep(step2Messages),
+          undefined,
+          0,
+        );
+        expect(res.messages).toHaveLength(5);
+      });
+
+      test("with watermark, step 2 saves only its 2 new messages", async () => {
+        const step1 = makeStep(step1Messages);
+        const res = await serializeNewMessagesInStep(
+          ctx,
+          component,
+          makeStep(step2Messages),
+          undefined,
+          step1.response.messages.length,
+        );
+        expect(res.messages).toHaveLength(2);
+        expect(contentTypes(res.messages[0].message)).toEqual(["tool-call"]);
+        expect(contentTypes(res.messages[1].message)).toEqual(["tool-result"]);
+      });
+    });
+  });
+
   describe("autoDenyUnresolvedApprovals", () => {
     test("returns messages unchanged when no unresolved approvals", () => {
       const messages = [
@@ -267,13 +507,21 @@ describe("mapping", () => {
           role: "assistant" as const,
           content: [
             { type: "tool-call", toolCallId: "tc1", toolName: "a", input: {} },
-            { type: "tool-approval-request", approvalId: "ap1", toolCallId: "tc1" },
+            {
+              type: "tool-approval-request",
+              approvalId: "ap1",
+              toolCallId: "tc1",
+            },
           ],
         },
         {
           role: "tool" as const,
           content: [
-            { type: "tool-approval-response", approvalId: "ap1", approved: true },
+            {
+              type: "tool-approval-response",
+              approvalId: "ap1",
+              approved: true,
+            },
           ],
         },
       ] as any;
@@ -289,7 +537,11 @@ describe("mapping", () => {
           role: "assistant" as const,
           content: [
             { type: "tool-call", toolCallId: "tc1", toolName: "a", input: {} },
-            { type: "tool-approval-request", approvalId: "ap1", toolCallId: "tc1" },
+            {
+              type: "tool-approval-request",
+              approvalId: "ap1",
+              toolCallId: "tc1",
+            },
           ],
         },
         { role: "user" as const, content: "new message" },
@@ -304,7 +556,9 @@ describe("mapping", () => {
       expect(denialContent[0].type).toBe("tool-approval-response");
       expect(denialContent[0].approvalId).toBe("ap1");
       expect(denialContent[0].approved).toBe(false);
-      expect(denialContent[0].reason).toBe("auto-denied: new generation started");
+      expect(denialContent[0].reason).toBe(
+        "auto-denied: new generation started",
+      );
       // The new user message should follow
       expect(result[3].role).toBe("user");
       expect(result[3].content).toBe("new message");
@@ -317,8 +571,16 @@ describe("mapping", () => {
           content: [
             { type: "tool-call", toolCallId: "tc1", toolName: "a", input: {} },
             { type: "tool-call", toolCallId: "tc2", toolName: "b", input: {} },
-            { type: "tool-approval-request", approvalId: "ap1", toolCallId: "tc1" },
-            { type: "tool-approval-request", approvalId: "ap2", toolCallId: "tc2" },
+            {
+              type: "tool-approval-request",
+              approvalId: "ap1",
+              toolCallId: "tc1",
+            },
+            {
+              type: "tool-approval-request",
+              approvalId: "ap2",
+              toolCallId: "tc2",
+            },
           ],
         },
       ] as any;
@@ -341,14 +603,26 @@ describe("mapping", () => {
           content: [
             { type: "tool-call", toolCallId: "tc1", toolName: "a", input: {} },
             { type: "tool-call", toolCallId: "tc2", toolName: "b", input: {} },
-            { type: "tool-approval-request", approvalId: "ap1", toolCallId: "tc1" },
-            { type: "tool-approval-request", approvalId: "ap2", toolCallId: "tc2" },
+            {
+              type: "tool-approval-request",
+              approvalId: "ap1",
+              toolCallId: "tc1",
+            },
+            {
+              type: "tool-approval-request",
+              approvalId: "ap2",
+              toolCallId: "tc2",
+            },
           ],
         },
         {
           role: "tool" as const,
           content: [
-            { type: "tool-approval-response", approvalId: "ap1", approved: true },
+            {
+              type: "tool-approval-response",
+              approvalId: "ap1",
+              approved: true,
+            },
           ],
         },
         { role: "user" as const, content: "next question" },
@@ -381,8 +655,16 @@ describe("mapping", () => {
           content: [
             { type: "tool-call", toolCallId: "tc1", toolName: "a", input: {} },
             { type: "tool-call", toolCallId: "tc2", toolName: "b", input: {} },
-            { type: "tool-approval-request", approvalId: "ap1", toolCallId: "tc1" },
-            { type: "tool-approval-request", approvalId: "ap2", toolCallId: "tc2" },
+            {
+              type: "tool-approval-request",
+              approvalId: "ap1",
+              toolCallId: "tc1",
+            },
+            {
+              type: "tool-approval-request",
+              approvalId: "ap2",
+              toolCallId: "tc2",
+            },
           ],
         },
       ] as any;
@@ -390,12 +672,8 @@ describe("mapping", () => {
       autoDenyUnresolvedApprovals(messages);
 
       expect(warnSpy).toHaveBeenCalledTimes(2);
-      expect(warnSpy).toHaveBeenCalledWith(
-        expect.stringContaining("ap1"),
-      );
-      expect(warnSpy).toHaveBeenCalledWith(
-        expect.stringContaining("ap2"),
-      );
+      expect(warnSpy).toHaveBeenCalledWith(expect.stringContaining("ap1"));
+      expect(warnSpy).toHaveBeenCalledWith(expect.stringContaining("ap2"));
       warnSpy.mockRestore();
     });
   });

package/src/mapping.ts

@@ -313,24 +313,30 @@ export async function serializeResponseMessages<TOOLS extends ToolSet>(
 }
 
 /**
- * Serialize the new messages from a step using a heuristic to determine
- * which response messages are new (last 1-2 messages).
+ * Serialize the new response messages produced by this step.
+ *
+ * `step.response.messages` is cumulative across steps in AI SDK v6 — each
+ * step's array contains all messages from prior steps too. Pass
+ * `previousResponseMessageCount` (the prior step's `response.messages.length`,
+ * or `0` for the first step) so we slice only the new tail. The parameter is
+ * required: defaulting it would silently duplicate every prior message on
+ * every multi-step save.
  */
 export async function serializeNewMessagesInStep<TOOLS extends ToolSet>(
   ctx: ActionCtx,
   component: AgentComponent,
   step: StepResult<TOOLS>,
   model: ModelOrMetadata | undefined,
+  previousResponseMessageCount: number,
 ): 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: [] }];
-  }
+  const newMessages = step.response.messages.slice(previousResponseMessageCount);
+  // Keep at least one message in the output so the step still anchors an
+  // order slot — downstream `addMessages` relies on each step contributing a
+  // row even when AI SDK produced no response messages.
+  const messagesToSerialize: ModelMessage[] =
+    newMessages.length > 0
+      ? newMessages
+      : [{ role: "assistant" as const, content: [] }];
   return serializeStepMessages(ctx, component, step, model, messagesToSerialize);
 }
 

package/src/react/useStreamingUIMessages.ts

@@ -4,10 +4,12 @@ import { type UIDataTypes, type UIMessageChunk, type UITools } from "ai";
 import type { StreamQuery, StreamQueryArgs } from "./types.js";
 import { type UIMessage } from "../UIMessages.js";
 import {
+  applyUIMessageChunksIncremental,
   blankUIMessage,
+  emptyIncrementalStreamState,
   getParts,
-  updateFromUIMessageChunks,
-  deriveUIMessagesFromTextStreamParts,
+  statusFromStreamStatus,
+  type IncrementalStreamState,
 } from "../deltas.js";
 import { useDeltaStreams } from "./useDeltaStreams.js";
 
@@ -53,6 +55,7 @@ export function useStreamingUIMessages<
       {
         uiMessage: UIMessage<METADATA, DATA_PARTS, TOOLS>;
         cursor: number;
+        streamState: IncrementalStreamState;
       }
     >
   >({});
@@ -63,16 +66,23 @@ export function useStreamingUIMessages<
 
   useEffect(() => {
     if (!streams) return;
-    // return if there are no new deltas beyond the cursors
     let noNewDeltas = true;
     for (const stream of streams) {
-      const lastDelta = stream.deltas.at(-1);
-      const cursor = messageState[stream.streamMessage.streamId]?.cursor;
-      if (!cursor) {
+      const existingStreamState = messageState[stream.streamMessage.streamId];
+      const cursor = existingStreamState?.cursor;
+      if (existingStreamState === undefined || cursor === undefined) {
         noNewDeltas = false;
         break;
       }
-      if (lastDelta && lastDelta.start >= cursor) {
+      if (stream.deltas.some((d) => d.parts.length > 0 && d.end > cursor)) {
+        noNewDeltas = false;
+        break;
+      }
+      if (
+        existingStreamState &&
+        existingStreamState.uiMessage.status !==
+          statusFromStreamStatus(stream.streamMessage.status)
+      ) {
         noNewDeltas = false;
         break;
       }
@@ -87,40 +97,58 @@ export function useStreamingUIMessages<
         {
           uiMessage: UIMessage<METADATA, DATA_PARTS, TOOLS>;
           cursor: number;
+          streamState: IncrementalStreamState;
         }
       > = Object.fromEntries(
         await Promise.all(
           streams.map(async ({ deltas, streamMessage }) => {
-            const { parts, cursor } = getParts<UIMessageChunk>(deltas, 0);
-            if (streamMessage.format === "UIMessageChunk") {
-              // Unfortunately this can't handle resuming from a UIMessage and
-              // adding more chunks, so we re-create it from scratch each time.
-              const uiMessage = await updateFromUIMessageChunks(
-                blankUIMessage(streamMessage, threadId),
-                parts,
-              );
-              return [
-                streamMessage.streamId,
-                {
-                  uiMessage,
-                  cursor,
-                },
-              ];
-            } else {
-              const [uiMessages] = deriveUIMessagesFromTextStreamParts(
-                threadId,
-                [streamMessage],
-                [],
-                deltas,
-              );
+            const streamId = streamMessage.streamId;
+            const existing = messageState[streamId];
+            const fromCursor = existing?.cursor ?? 0;
+            const status = statusFromStreamStatus(streamMessage.status);
+            const prevState =
+              existing?.streamState ?? emptyIncrementalStreamState();
+
+            const { parts: newParts, cursor } = getParts<UIMessageChunk>(
+              deltas,
+              fromCursor,
+            );
+
+            const base =
+              existing?.uiMessage ??
+              blankUIMessage(streamMessage, threadId as string);
+
+            if (newParts.length === 0) {
+              if (existing && existing.uiMessage.status !== status) {
+                return [
+                  streamId,
+                  {
+                    uiMessage: { ...existing.uiMessage, status },
+                    cursor: existing.cursor,
+                    streamState: prevState,
+                  },
+                ];
+              }
               return [
-                streamMessage.streamId,
-                {
-                  uiMessage: uiMessages[0],
-                  cursor,
-                },
+                streamId,
+                existing ?? { uiMessage: base, cursor: 0, streamState: prevState },
               ];
             }
+
+            const { message, streamState } = applyUIMessageChunksIncremental(
+              base as UIMessage,
+              newParts,
+              prevState,
+            );
+            message.status = status;
+            return [
+              streamId,
+              {
+                uiMessage: message as UIMessage<METADATA, DATA_PARTS, TOOLS>,
+                cursor,
+                streamState,
+              },
+            ];
           }),
         ),
       );

package/src/react/useUIMessages.test.ts

@@ -1,5 +1,9 @@
 import { describe, it, expect } from "vitest";
-import { dedupeMessages } from "./useUIMessages.js";
+import {
+  dedupeMessages,
+  mergeUIMessages,
+  type UIMessageLike,
+} from "./useUIMessages.js";
 
 type TestMessage = {
   order: number;
@@ -8,6 +12,39 @@ type TestMessage = {
   id: string;
 };
 
+type TestUIMessage = UIMessageLike & {
+  id: string;
+  key: string;
+  text: string;
+  _creationTime: number;
+};
+
+function testUIMessage({
+  id,
+  order,
+  stepOrder,
+  status,
+  text,
+}: {
+  id: string;
+  order: number;
+  stepOrder: number;
+  status: TestUIMessage["status"];
+  text: string;
+}): TestUIMessage {
+  return {
+    id,
+    key: `thread-${order}-${stepOrder}`,
+    order,
+    stepOrder,
+    status,
+    role: "assistant",
+    parts: [{ type: "text", text }],
+    text,
+    _creationTime: 0,
+  };
+}
+
 describe("dedupeMessages", () => {
   it("should prefer messages from messages list when streaming messages are absent", () => {
     const messages: TestMessage[] = [
@@ -253,3 +290,45 @@ describe("dedupeMessages", () => {
     expect(result[0].id).toBe("messages-success");
   });
 });
+
+describe("mergeUIMessages", () => {
+  it("dedupes streaming steps before combining assistant messages", () => {
+    const messages = [
+      testUIMessage({
+        id: "persisted-step-1",
+        order: 1,
+        stepOrder: 1,
+        status: "success",
+        text: "7",
+      }),
+      testUIMessage({
+        id: "persisted-step-2",
+        order: 1,
+        stepOrder: 2,
+        status: "success",
+        text: "8",
+      }),
+    ];
+    const streamMessages = [
+      testUIMessage({
+        id: "stream-step-2",
+        order: 1,
+        stepOrder: 2,
+        status: "streaming",
+        text: "8",
+      }),
+    ];
+
+    const result = mergeUIMessages(messages, streamMessages);
+
+    expect(result).toHaveLength(1);
+    expect(result[0].id).toBe("persisted-step-1");
+    expect(result[0].stepOrder).toBe(1);
+    expect(result[0].status).toBe("success");
+    expect(result[0].text).toBe("7 8");
+    expect(result[0].parts).toEqual([
+      { type: "text", text: "7" },
+      { type: "text", text: "8" },
+    ]);
+  });
+});