@convex-dev/agent 0.1.18 → 0.2.0

package/src/react/deltas.test.ts

@@ -1,36 +1,24 @@
 import { describe, it, expect } from "vitest";
 import { mergeDeltas, applyDeltasToStreamMessage } from "./deltas.js";
-import type {
-  StreamMessage,
-  StreamDelta,
-  TextStreamPart,
-} from "../validators.js";
+import type { StreamMessage, StreamDelta } from "../validators.js";
 import { omit } from "convex-helpers";
+import type { TextStreamPart, ToolSet } from "ai";
 
 function makeStreamMessage(
   streamId: string,
   order: number,
   stepOrder: number,
 ): StreamMessage {
-  return {
-    streamId,
-    order,
-    stepOrder,
-  } as StreamMessage;
+  return { streamId, order, stepOrder } as StreamMessage;
 }
 
 function makeDelta(
   streamId: string,
   start: number,
   end: number,
-  parts: TextStreamPart[],
+  parts: TextStreamPart<ToolSet>[],
 ): StreamDelta {
-  return {
-    streamId,
-    start,
-    end,
-    parts,
-  };
+  return { streamId, start, end, parts };
 }
 
 describe("mergeDeltas", () => {
@@ -38,7 +26,9 @@ describe("mergeDeltas", () => {
     const streamId = "s1";
     const streamMessages = [makeStreamMessage(streamId, 1, 0)];
     const deltas = [
-      makeDelta(streamId, 0, 5, [{ type: "text-delta", textDelta: "Hello" }]),
+      makeDelta(streamId, 0, 5, [
+        { type: "text-delta", id: "1", text: "Hello" },
+      ]),
     ];
     const [messages, newStreams, changed] = mergeDeltas(
       "thread1",
@@ -57,9 +47,11 @@ describe("mergeDeltas", () => {
     const streamId = "s1";
     const streamMessages = [makeStreamMessage(streamId, 1, 0)];
     const deltas = [
-      makeDelta(streamId, 0, 5, [{ type: "text-delta", textDelta: "Hello" }]),
+      makeDelta(streamId, 0, 5, [
+        { type: "text-delta", id: "1", text: "Hello" },
+      ]),
       makeDelta(streamId, 5, 11, [
-        { type: "text-delta", textDelta: " World!" },
+        { type: "text-delta", id: "2", text: " World!" },
       ]),
     ];
     const [messages, newStreams, changed] = mergeDeltas(
@@ -83,7 +75,7 @@ describe("mergeDeltas", () => {
           type: "tool-call",
           toolCallId: "call1",
           toolName: "myTool",
-          args: "",
+          input: "",
         },
       ]),
       makeDelta(streamId, 1, 2, [
@@ -91,7 +83,8 @@ describe("mergeDeltas", () => {
           type: "tool-result",
           toolCallId: "call1",
           toolName: "myTool",
-          result: "42",
+          input: undefined,
+          output: "42",
         },
       ]),
     ];
@@ -106,12 +99,7 @@ describe("mergeDeltas", () => {
     expect(messages[0].tool).toBe(true);
     const content = messages[0].message?.content;
     expect(content).toEqual([
-      {
-        type: "tool-call",
-        toolCallId: "call1",
-        toolName: "myTool",
-        args: "",
-      },
+      { type: "tool-call", toolCallId: "call1", toolName: "myTool", args: "" },
     ]);
     expect(messages[1].message?.role).toBe("tool");
     expect(messages[1].tool).toBe(true);
@@ -145,8 +133,8 @@ describe("mergeDeltas", () => {
     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" }]),
+      makeDelta("s2", 0, 3, [{ type: "text-delta", id: "1", text: "B" }]),
+      makeDelta("s1", 0, 3, [{ type: "text-delta", id: "2", text: "A" }]),
     ];
     const [messages, _, changed] = mergeDeltas("thread1", [s2, s1], [], deltas);
     expect(messages).toHaveLength(2);
@@ -162,11 +150,13 @@ describe("mergeDeltas", () => {
     const streamId = "s4";
     const streamMessages = [makeStreamMessage(streamId, 4, 0)];
     const deltas = [
-      makeDelta(streamId, 0, 5, [{ type: "text-delta", textDelta: "Hello" }]),
+      makeDelta(streamId, 0, 5, [
+        { type: "text-delta", id: "1", text: "Hello" },
+      ]),
       makeDelta(streamId, 5, 11, [
-        { type: "text-delta", textDelta: " World!" },
+        { type: "text-delta", id: "2", text: " World!" },
       ]),
-      makeDelta(streamId, 11, 12, [{ type: "text-delta", textDelta: "!" }]),
+      makeDelta(streamId, 11, 12, [{ type: "text-delta", id: "3", text: "!" }]),
     ];
     const [messages] = mergeDeltas("thread1", streamMessages, [], deltas);
     expect(messages).toHaveLength(1);
@@ -185,10 +175,10 @@ describe("mergeDeltas", () => {
     const streamMessages = [makeStreamMessage(streamId, 6, 0)];
     const deltas = [
       makeDelta(streamId, 0, 1, [
-        { type: "reasoning", textDelta: "I'm thinking..." },
+        { type: "reasoning-delta", id: "1", text: "I'm thinking..." },
       ]),
       makeDelta(streamId, 1, 2, [
-        { type: "reasoning", textDelta: " Still thinking..." },
+        { type: "reasoning-delta", id: "2", text: " Still thinking..." },
       ]),
     ];
     const [messages] = mergeDeltas("thread1", streamMessages, [], deltas);
@@ -206,9 +196,11 @@ describe("mergeDeltas", () => {
     const streamId = "s7";
     const streamMessage = makeStreamMessage(streamId, 7, 0);
     const deltas = [
-      makeDelta(streamId, 0, 5, [{ type: "text-delta", textDelta: "Hello" }]),
+      makeDelta(streamId, 0, 5, [
+        { type: "text-delta", id: "1", text: "Hello" },
+      ]),
       makeDelta(streamId, 5, 11, [
-        { type: "text-delta", textDelta: " World!" },
+        { type: "text-delta", id: "2", text: " World!" },
       ]),
     ];
     // First call: apply both deltas
@@ -232,7 +224,7 @@ describe("mergeDeltas", () => {
     // Third call: add a new delta
     const moreDeltas = [
       ...deltas,
-      makeDelta(streamId, 11, 12, [{ type: "text-delta", textDelta: "!" }]),
+      makeDelta(streamId, 11, 12, [{ type: "text-delta", id: "3", text: "!" }]),
     ];
     [result, changed] = applyDeltasToStreamMessage(
       "thread1",
@@ -259,9 +251,11 @@ describe("mergeDeltas", () => {
     const streamId = "s8";
     const streamMessages = [makeStreamMessage(streamId, 8, 0)];
     const deltas = [
-      makeDelta(streamId, 0, 5, [{ type: "text-delta", textDelta: "Hello" }]),
+      makeDelta(streamId, 0, 5, [
+        { type: "text-delta", id: "1", text: "Hello" },
+      ]),
       makeDelta(streamId, 5, 11, [
-        { type: "text-delta", textDelta: " World!" },
+        { type: "text-delta", id: "2", text: " World!" },
       ]),
     ];
     // Deep freeze inputs to catch mutation
@@ -306,9 +300,11 @@ describe("mergeDeltas", () => {
     // 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, 0, 5, [
+        { type: "text-delta", id: "1", text: "Hello" },
+      ]),
       makeDelta(streamId, 5, 11, [
-        { type: "text-delta", textDelta: " World!" },
+        { type: "text-delta", id: "2", text: " World!" },
       ]),
     ]);
   });

package/src/react/deltas.ts

@@ -1,16 +1,16 @@
-import type { TextPart, ToolCallPart, ToolResultPart } from "ai";
+import type { TextStreamPart, ToolSet } from "ai";
 import type { MessageDoc } from "../client/index.js";
 import type {
   Message,
   MessageStatus,
   StreamDelta,
   StreamMessage,
-  TextStreamPart,
+  vReasoningPart,
+  vTextPart,
+  vToolCallPart,
+  vToolResultPart,
 } from "../validators.js";
-import type { UIMessage } from "./toUIMessages.js";
-import { toUIMessages } from "./toUIMessages.js";
-
-export { toUIMessages, type UIMessage };
+import type { Infer } from "convex/values";
 
 export function mergeDeltas(
   threadId: string,
@@ -73,7 +73,7 @@ export function applyDeltasToStreamMessage(
 ): [{ streamId: string; cursor: number; messages: MessageDoc[] }, boolean] {
   let changed = false;
   let cursor = existing?.cursor ?? 0;
-  let parts: TextStreamPart[] = [];
+  let parts: TextStreamPart<ToolSet>[] = [];
   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)}`);
@@ -142,10 +142,10 @@ export function applyDeltasToStreamMessage(
   let lastContent = getLastContent(currentMessage);
   for (const part of parts) {
     let contentToAdd:
-      | TextPart
-      | ToolCallPart
-      | { type: "reasoning"; text: string }
-      | ToolResultPart
+      | Infer<typeof vTextPart>
+      | Infer<typeof vToolCallPart>
+      | Infer<typeof vToolResultPart>
+      | Infer<typeof vReasoningPart>
       | undefined;
     const isToolRole = part.type === "source" || part.type === "tool-result";
     if (isToolRole !== (currentMessage.message!.role === "tool")) {
@@ -160,27 +160,34 @@ export function applyDeltasToStreamMessage(
       continue;
     }
     switch (part.type) {
-      case "text-delta":
-        currentMessage.text = (currentMessage.text ?? "") + part.textDelta;
+      case "text-delta": {
+        const text = part.text;
+        currentMessage.text = (currentMessage.text ?? "") + text;
         if (lastContent?.type === "text") {
-          lastContent.text = (lastContent.text ?? "") + part.textDelta;
+          lastContent.text = (lastContent.text ?? "") + text;
         } else {
-          contentToAdd = {
-            type: "text",
-            text: part.textDelta,
-          };
+          contentToAdd = { type: "text", text } satisfies Infer<
+            typeof vTextPart
+          >;
         }
         break;
-      case "tool-call-streaming-start":
+      }
+      case "tool-input-start": {
+        const toolCallId = part.id;
         currentMessage.tool = true;
         contentToAdd = {
           type: "tool-call",
-          toolCallId: part.toolCallId,
+          toolCallId,
           toolName: part.toolName,
           args: "",
-        };
+          providerExecuted:
+            "providerExecuted" in part ? part.providerExecuted : undefined,
+          providerOptions:
+            "providerMetadata" in part ? part.providerMetadata : undefined,
+        } satisfies Infer<typeof vToolCallPart>;
         break;
-      case "tool-call-delta":
+      }
+      case "tool-input-delta":
         {
           currentMessage.tool = true;
           if (lastContent?.type !== "tool-call") {
@@ -189,33 +196,43 @@ export function applyDeltasToStreamMessage(
           if (typeof lastContent.args !== "string") {
             throw new Error("Expected args to be a string");
           }
-          lastContent.args = (lastContent.args ?? "") + part.argsTextDelta;
+          const delta =
+            "argsTextDelta" in part ? part.argsTextDelta : part.delta;
+          lastContent.args = (lastContent.args ?? "") + delta;
         }
         break;
-      case "tool-call":
+      case "tool-call": {
         currentMessage.tool = true;
-        contentToAdd = part;
+        contentToAdd = toolCallContent(part);
+        break;
+      }
+      case "tool-result": {
+        contentToAdd = toolResultContent(part);
         break;
-      case "reasoning":
-        currentMessage.reasoning =
-          (currentMessage.reasoning ?? "") + part.textDelta;
+      }
+      case "reasoning-delta": {
+        currentMessage.reasoning = (currentMessage.reasoning ?? "") + part.text;
         if (lastContent?.type === "reasoning") {
-          lastContent.text = (lastContent.text ?? "") + part.textDelta;
+          lastContent.text = (lastContent.text ?? "") + part.text;
         } else {
           contentToAdd = {
             type: "reasoning",
-            text: part.textDelta,
-          };
+            text: part.text,
+            providerOptions:
+              "providerMetadata" in part ? part.providerMetadata : undefined,
+            state: "streaming",
+          } satisfies Infer<typeof vReasoningPart>;
         }
         break;
+      }
       case "source":
         if (!currentMessage.sources) {
           currentMessage.sources = [];
         }
-        currentMessage.sources.push(part.source);
+        currentMessage.sources.push(part);
         break;
-      case "tool-result":
-        contentToAdd = part;
+      case "raw":
+        // ignore
         break;
       default:
         console.warn(`Received unexpected part: ${JSON.stringify(part)}`);
@@ -236,6 +253,31 @@ export function applyDeltasToStreamMessage(
   return [newStream, true];
 }
 
+function toolCallContent(
+  part: Extract<TextStreamPart<ToolSet>, { type: "tool-call" }>,
+): Infer<typeof vToolCallPart> {
+  const args = "args" in part ? part.args : part.input;
+  return {
+    type: "tool-call",
+    toolCallId: part.toolCallId,
+    toolName: part.toolName,
+    args,
+    providerExecuted: part.providerExecuted,
+  } satisfies Infer<typeof vToolCallPart>;
+}
+
+function toolResultContent(
+  part: Extract<TextStreamPart<ToolSet>, { type: "tool-result" }>,
+): Infer<typeof vToolResultPart> {
+  return {
+    type: "tool-result",
+    toolCallId: part.toolCallId,
+    toolName: part.toolName,
+    result: part.output,
+    args: part.input,
+    providerExecuted: part.providerExecuted,
+  } satisfies Infer<typeof vToolResultPart>;
+}
 function cloneMessageAndContent(
   message: Message | undefined,
 ): Message | undefined {
@@ -272,10 +314,11 @@ function statusFromStreamStatus(
   }
 }
 
+// TODO: share more code with applyDeltasToStreamMessage
 export function createStreamingMessage(
   threadId: string,
   message: StreamMessage,
-  part: TextStreamPart,
+  part: TextStreamPart<ToolSet>,
   index: number,
 ): MessageDoc {
   const { streamId, ...rest } = message;
@@ -288,16 +331,15 @@ export function createStreamingMessage(
     tool: false,
   };
   switch (part.type) {
-    case "text-delta":
+    case "text-delta": {
+      const text = part.text || "";
       return {
         ...metadata,
-        message: {
-          role: "assistant",
-          content: [{ type: "text", text: part.textDelta }],
-        },
-        text: part.textDelta,
+        message: { role: "assistant", content: [{ type: "text", text }] },
+        text,
       };
-    case "tool-call-streaming-start":
+    }
+    case "tool-input-start": {
       return {
         ...metadata,
         tool: true,
@@ -307,58 +349,67 @@ export function createStreamingMessage(
             {
               type: "tool-call",
               toolName: part.toolName,
-              toolCallId: part.toolCallId,
+              toolCallId: part.id,
               args: "", // when it's a string, it's a partial call
+              providerExecuted:
+                "providerExecuted" in part ? part.providerExecuted : undefined,
+              providerOptions:
+                "providerMetadata" in part ? part.providerMetadata : undefined,
             },
           ],
         },
       };
-    case "reasoning":
+    }
+    case "tool-input-delta": {
+      console.warn("Received tool call delta part first??");
+      const delta = part.delta;
+      const toolCallId = part.id;
+      const toolName = part.type.slice("tool-".length);
       return {
         ...metadata,
+        tool: true,
         message: {
           role: "assistant",
-          content: [{ type: "reasoning", text: part.textDelta }],
+          content: [{ type: "tool-call", toolCallId, toolName, args: delta }],
         },
-        reasoning: part.textDelta,
       };
-    case "source":
-      console.warn("Received source part first??");
+    }
+    case "tool-call": {
       return {
         ...metadata,
         tool: true,
-        message: { role: "tool", content: [] },
-        sources: [part.source],
+        message: { role: "assistant", content: [toolCallContent(part)] },
       };
-    case "tool-call":
+    }
+    case "tool-result":
       return {
         ...metadata,
         tool: true,
-        message: { role: "assistant", content: [part] },
+        message: { role: "tool", content: [toolResultContent(part)] },
       };
-    case "tool-call-delta":
-      console.warn("Received tool call delta part first??");
+    case "reasoning-delta": {
       return {
         ...metadata,
-        tool: true,
         message: {
           role: "assistant",
-          content: [
-            {
-              type: "tool-call",
-              toolCallId: part.toolCallId,
-              toolName: part.toolName,
-              args: part.argsTextDelta,
-            },
-          ],
+          content: [{ type: "reasoning", text: part.text }],
         },
+        reasoning: part.text,
       };
-    case "tool-result":
+    }
+    case "source":
+      console.warn("Received source part first??");
       return {
         ...metadata,
         tool: true,
-        message: { role: "tool", content: [part] },
+        message: { role: "tool", content: [] },
+        sources: [part],
       };
+    // case "raw":
+    //   return {
+    //     ...metadata,
+    //     message: { role: "assistant", content: [part.rawValue] },
+    //   };
     default:
       throw new Error(`Unexpected part type: ${JSON.stringify(part)}`);
   }

package/src/react/index.ts

@@ -7,7 +7,7 @@ import {
 } from "convex/react";
 import { usePaginatedQuery } from "convex-helpers/react";
 import type { FunctionArgs } from "convex/server";
-import { useMemo, useRef, useState } from "react";
+import { useMemo, useState } from "react";
 import type { MessageDoc } from "../client/index.js";
 import type { SyncStreamsReturnValue } from "../client/types.js";
 import type { StreamArgs } from "../validators.js";
@@ -23,6 +23,7 @@ import type {
 
 export { optimisticallySendMessage } from "./optimisticallySendMessage.js";
 export { useSmoothText } from "./useSmoothText.js";
+export { SmoothText } from "./SmoothText.js";
 export { toUIMessages, type UIMessage, type ThreadQuery };
 
 /**
@@ -97,6 +98,15 @@ export function useThreadMessages<
     { initialNumItems: options.initialNumItems },
   );
 
+  let startOrder = paginated.results.at(-1)?.order ?? 0;
+  for (let i = paginated.results.length - 1; i >= 0; i--) {
+    const m = paginated.results[i];
+    if (!m.streaming && m.status === "pending") {
+      // round down to the nearest 10 for some cache benefits
+      startOrder = m.order - (m.order % 10);
+      break;
+    }
+  }
   // These are streaming messages that will not include full messages.
   const streamMessages = useStreamingThreadMessages(
     query as ThreadStreamQuery<
@@ -107,7 +117,7 @@ export function useThreadMessages<
       args === "skip" ||
       paginated.status === "LoadingFirstPage"
       ? "skip"
-      : { ...args, startOrder: paginated.results.at(-1)?.order },
+      : { ...args, startOrder },
   );
 
   const merged = useMemo(() => {
@@ -120,16 +130,32 @@ export function useThreadMessages<
       ...paginated,
       results: paginated.results
         .map((m) => ({ ...m, streaming: false }))
+        // Note: this is intentionally after paginated results.
         .concat(streamListMessages)
         .sort((a, b) =>
           a.order === b.order ? a.stepOrder - b.stepOrder : a.order - b.order,
         )
-        // They shouldn't overlap, but check for duplicates just in case.
-        .filter(
-          (m, i, arr) =>
-            !arr[i - 1] ||
-            m.order !== arr[i - 1].order ||
-            m.stepOrder !== arr[i - 1].stepOrder,
+        .reduce(
+          (msgs, msg) => {
+            const last = msgs.at(-1);
+            if (!last) {
+              return [msg];
+            }
+            if (last.order !== msg.order || last.stepOrder !== msg.stepOrder) {
+              return [...msgs, msg];
+            }
+            if (
+              last.status === "pending" &&
+              (msg.streaming || msg.status !== "pending")
+            ) {
+              // Let's prefer a streaming or finalized message over a pending
+              // one.
+              return [...msgs.slice(0, -1), msg];
+            }
+            // skip the new one if the previous one (listed) was finalized
+            return msgs;
+          },
+          [] as (ThreadMessagesResult<Query> & { streaming: boolean })[],
         ),
     };
   }, [paginated, streamMessages]);
@@ -162,11 +188,7 @@ export function useStreamingThreadMessages<
   const [streams, setStreams] = useState<
     Array<{ streamId: string; cursor: number; messages: MessageDoc[] }>
   >([]);
-  const startOrderRef = useRef<number>(0);
   const queryArgs = args === "skip" ? args : omit(args, ["startOrder"]);
-  if (args !== "skip" && !startOrderRef.current && args.startOrder) {
-    startOrderRef.current = args.startOrder;
-  }
   // Get all the active streams
   const streamList = useQuery(
     query,
@@ -177,7 +199,7 @@ export function useStreamingThreadMessages<
           paginationOpts: { cursor: null, numItems: 0 },
           streamArgs: {
             kind: "list",
-            startOrder: startOrderRef.current,
+            startOrder: queryArgs.startOrder ?? 0,
           } as StreamArgs,
         } as FunctionArgs<Query>),
   ) as

package/src/react/optimisticallySendMessage.ts

@@ -12,17 +12,14 @@ export function optimisticallySendMessage(
   return (store, args) => {
     const queries = store.getAllQueries(query);
     let maxOrder = -1;
-    let maxStepOrder = 0;
     for (const q of queries) {
       if (q.args?.threadId !== args.threadId) continue;
       if (q.args.streamArgs) continue;
       for (const m of q.value?.page ?? []) {
         maxOrder = Math.max(maxOrder, m.order);
-        maxStepOrder = Math.max(maxStepOrder, m.stepOrder);
       }
     }
     const order = maxOrder + 1;
-    const stepOrder = 0;
     insertAtTop({
       paginatedQuery: query,
       argsToMatch: { threadId: args.threadId, streamArgs: undefined },
@@ -30,14 +27,11 @@ export function optimisticallySendMessage(
         _creationTime: Date.now(),
         _id: randomUUID(),
         order,
-        stepOrder,
+        stepOrder: 0,
         status: "pending",
         threadId: args.threadId,
         tool: false,
-        message: {
-          role: "user",
-          content: args.prompt,
-        },
+        message: { role: "user", content: args.prompt },
         text: args.prompt,
       },
       localQueryStore: store,