@convex-dev/agent 0.2.7 → 0.2.8-alpha.0

package/src/client/streaming.ts

@@ -1,4 +1,10 @@
-import { type ChunkDetector } from "ai";
+import {
+  type ChunkDetector,
+  type AsyncIterableStream,
+  type StreamTextTransform,
+  type ToolSet,
+  smoothStream,
+} from "ai";
 import {
   vStreamDelta,
   vStreamMessage,
@@ -152,6 +158,36 @@ export const DEFAULT_STREAMING_OPTIONS = {
   returnImmediately: false,
 } satisfies StreamingOptions;
 
+/**
+ *
+ * @param options The options passed to `agent.streamText` to decide whether to
+ * save deltas while streaming.
+ * @param existing The transforms passed to `agent.streamText` to merge with.
+ * @returns The merged transforms to pass to the underlying `streamText` call.
+ */
+export function mergeTransforms<TOOLS extends ToolSet>(
+  options: { chunking?: StreamingOptions["chunking"] } | boolean | undefined,
+  existing:
+    | StreamTextTransform<TOOLS>
+    | Array<StreamTextTransform<TOOLS>>
+    | undefined,
+) {
+  if (!options) {
+    return existing;
+  }
+  const chunking =
+    typeof options === "boolean"
+      ? DEFAULT_STREAMING_OPTIONS.chunking
+      : options.chunking;
+  const transforms = Array.isArray(existing)
+    ? existing
+    : existing
+      ? [existing]
+      : [];
+  transforms.push(smoothStream({ delayInMs: null, chunking }));
+  return transforms;
+}
+
 /**
  * DeltaStreamer can be used to save a stream of "parts" by writing
  * batches of them in "deltas" to the database so clients can subscribe
@@ -162,7 +198,7 @@ export const DEFAULT_STREAMING_OPTIONS = {
 export class DeltaStreamer<T> {
   public streamId: string | undefined;
   public readonly config: {
-    stream: Required<StreamingOptions>;
+    throttleMs: number;
     onAsyncAbort: (reason: string) => Promise<void>;
     compress: ((parts: T[]) => T[]) | null;
   };
@@ -176,7 +212,7 @@ export class DeltaStreamer<T> {
     public readonly component: AgentComponent,
     public readonly ctx: RunMutationCtx,
     config: {
-      stream: true | StreamingOptions;
+      throttleMs: number | undefined;
       onAsyncAbort: (reason: string) => Promise<void>;
       abortSignal: AbortSignal | undefined;
       compress: ((parts: T[]) => T[]) | null;
@@ -190,13 +226,11 @@ export class DeltaStreamer<T> {
       model?: string;
       provider?: string;
       providerOptions?: ProviderOptions;
+      format: "UIMessageChunk" | "TextStreamPart" | undefined;
     },
   ) {
     this.config = {
-      stream:
-        config.stream === true
-          ? DEFAULT_STREAMING_OPTIONS
-          : { ...DEFAULT_STREAMING_OPTIONS, ...config },
+      throttleMs: config.throttleMs ?? DEFAULT_STREAMING_OPTIONS.throttleMs,
       onAsyncAbort: config.onAsyncAbort,
       compress: config.compress,
     };
@@ -232,12 +266,19 @@ export class DeltaStreamer<T> {
     this.#nextParts.push(...parts);
     if (
       !this.#ongoingWrite &&
-      Date.now() - this.#latestWrite >= this.config.stream.throttleMs
+      Date.now() - this.#latestWrite >= this.config.throttleMs
     ) {
       this.#ongoingWrite = this.#sendDelta();
     }
   }
 
+  public async consumeStream(stream: AsyncIterableStream<T>) {
+    for await (const chunk of stream) {
+      await this.addParts([chunk]);
+    }
+    await this.finish();
+  }
+
   async #sendDelta() {
     if (this.abortController.signal.aborted) {
       return;
@@ -267,7 +308,7 @@ export class DeltaStreamer<T> {
     // Now that we've sent the delta, check if we need to send another one.
     if (
       this.#nextParts.length > 0 &&
-      Date.now() - this.#latestWrite >= this.config.stream.throttleMs
+      Date.now() - this.#latestWrite >= this.config.throttleMs
     ) {
       // We send again immediately with the accumulated deltas.
       this.#ongoingWrite = this.#sendDelta();

package/src/component/_generated/api.d.ts

@@ -2189,6 +2189,7 @@ export type Mounts = {
       "public",
       {
         agentName?: string;
+        format?: "UIMessageChunk" | "TextStreamPart";
         model?: string;
         order: number;
         provider?: string;
@@ -2253,6 +2254,7 @@ export type Mounts = {
       },
       Array<{
         agentName?: string;
+        format?: "UIMessageChunk" | "TextStreamPart";
         model?: string;
         order: number;
         provider?: string;

package/src/component/messages.ts

@@ -344,7 +344,7 @@ export const finalizeMessage = mutation({
     const message = await ctx.db.get(messageId);
     assert(message, `Message ${messageId} not found`);
     if (message.status !== "pending") {
-      console.log(
+      console.debug(
         "Trying to finalize a message that's already",
         message.status,
       );

package/src/component/schema.ts

@@ -91,6 +91,12 @@ export const schema = defineSchema({
     model: v.optional(v.string()),
     provider: v.optional(v.string()),
     providerOptions: v.optional(vProviderOptions), // Sent to model
+    // The data format for the deltas. By default, we use UIMessageChunks.
+    // This format dictates how the messages are materialized for capturing
+    // partial messages during failure, as well as on the client side.
+    format: v.optional(
+      v.union(v.literal("UIMessageChunk"), v.literal("TextStreamPart")),
+    ),
 
     threadId: v.id("threads"),
     order: v.number(),

package/src/component/streams.ts

@@ -20,8 +20,8 @@ import { stream } from "convex-helpers/server/stream";
 import { mergedStream } from "convex-helpers/server/stream";
 import { paginator } from "convex-helpers/server/pagination";
 import type { WithoutSystemFields } from "convex/server";
-import { mergeDeltas } from "../react/deltas.js";
-import { serializeOrThrow } from "../mapping.js";
+import { deriveUIMessagesFromDeltas } from "../deltas.js";
+import { fromUIMessages } from "../UIMessages.js";
 
 const SECOND = 1000;
 const MINUTE = 60 * SECOND;
@@ -140,6 +140,7 @@ export const list = query({
       streamId: m._id,
       status: m.state.kind,
       ...pick(m, [
+        "format",
         "order",
         "stepOrder",
         "userId",
@@ -516,72 +517,54 @@ export async function getStreamingMessagesWithMetadata(
   );
   const messages = (
     await Promise.all(
-      streamingMessages.map((m) =>
-        getMessagesWithMetadataForStreamingMessage(
-          ctx,
+      streamingMessages.map(async (streamingMessage) => {
+        const deltas = await ctx.db
+          .query("streamDeltas")
+          .withIndex("streamId_start_end", (q) =>
+            q.eq("streamId", streamingMessage._id),
+          )
+          .take(1000);
+        const uiMessages = await deriveUIMessagesFromDeltas(
           threadId,
-          stepOrder,
-          m,
-          metadata,
-        ),
-      ),
+          [
+            {
+              ...streamingMessage,
+              status: "streaming",
+              streamId: streamingMessage._id,
+            },
+          ],
+          deltas,
+        );
+        // We don't save messages that have already been saved
+        const numToSkip = stepOrder - streamingMessage.stepOrder;
+        const messages = await Promise.all(
+          fromUIMessages(uiMessages, streamingMessage)
+            .slice(numToSkip)
+            .filter((m) => m.message !== undefined)
+            .map(async (msg) => {
+              return {
+                ...pick(msg, [
+                  "message",
+                  "fileIds",
+                  "status",
+                  "finishReason",
+                  "model",
+                  "provider",
+                  "providerMetadata",
+                  "sources",
+                  "reasoning",
+                  "reasoningDetails",
+                  "usage",
+                  "warnings",
+                  "error",
+                ]),
+                ...metadata,
+              } as MessageWithMetadataInternal;
+            }),
+        );
+        return messages;
+      }),
     )
   ).flat();
   return messages;
 }
-
-export async function getMessagesWithMetadataForStreamingMessage(
-  ctx: MutationCtx,
-  threadId: Id<"threads">,
-  stepOrder: number,
-  streamingMessage: Doc<"streamingMessages">,
-  metadata: { status: "success" | "failed"; error?: string },
-): Promise<MessageWithMetadataInternal[]> {
-  const deltas = await ctx.db
-    .query("streamDeltas")
-    .withIndex("streamId_start_end", (q) =>
-      q.eq("streamId", streamingMessage._id),
-    )
-    .take(1000);
-  const [messageDocs] = mergeDeltas(
-    threadId,
-    [
-      {
-        ...streamingMessage,
-        status: "streaming",
-        streamId: streamingMessage._id,
-      },
-    ],
-    [],
-    deltas,
-  );
-  // We don't save messages that have already been saved
-  const numToSkip = stepOrder - streamingMessage.stepOrder;
-  const messages = await Promise.all(
-    messageDocs
-      .slice(numToSkip)
-      .filter((m) => m.message !== undefined)
-      .map(async (msg) => {
-        const message = await serializeOrThrow(msg.message!);
-        return {
-          message,
-          ...pick(msg, [
-            "fileIds",
-            "status",
-            "finishReason",
-            "model",
-            "provider",
-            "providerMetadata",
-            "sources",
-            "reasoning",
-            "reasoningDetails",
-            "usage",
-            "warnings",
-            "error",
-          ]),
-          ...metadata,
-        } as MessageWithMetadataInternal;
-      }),
-  );
-  return messages;
-}

package/src/{react/deltas.test.ts → deltas.test.ts}

@@ -1,8 +1,11 @@
 import { describe, it, expect } from "vitest";
-import { mergeDeltas, applyDeltasToStreamMessage } from "./deltas.js";
-import type { StreamMessage, StreamDelta } from "../validators.js";
+import {
+  deriveUIMessagesFromTextStreamParts,
+  updateFromTextStreamParts,
+} from "./deltas.js";
+import type { StreamMessage, StreamDelta } from "./validators.js";
 import { omit } from "convex-helpers";
-import type { TextStreamPart, ToolSet } from "ai";
+import type { TextStreamPart, ToolSet, ToolUIPart } from "ai";
 
 function makeStreamMessage(
   streamId: string,
@@ -30,7 +33,7 @@ describe("mergeDeltas", () => {
         { type: "text-delta", id: "1", text: "Hello" },
       ]),
     ];
-    const [messages, newStreams, changed] = mergeDeltas(
+    const [messages, newStreams, changed] = deriveUIMessagesFromTextStreamParts(
       "thread1",
       streamMessages,
       [],
@@ -38,7 +41,7 @@ describe("mergeDeltas", () => {
     );
     expect(messages).toHaveLength(1);
     expect(messages[0].text).toBe("Hello");
-    expect(messages[0].message?.role).toBe("assistant");
+    expect(messages[0].role).toBe("assistant");
     expect(changed).toBe(true);
     expect(newStreams[0].cursor).toBe(5);
   });
@@ -54,7 +57,7 @@ describe("mergeDeltas", () => {
         { type: "text-delta", id: "2", text: " World!" },
       ]),
     ];
-    const [messages, newStreams, changed] = mergeDeltas(
+    const [messages, newStreams, changed] = deriveUIMessagesFromTextStreamParts(
       "thread1",
       streamMessages,
       [],
@@ -75,7 +78,7 @@ describe("mergeDeltas", () => {
           type: "tool-call",
           toolCallId: "call1",
           toolName: "myTool",
-          input: "",
+          input: "What's the meaning of life?",
         },
       ]),
       makeDelta(streamId, 1, 2, [
@@ -88,28 +91,23 @@ describe("mergeDeltas", () => {
         },
       ]),
     ];
-    const [messages, _, changed] = mergeDeltas(
+    const [[message], _, changed] = deriveUIMessagesFromTextStreamParts(
       "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(message).toBeDefined();
+    expect(message.role).toBe("assistant");
+    const content = message.parts;
     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",
+        type: "tool-myTool",
         toolCallId: "call1",
-        toolName: "myTool",
-        result: "42",
-      },
+        input: "What's the meaning of life?",
+        output: "42",
+        state: "output-available",
+      } satisfies ToolUIPart,
     ]);
     expect(changed).toBe(true);
   });
@@ -118,13 +116,12 @@ describe("mergeDeltas", () => {
     const streamId = "s3";
     const streamMessages = [makeStreamMessage(streamId, 3, 0)];
     const deltas: StreamDelta[] = [];
-    const [messages, newStreams, changed] = mergeDeltas(
+    const [, newStreams, changed] = deriveUIMessagesFromTextStreamParts(
       "thread1",
       streamMessages,
       [],
       deltas,
     );
-    expect(messages).toHaveLength(0);
     expect(changed).toBe(false);
     expect(newStreams[0].cursor).toBe(0);
   });
@@ -136,7 +133,12 @@ describe("mergeDeltas", () => {
       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);
+    const [messages, _, changed] = deriveUIMessagesFromTextStreamParts(
+      "thread1",
+      [s2, s1],
+      [],
+      deltas,
+    );
     expect(messages).toHaveLength(2);
     expect(messages[0].text).toBe("A");
     expect(messages[1].text).toBe("B");
@@ -158,11 +160,16 @@ describe("mergeDeltas", () => {
       ]),
       makeDelta(streamId, 11, 12, [{ type: "text-delta", id: "3", text: "!" }]),
     ];
-    const [messages] = mergeDeltas("thread1", streamMessages, [], deltas);
+    const [messages] = deriveUIMessagesFromTextStreamParts(
+      "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;
+    const content = messages[0].parts;
     if (Array.isArray(content)) {
       const textParts = content.filter((p) => p.type === "text");
       expect(textParts).toHaveLength(1);
@@ -175,20 +182,28 @@ describe("mergeDeltas", () => {
     const streamMessages = [makeStreamMessage(streamId, 6, 0)];
     const deltas = [
       makeDelta(streamId, 0, 1, [
+        { type: "reasoning-start", id: "1" },
         { type: "reasoning-delta", id: "1", text: "I'm thinking..." },
       ]),
       makeDelta(streamId, 1, 2, [
-        { type: "reasoning-delta", id: "2", text: " Still thinking..." },
+        { type: "reasoning-delta", id: "1", text: " Still thinking..." },
       ]),
+      makeDelta(streamId, 2, 3, [{ type: "reasoning-end", id: "1" }]),
     ];
-    const [messages] = mergeDeltas("thread1", streamMessages, [], deltas);
+    const [messages] = deriveUIMessagesFromTextStreamParts(
+      "thread1",
+      streamMessages,
+      [],
+      deltas,
+    );
     expect(messages).toHaveLength(1);
-    if (Array.isArray(messages[0].message?.content)) {
-      const reasoningParts = messages[0].message.content.filter(
+    if (Array.isArray(messages[0].parts)) {
+      const reasoningParts = messages[0].parts.filter(
         (p) => p.type === "reasoning",
       );
       expect(reasoningParts).toHaveLength(1);
       expect(reasoningParts[0].text).toBe("I'm thinking... Still thinking...");
+      expect(reasoningParts[0].state).toBe("done");
     }
   });
 
@@ -204,47 +219,43 @@ describe("mergeDeltas", () => {
       ]),
     ];
     // First call: apply both deltas
-    let [result, changed] = applyDeltasToStreamMessage(
+    let [result, changed] = updateFromTextStreamParts(
       "thread1",
       streamMessage,
       undefined,
       deltas,
     );
-    expect(result.messages).toHaveLength(1);
-    expect(result.messages[0].text).toBe("Hello World!");
+    expect(result.message.text).toBe("Hello World!");
     // Second call: re-apply the same deltas (should not duplicate)
-    [result, changed] = applyDeltasToStreamMessage(
+    [result, changed] = updateFromTextStreamParts(
       "thread1",
       streamMessage,
       result,
       deltas,
     );
-    expect(result.messages).toHaveLength(1);
-    expect(result.messages[0].text).toBe("Hello World!");
+    expect(result.message.text).toBe("Hello World!");
     // Third call: add a new delta
     const moreDeltas = [
       ...deltas,
       makeDelta(streamId, 11, 12, [{ type: "text-delta", id: "3", text: "!" }]),
     ];
-    [result, changed] = applyDeltasToStreamMessage(
+    [result, changed] = updateFromTextStreamParts(
       "thread1",
       streamMessage,
       result,
       moreDeltas,
     );
     expect(changed).toBe(true);
-    expect(result.messages).toHaveLength(1);
-    expect(result.messages[0].text).toBe("Hello World!!");
+    expect(result.message.text).toBe("Hello World!!");
     // Re-apply all deltas again (should still not duplicate)
-    [result, changed] = applyDeltasToStreamMessage(
+    [result, changed] = updateFromTextStreamParts(
       "thread1",
       streamMessage,
       result,
       moreDeltas,
     );
     expect(changed).toBe(false);
-    expect(result.messages).toHaveLength(1);
-    expect(result.messages[0].text).toBe("Hello World!!");
+    expect(result.message.text).toBe("Hello World!!");
   });
 
   it("mergeDeltas is pure and does not mutate inputs", () => {
@@ -270,13 +281,13 @@ describe("mergeDeltas", () => {
     }
     deepFreeze(streamMessages);
     deepFreeze(deltas);
-    const [messages1, streams1, changed1] = mergeDeltas(
+    const [messages1, streams1, changed1] = deriveUIMessagesFromTextStreamParts(
       "thread1",
       streamMessages,
       [],
       deltas,
     );
-    const [messages2, streams2, changed2] = mergeDeltas(
+    const [messages2, streams2, changed2] = deriveUIMessagesFromTextStreamParts(
       "thread1",
       streamMessages,
       [],
@@ -288,12 +299,12 @@ describe("mergeDeltas", () => {
     expect(
       streams1.map((s) => ({
         ...s,
-        messages: s.messages.map((m) => omit(m, ["_creationTime"])),
+        message: omit(s.message, ["_creationTime"]),
       })),
     ).toEqual(
       streams2.map((s) => ({
         ...s,
-        messages: s.messages.map((m) => omit(m, ["_creationTime"])),
+        message: omit(s.message, ["_creationTime"]),
       })),
     );
     expect(changed1).toBe(changed2);