@convex-dev/agent 0.1.3 → 0.1.4

package/src/react/deltas.test.ts

@@ -0,0 +1,298 @@
+import { describe, it, expect } from "vitest";
+import { mergeDeltas, applyDeltasToStreamMessage } from "./deltas";
+import type { StreamMessage, StreamDelta, TextStreamPart } from "../validators";
+
+function makeStreamMessage(
+  streamId: string,
+  order: number,
+  stepOrder: number
+): StreamMessage {
+  return {
+    streamId,
+    order,
+    stepOrder,
+  } as StreamMessage;
+}
+
+function makeDelta(
+  streamId: string,
+  start: number,
+  end: number,
+  parts: TextStreamPart[]
+): StreamDelta {
+  return {
+    streamId,
+    start,
+    end,
+    parts,
+  };
+}
+
+describe("mergeDeltas", () => {
+  it("merges a single text-delta into a message", () => {
+    const streamId = "s1";
+    const streamMessages = [makeStreamMessage(streamId, 1, 0)];
+    const deltas = [
+      makeDelta(streamId, 0, 5, [{ type: "text-delta", textDelta: "Hello" }]),
+    ];
+    const [messages, newStreams, changed] = mergeDeltas(
+      "thread1",
+      streamMessages,
+      [],
+      deltas
+    );
+    expect(messages).toHaveLength(1);
+    expect(messages[0].text).toBe("Hello");
+    expect(messages[0].message?.role).toBe("assistant");
+    expect(changed).toBe(true);
+    expect(newStreams[0].cursor).toBe(5);
+  });
+
+  it("merges multiple deltas for the same stream", () => {
+    const streamId = "s1";
+    const streamMessages = [makeStreamMessage(streamId, 1, 0)];
+    const deltas = [
+      makeDelta(streamId, 0, 5, [{ type: "text-delta", textDelta: "Hello" }]),
+      makeDelta(streamId, 5, 11, [
+        { type: "text-delta", textDelta: " World!" },
+      ]),
+    ];
+    const [messages, newStreams, changed] = mergeDeltas(
+      "thread1",
+      streamMessages,
+      [],
+      deltas
+    );
+    expect(messages).toHaveLength(1);
+    expect(messages[0].text).toBe("Hello World!");
+    expect(changed).toBe(true);
+    expect(newStreams[0].cursor).toBe(11);
+  });
+
+  it("handles tool-call and tool-result parts", () => {
+    const streamId = "s2";
+    const streamMessages = [makeStreamMessage(streamId, 2, 0)];
+    const deltas = [
+      makeDelta(streamId, 0, 1, [
+        {
+          type: "tool-call",
+          toolCallId: "call1",
+          toolName: "myTool",
+          args: "",
+        },
+      ]),
+      makeDelta(streamId, 1, 2, [
+        {
+          type: "tool-result",
+          toolCallId: "call1",
+          toolName: "myTool",
+          result: "42",
+        },
+      ]),
+    ];
+    const [messages, _, changed] = mergeDeltas(
+      "thread1",
+      streamMessages,
+      [],
+      deltas
+    );
+    expect(messages).toHaveLength(2);
+    expect(messages[0].message?.role).toBe("assistant");
+    expect(messages[0].tool).toBe(true);
+    const content = messages[0].message?.content;
+    expect(content).toEqual([
+      {
+        type: "tool-call",
+        toolCallId: "call1",
+        toolName: "myTool",
+        args: "",
+      },
+    ]);
+    expect(messages[1].message?.role).toBe("tool");
+    expect(messages[1].tool).toBe(true);
+    expect(messages[1].message?.content).toEqual([
+      {
+        type: "tool-result",
+        toolCallId: "call1",
+        toolName: "myTool",
+        result: "42",
+      },
+    ]);
+    expect(changed).toBe(true);
+  });
+
+  it("returns changed=false if no new deltas", () => {
+    const streamId = "s3";
+    const streamMessages = [makeStreamMessage(streamId, 3, 0)];
+    const deltas: StreamDelta[] = [];
+    const [messages, newStreams, changed] = mergeDeltas(
+      "thread1",
+      streamMessages,
+      [],
+      deltas
+    );
+    expect(messages).toHaveLength(0);
+    expect(changed).toBe(false);
+    expect(newStreams[0].cursor).toBe(0);
+  });
+
+  it("handles multiple streams and sorts by order/stepOrder", () => {
+    const s1 = makeStreamMessage("s1", 1, 0);
+    const s2 = makeStreamMessage("s2", 2, 0);
+    const deltas = [
+      makeDelta("s2", 0, 3, [{ type: "text-delta", textDelta: "B" }]),
+      makeDelta("s1", 0, 3, [{ type: "text-delta", textDelta: "A" }]),
+    ];
+    const [messages, _, changed] = mergeDeltas("thread1", [s2, s1], [], deltas);
+    expect(messages).toHaveLength(2);
+    expect(messages[0].text).toBe("A");
+    expect(messages[1].text).toBe("B");
+    expect(changed).toBe(true);
+    // Sorted by order
+    expect(messages[0].order).toBe(1);
+    expect(messages[1].order).toBe(2);
+  });
+
+  it("does not duplicate text content when merging sequential text-deltas", () => {
+    const streamId = "s4";
+    const streamMessages = [makeStreamMessage(streamId, 4, 0)];
+    const deltas = [
+      makeDelta(streamId, 0, 5, [{ type: "text-delta", textDelta: "Hello" }]),
+      makeDelta(streamId, 5, 11, [
+        { type: "text-delta", textDelta: " World!" },
+      ]),
+      makeDelta(streamId, 11, 12, [{ type: "text-delta", textDelta: "!" }]),
+    ];
+    const [messages] = mergeDeltas("thread1", streamMessages, [], deltas);
+    expect(messages).toHaveLength(1);
+    expect(messages[0].text).toBe("Hello World!!");
+    // There should only be one text part per message
+    const content = messages[0].message?.content;
+    if (Array.isArray(content)) {
+      const textParts = content.filter((p) => p.type === "text");
+      expect(textParts).toHaveLength(1);
+      expect(textParts[0].text).toBe("Hello World!!");
+    }
+  });
+
+  it("does not duplicate reasoning parts", () => {
+    const streamId = "s6";
+    const streamMessages = [makeStreamMessage(streamId, 6, 0)];
+    const deltas = [
+      makeDelta(streamId, 0, 1, [
+        { type: "reasoning", textDelta: "I'm thinking..." },
+      ]),
+      makeDelta(streamId, 1, 2, [
+        { type: "reasoning", textDelta: " Still thinking..." },
+      ]),
+    ];
+    const [messages] = mergeDeltas("thread1", streamMessages, [], deltas);
+    expect(messages).toHaveLength(1);
+    if (Array.isArray(messages[0].message?.content)) {
+      const reasoningParts = messages[0].message.content.filter(
+        (p) => p.type === "reasoning"
+      );
+      expect(reasoningParts).toHaveLength(1);
+      expect(reasoningParts[0].text).toBe("I'm thinking... Still thinking...");
+    }
+  });
+
+  it("applyDeltasToStreamMessage is idempotent and does not duplicate content", () => {
+    const streamId = "s7";
+    const streamMessage = makeStreamMessage(streamId, 7, 0);
+    const deltas = [
+      makeDelta(streamId, 0, 5, [{ type: "text-delta", textDelta: "Hello" }]),
+      makeDelta(streamId, 5, 11, [
+        { type: "text-delta", textDelta: " World!" },
+      ]),
+    ];
+    // First call: apply both deltas
+    let [result, changed] = applyDeltasToStreamMessage(
+      "thread1",
+      streamMessage,
+      undefined,
+      deltas
+    );
+    expect(result.messages).toHaveLength(1);
+    expect(result.messages[0].text).toBe("Hello World!");
+    // Second call: re-apply the same deltas (should not duplicate)
+    [result, changed] = applyDeltasToStreamMessage(
+      "thread1",
+      streamMessage,
+      result,
+      deltas
+    );
+    expect(result.messages).toHaveLength(1);
+    expect(result.messages[0].text).toBe("Hello World!");
+    // Third call: add a new delta
+    const moreDeltas = [
+      ...deltas,
+      makeDelta(streamId, 11, 12, [{ type: "text-delta", textDelta: "!" }]),
+    ];
+    [result, changed] = applyDeltasToStreamMessage(
+      "thread1",
+      streamMessage,
+      result,
+      moreDeltas
+    );
+    expect(changed).toBe(true);
+    expect(result.messages).toHaveLength(1);
+    expect(result.messages[0].text).toBe("Hello World!!");
+    // Re-apply all deltas again (should still not duplicate)
+    [result, changed] = applyDeltasToStreamMessage(
+      "thread1",
+      streamMessage,
+      result,
+      moreDeltas
+    );
+    expect(changed).toBe(false);
+    expect(result.messages).toHaveLength(1);
+    expect(result.messages[0].text).toBe("Hello World!!");
+  });
+
+  it("mergeDeltas is pure and does not mutate inputs", () => {
+    const streamId = "s8";
+    const streamMessages = [makeStreamMessage(streamId, 8, 0)];
+    const deltas = [
+      makeDelta(streamId, 0, 5, [{ type: "text-delta", textDelta: "Hello" }]),
+      makeDelta(streamId, 5, 11, [
+        { type: "text-delta", textDelta: " World!" },
+      ]),
+    ];
+    // Deep freeze inputs to catch mutation
+    function deepFreeze(obj: unknown): unknown {
+      if (obj && typeof obj === "object" && !Object.isFrozen(obj)) {
+        Object.freeze(obj);
+        for (const key of Object.keys(obj)) {
+          deepFreeze((obj as Record<string, unknown>)[key]);
+        }
+      }
+      return obj;
+    }
+    deepFreeze(streamMessages);
+    deepFreeze(deltas);
+    const [messages1, streams1, changed1] = mergeDeltas(
+      "thread1",
+      streamMessages,
+      [],
+      deltas
+    );
+    const [messages2, streams2, changed2] = mergeDeltas(
+      "thread1",
+      streamMessages,
+      [],
+      deltas
+    );
+    expect(messages1).toEqual(messages2);
+    expect(streams1).toEqual(streams2);
+    expect(changed1).toBe(changed2);
+    // Inputs should remain unchanged
+    expect(streamMessages).toEqual([makeStreamMessage(streamId, 8, 0)]);
+    expect(deltas).toEqual([
+      makeDelta(streamId, 0, 5, [{ type: "text-delta", textDelta: "Hello" }]),
+      makeDelta(streamId, 5, 11, [
+        { type: "text-delta", textDelta: " World!" },
+      ]),
+    ]);
+  });
+});

package/src/react/deltas.ts

@@ -0,0 +1,340 @@
+import type { TextPart, ToolCallPart, ToolResultPart } from "ai";
+import type { MessageDoc } from "../client";
+import type {
+  Message,
+  StreamDelta,
+  StreamMessage,
+  TextStreamPart,
+} from "../validators";
+import type { UIMessage } from "./toUIMessages";
+import { toUIMessages } from "./toUIMessages";
+
+export { toUIMessages, type UIMessage };
+
+export function mergeDeltas(
+  threadId: string,
+  streamMessages: StreamMessage[],
+  existingStreams: Array<{
+    streamId: string;
+    cursor: number;
+    messages: MessageDoc[];
+  }>,
+  allDeltas: StreamDelta[]
+): [
+  MessageDoc[],
+  Array<{ streamId: string; cursor: number; messages: MessageDoc[] }>,
+  boolean,
+] {
+  const newStreams: Array<{
+    streamId: string;
+    cursor: number;
+    messages: MessageDoc[];
+  }> = [];
+  // Seed the existing chunks
+  let changed = false;
+  for (const streamMessage of streamMessages) {
+    const deltas = allDeltas.filter(
+      (d) => d.streamId === streamMessage.streamId
+    );
+    const existing = existingStreams.find(
+      (s) => s.streamId === streamMessage.streamId
+    );
+    const [newStream, messageChanged] = applyDeltasToStreamMessage(
+      threadId,
+      streamMessage,
+      existing,
+      deltas
+    );
+    newStreams.push(newStream);
+    if (messageChanged) changed = true;
+  }
+  for (const { streamId } of existingStreams) {
+    if (!newStreams.find((s) => s.streamId === streamId)) {
+      // There's a stream that's no longer active.
+      changed = true;
+    }
+  }
+  const messages = newStreams
+    .map((s) => s.messages)
+    .flat()
+    .sort((a, b) => a.order - b.order || a.stepOrder - b.stepOrder);
+  return [messages, newStreams, changed];
+}
+
+// exported for testing
+export function applyDeltasToStreamMessage(
+  threadId: string,
+  streamMessage: StreamMessage,
+  existing:
+    | { streamId: string; cursor: number; messages: MessageDoc[] }
+    | undefined,
+  deltas: StreamDelta[]
+): [{ streamId: string; cursor: number; messages: MessageDoc[] }, boolean] {
+  let changed = false;
+  let cursor = existing?.cursor ?? 0;
+  let parts: TextStreamPart[] = [];
+  for (const delta of deltas.sort((a, b) => a.start - b.start)) {
+    if (delta.parts.length === 0) {
+      console.warn(`Got delta with no parts: ${JSON.stringify(delta)}`);
+      continue;
+    }
+    if (cursor !== delta.start) {
+      if (cursor >= delta.end) {
+        console.debug(
+          `Got duplicate delta for stream ${delta.streamId} at ${delta.start}`
+        );
+        continue;
+      } else if (cursor < delta.start) {
+        console.warn(
+          `Got delta for stream ${delta.streamId} that has a gap ${cursor} -> ${delta.start}`
+        );
+        continue;
+      } else {
+        throw new Error(
+          `Got unexpected delta for stream ${delta.streamId}: delta: ${delta.start} -> ${delta.end} existing cursor: ${cursor}`
+        );
+      }
+    }
+    changed = true;
+    cursor = delta.end;
+    parts.push(...delta.parts);
+  }
+  if (!changed) {
+    return [
+      existing ?? { streamId: streamMessage.streamId, cursor, messages: [] },
+      false,
+    ];
+  }
+
+  const existingMessages = existing?.messages ?? [];
+
+  let currentMessage: MessageDoc;
+  if (existingMessages.length > 0) {
+    // replace the last message with a new one
+    const lastMessage = existingMessages.at(-1)!;
+    currentMessage = {
+      ...lastMessage,
+      message: cloneMessageAndContent(lastMessage.message),
+    };
+  } else {
+    const newMessage = createStreamingMessage(
+      threadId,
+      streamMessage,
+      parts[0]!,
+      existingMessages.length
+    );
+    parts = parts.slice(1);
+    currentMessage = newMessage;
+  }
+  const newStream = {
+    streamId: streamMessage.streamId,
+    cursor,
+    messages: [...existingMessages.slice(0, -1), currentMessage],
+  };
+  let lastContent = getLastContent(currentMessage);
+  for (const part of parts) {
+    let contentToAdd:
+      | TextPart
+      | ToolCallPart
+      | { type: "reasoning"; text: string }
+      | ToolResultPart
+      | undefined;
+    const isToolRole = part.type === "source" || part.type === "tool-result";
+    if (isToolRole !== (currentMessage.message!.role === "tool")) {
+      currentMessage = createStreamingMessage(
+        threadId,
+        streamMessage,
+        part,
+        newStream.messages.length
+      );
+      lastContent = getLastContent(currentMessage);
+      newStream.messages.push(currentMessage);
+      continue;
+    }
+    switch (part.type) {
+      case "text-delta":
+        currentMessage.text += part.textDelta;
+        if (lastContent?.type === "text") {
+          lastContent.text += part.textDelta;
+        } else {
+          contentToAdd = {
+            type: "text",
+            text: part.textDelta,
+          };
+        }
+        break;
+      case "tool-call-streaming-start":
+        currentMessage.tool = true;
+        contentToAdd = {
+          type: "tool-call",
+          toolCallId: part.toolCallId,
+          toolName: part.toolName,
+          args: "",
+        };
+        break;
+      case "tool-call-delta":
+        {
+          currentMessage.tool = true;
+          if (lastContent?.type !== "tool-call") {
+            throw new Error("Expected last content to be a tool call");
+          }
+          if (typeof lastContent.args !== "string") {
+            throw new Error("Expected args to be a string");
+          }
+          lastContent.args += part.argsTextDelta;
+        }
+        break;
+      case "tool-call":
+        currentMessage.tool = true;
+        contentToAdd = part;
+        break;
+      case "reasoning":
+        if (lastContent?.type === "reasoning") {
+          lastContent.text += part.textDelta;
+        } else {
+          contentToAdd = {
+            type: "reasoning",
+            text: part.textDelta,
+          };
+        }
+        break;
+      case "source":
+        if (!currentMessage.sources) {
+          currentMessage.sources = [];
+        }
+        currentMessage.sources.push(part.source);
+        break;
+      case "tool-result":
+        contentToAdd = part;
+        break;
+      default:
+        console.warn(`Received unexpected part: ${JSON.stringify(part)}`);
+        break;
+    }
+    if (contentToAdd) {
+      if (!currentMessage.message!.content) {
+        currentMessage.message!.content = [];
+      }
+      if (!Array.isArray(currentMessage.message?.content)) {
+        throw new Error("Expected message content to be an array");
+      }
+      // eslint-disable-next-line @typescript-eslint/no-explicit-any
+      currentMessage.message.content.push(contentToAdd as any);
+      lastContent = contentToAdd;
+    }
+  }
+  return [newStream, true];
+}
+
+function cloneMessageAndContent(
+  message: Message | undefined
+): Message | undefined {
+  return (
+    message &&
+    ({
+      ...message,
+      content: Array.isArray(message.content)
+        ? [...message.content]
+        : message.content,
+    } as typeof message)
+  );
+}
+
+function getLastContent(message: MessageDoc) {
+  if (Array.isArray(message.message?.content)) {
+    return message.message.content.at(-1);
+  }
+  return undefined;
+}
+
+export function createStreamingMessage(
+  threadId: string,
+  message: StreamMessage,
+  part: TextStreamPart,
+  index: number
+): MessageDoc {
+  const { streamId, ...rest } = message;
+  const metadata: MessageDoc = {
+    _id: `${streamId}-${index}`,
+    _creationTime: Date.now(),
+    status: "pending",
+    threadId,
+    tool: false,
+    ...rest,
+  };
+  switch (part.type) {
+    case "text-delta":
+      return {
+        ...metadata,
+        message: {
+          role: "assistant",
+          content: [{ type: "text", text: part.textDelta }],
+        },
+        text: part.textDelta,
+      };
+    case "tool-call-streaming-start":
+      return {
+        ...metadata,
+        tool: true,
+        message: {
+          role: "assistant",
+          content: [
+            {
+              type: "tool-call",
+              toolName: part.toolName,
+              toolCallId: part.toolCallId,
+              args: "", // when it's a string, it's a partial call
+            },
+          ],
+        },
+      };
+    case "reasoning":
+      return {
+        ...metadata,
+        message: {
+          role: "assistant",
+          content: [{ type: "reasoning", text: part.textDelta }],
+        },
+        reasoning: part.textDelta,
+      };
+    case "source":
+      console.warn("Received source part first??");
+      return {
+        ...metadata,
+        tool: true,
+        message: { role: "tool", content: [] },
+        sources: [part.source],
+      };
+    case "tool-call":
+      return {
+        ...metadata,
+        tool: true,
+        message: { role: "assistant", content: [part] },
+      };
+    case "tool-call-delta":
+      console.warn("Received tool call delta part first??");
+      return {
+        ...metadata,
+        tool: true,
+        message: {
+          role: "assistant",
+          content: [
+            {
+              type: "tool-call",
+              toolCallId: part.toolCallId,
+              toolName: part.toolName,
+              args: part.argsTextDelta,
+            },
+          ],
+        },
+      };
+    case "tool-result":
+      return {
+        ...metadata,
+        tool: true,
+        message: { role: "tool", content: [part] },
+      };
+    default:
+      throw new Error(`Unexpected part type: ${JSON.stringify(part)}`);
+  }
+}