@convex-dev/agent 0.2.0-alpha.2 → 0.2.1

package/src/component/streams.ts

@@ -1,6 +1,7 @@
 import { omit, pick } from "convex-helpers";
 import { v } from "convex/values";
 import {
+  type MessageWithMetadataInternal,
   type StreamDelta,
   vStreamDelta,
   vStreamMessage,
@@ -18,6 +19,9 @@ import schema from "./schema.js";
 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";
 
 const SECOND = 1000;
 const MINUTE = 60 * SECOND;
@@ -33,13 +37,16 @@ export const addDelta = mutation({
   args: deltaValidator,
   returns: v.boolean(),
   handler: async (ctx, args) => {
-    await ctx.db.insert("streamDeltas", args);
-    await heartbeatStream(ctx, { streamId: args.streamId });
     const stream = await ctx.db.get(args.streamId);
-    if (stream?.state.kind !== "streaming") {
-      console.warn(`Stream is not streaming: ${args.streamId}`);
+    if (!stream) {
+      console.warn("Stream not found", args.streamId);
+      return false;
+    }
+    if (stream.state.kind !== "streaming") {
       return false;
     }
+    await ctx.db.insert("streamDeltas", args);
+    await heartbeatStream(ctx, { streamId: args.streamId });
     return true;
   },
 });
@@ -48,10 +55,7 @@ export const listDeltas = query({
   args: {
     threadId: v.id("threads"),
     cursors: v.array(
-      v.object({
-        streamId: v.id("streamingMessages"),
-        cursor: v.number(),
-      }),
+      v.object({ streamId: v.id("streamingMessages"), cursor: v.number() }),
     ),
   },
   returns: v.array(vStreamDelta),
@@ -85,10 +89,7 @@ export const create = mutation({
   args: omit(schema.tables.streamingMessages.validator.fields, ["state"]),
   returns: v.id("streamingMessages"),
   handler: async (ctx, args) => {
-    const state = {
-      kind: "streaming" as const,
-      lastHeartbeat: Date.now(),
-    };
+    const state = { kind: "streaming" as const, lastHeartbeat: Date.now() };
     const streamId = await ctx.db.insert("streamingMessages", {
       ...args,
       state,
@@ -152,11 +153,7 @@ export const list = query({
 });
 
 export const abortByOrder = mutation({
-  args: {
-    threadId: v.id("threads"),
-    order: v.number(),
-    reason: v.string(),
-  },
+  args: { threadId: v.id("threads"), order: v.number(), reason: v.string() },
   returns: v.boolean(),
   handler: async (ctx, args) => {
     const streams = await ctx.db
@@ -169,10 +166,7 @@ export const abortByOrder = mutation({
       )
       .take(100);
     for (const stream of streams) {
-      await abortById(ctx, {
-        streamId: stream._id,
-        reason: args.reason,
-      });
+      await abortById(ctx, { streamId: stream._id, reason: args.reason });
     }
     return streams.length > 0;
   },
@@ -182,6 +176,7 @@ export const abort = mutation({
   args: {
     streamId: v.id("streamingMessages"),
     reason: v.string(),
+    finalDelta: v.optional(deltaValidator),
   },
   returns: v.boolean(),
   handler: abortById,
@@ -189,16 +184,20 @@ export const abort = mutation({
 
 async function abortById(
   ctx: MutationCtx,
-  args: { streamId: Id<"streamingMessages">; reason: string },
+  args: {
+    streamId: Id<"streamingMessages">;
+    reason: string;
+    finalDelta?: WithoutSystemFields<Doc<"streamDeltas">>;
+  },
 ) {
   const stream = await ctx.db.get(args.streamId);
   if (!stream) {
     throw new Error(`Stream not found: ${args.streamId}`);
   }
+  if (args.finalDelta) {
+    await ctx.db.insert("streamDeltas", args.finalDelta);
+  }
   if (stream.state.kind !== "streaming") {
-    console.warn(
-      `Stream trying to abort but not currently streaming (${stream.state.kind}): ${args.streamId}`,
-    );
     return false;
   }
   await cleanupTimeoutFn(ctx, stream);
@@ -220,49 +219,65 @@ async function cleanupTimeoutFn(
   }
 }
 
+// No longer used from the DeltaStreamer
 export const finish = mutation({
   args: {
     streamId: v.id("streamingMessages"),
     finalDelta: v.optional(deltaValidator),
   },
   returns: v.null(),
-  handler: async (ctx, args) => {
-    if (args.finalDelta) {
-      await ctx.db.insert("streamDeltas", args.finalDelta);
-    }
-    const stream = await ctx.db.get(args.streamId);
-    if (!stream) {
-      throw new Error(`Stream not found: ${args.streamId}`);
-    }
-    if (stream.state.kind !== "streaming") {
-      console.warn(
-        `Stream trying to finish but not currently streaming: ${args.streamId}`,
-      );
-      return;
-    }
-    await cleanupTimeoutFn(ctx, stream);
-    const cleanupFnId = await ctx.scheduler.runAfter(
-      DELETE_STREAM_DELAY,
-      api.streams.deleteStreamAsync,
-      { streamId: args.streamId },
-    );
-    await ctx.db.patch(args.streamId, {
-      state: { kind: "finished", endedAt: Date.now(), cleanupFnId },
-    });
+  handler: finishHandler,
+});
+
+export async function finishHandler(
+  ctx: MutationCtx,
+  args: {
+    streamId: Id<"streamingMessages">;
+    finalDelta?: WithoutSystemFields<Doc<"streamDeltas">>;
   },
+) {
+  if (args.finalDelta) {
+    await ctx.db.insert("streamDeltas", args.finalDelta);
+  }
+  const stream = await ctx.db.get(args.streamId);
+  if (!stream) {
+    throw new Error(`Stream not found: ${args.streamId}`);
+  }
+  if (stream.state.kind !== "streaming") {
+    console.warn(
+      `Stream trying to finish ${args.streamId} but is ${stream.state.kind}`,
+    );
+    return;
+  }
+  await cleanupTimeoutFn(ctx, stream);
+  const cleanupFnId = await ctx.scheduler.runAfter(
+    DELETE_STREAM_DELAY,
+    api.streams.deleteStreamAsync,
+    { streamId: args.streamId },
+  );
+  await ctx.db.patch(args.streamId, {
+    state: { kind: "finished", endedAt: Date.now(), cleanupFnId },
+  });
+}
+
+// TODO: use this heartbeat while streaming, every 30 seconds or so,
+// then reduce the timeout to 60 seconds.
+export const heartbeat = mutation({
+  args: { streamId: v.id("streamingMessages") },
+  returns: v.null(),
+  handler: heartbeatStream,
 });
 
 async function heartbeatStream(
   ctx: MutationCtx,
   args: { streamId: Id<"streamingMessages"> },
-) {
+): Promise<void> {
   const stream = await ctx.db.get(args.streamId);
   if (!stream) {
     console.warn("Stream not found", args.streamId);
     return;
   }
   if (stream.state.kind !== "streaming") {
-    console.warn("Stream is not streaming", args.streamId);
     return;
   }
   if (Date.now() - stream.state.lastHeartbeat < TIMEOUT_INTERVAL / 4) {
@@ -286,11 +301,7 @@ async function heartbeatStream(
     { streamId: args.streamId },
   );
   await ctx.db.patch(args.streamId, {
-    state: {
-      kind: "streaming",
-      lastHeartbeat: Date.now(),
-      timeoutFnId,
-    },
+    state: { kind: "streaming", lastHeartbeat: Date.now(), timeoutFnId },
   });
 }
 
@@ -304,10 +315,7 @@ export const timeoutStream = internalMutation({
       return;
     }
     await ctx.db.patch(args.streamId, {
-      state: {
-        kind: "aborted",
-        reason: "timeout",
-      },
+      state: { kind: "aborted", reason: "timeout" },
     });
   },
 });
@@ -362,11 +370,7 @@ export async function deleteStreamsPageForThreadId(
     "stepOrder",
   ]).first();
   if (!streamMessage) {
-    return {
-      isDone: true,
-      streamOrder: undefined,
-      deltaCursor: undefined,
-    };
+    return { isDone: true, streamOrder: undefined, deltaCursor: undefined };
   }
   const result = await deletePageForStreamId(ctx, {
     streamId: streamMessage._id,
@@ -375,11 +379,7 @@ export async function deleteStreamsPageForThreadId(
   if (result.isDone) {
     deltaCursor = undefined;
   }
-  return {
-    isDone: false,
-    streamOrder: streamMessage.order,
-    deltaCursor,
-  };
+  return { isDone: false, streamOrder: streamMessage.order, deltaCursor };
 }
 
 export const deleteStreamsPageForThreadIdMutation = internalMutation({
@@ -474,3 +474,114 @@ export const deleteAllStreamsForThreadIdSync = action({
     }
   },
 });
+
+export async function getStreamingMessages(
+  ctx: MutationCtx,
+  threadId: Id<"threads">,
+  order: number,
+  stepOrder: number,
+): Promise<Doc<"streamingMessages">[]> {
+  return mergedStream(
+    (["aborted", "streaming", "finished"] as const).map((state) =>
+      stream(ctx.db, schema)
+        .query("streamingMessages")
+        .withIndex("threadId_state_order_stepOrder", (q) =>
+          q
+            .eq("threadId", threadId)
+            .eq("state.kind", state)
+            .eq("order", order)
+            .lte("stepOrder", stepOrder),
+        )
+        .order("desc"),
+    ),
+    ["stepOrder"],
+  ).take(10);
+}
+
+export async function getStreamingMessagesWithMetadata(
+  ctx: MutationCtx,
+  {
+    threadId,
+    order,
+    stepOrder,
+  }: { threadId: Id<"threads">; order: number; stepOrder: number },
+  metadata: { status: "success" | "failed"; error?: string },
+): Promise<MessageWithMetadataInternal[]> {
+  // See if there are any streaming messages for this order
+  const streamingMessages = await getStreamingMessages(
+    ctx,
+    threadId,
+    order,
+    stepOrder,
+  );
+  const messages = (
+    await Promise.all(
+      streamingMessages.map((m) =>
+        getMessagesWithMetadataForStreamingMessage(
+          ctx,
+          threadId,
+          stepOrder,
+          m,
+          metadata,
+        ),
+      ),
+    )
+  ).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/mapping.ts

@@ -76,6 +76,25 @@ export async function serializeMessage(
   };
 }
 
+export async function serializeOrThrow(
+  message: ModelMessage | Message,
+): Promise<SerializedMessage> {
+  const { content } = await serializeContent(
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    {} as any,
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    {} as any,
+    message.content,
+  );
+  return {
+    role: message.role,
+    content,
+    ...(message.providerOptions
+      ? { providerOptions: message.providerOptions }
+      : {}),
+  } as SerializedMessage;
+}
+
 export function deserializeMessage(message: SerializedMessage): ModelMessage {
   return {
     ...message,
@@ -128,7 +147,7 @@ export async function serializeNewMessagesInStep<TOOLS extends ToolSet>(
   component: AgentComponent,
   step: StepResult<TOOLS>,
   metadata: { model: string; provider: string },
-): Promise<MessageWithMetadata[]> {
+): Promise<{ messages: MessageWithMetadata[] }> {
   // If there are tool results, there's another message with the tool results
   // ref: https://github.com/vercel/ai/blob/main/packages/ai/core/generate-text/to-response-messages.ts
   const assistantFields = {
@@ -158,7 +177,8 @@ export async function serializeNewMessagesInStep<TOOLS extends ToolSet>(
       };
     }),
   );
-  return messages;
+  // TODO: capture step.files separately?
+  return { messages };
 }
 
 export async function serializeObjectResult(
@@ -288,6 +308,7 @@ export async function serializeContent(
             providerOptions: part.providerOptions,
           } satisfies Infer<typeof vReasoningPart>;
         }
+        // Not in current generation output, but could be in historical messages
         case "redacted-reasoning": {
           return {
             type: part.type,

package/src/parts.ts

@@ -0,0 +1,39 @@
+import type { TextStreamPart, ToolSet } from "ai";
+
+export function serializeTextStreamingPartsV5(
+  parts: TextStreamPart<ToolSet>[],
+): TextStreamPart<ToolSet>[] {
+  const compressed: TextStreamPart<ToolSet>[] = [];
+  for (const part of parts) {
+    const last = compressed.at(-1);
+    if (part.type === "text-delta" && last?.type === "text-delta") {
+      last.text += part.text;
+    } else if (
+      part.type === "reasoning-delta" &&
+      last?.type === "reasoning-delta"
+    ) {
+      last.text += part.text;
+    } else {
+      if (
+        part.type === "start-step" ||
+        part.type === "finish-step" ||
+        part.type === "start" ||
+        part.type === "finish"
+      ) {
+        continue;
+      }
+      if (part.type === "file") {
+        compressed.push({
+          type: "file",
+          file: {
+            mediaType: part.file.mediaType,
+            base64: part.file.base64,
+            uint8Array: new Uint8Array([]),
+          },
+        });
+      }
+      compressed.push(part);
+    }
+  }
+  return compressed;
+}

package/src/react/SmoothText.tsx

@@ -0,0 +1,9 @@
+import { useSmoothText, type SmoothTextOptions } from "./useSmoothText.js";
+
+export function SmoothText({
+  text,
+  ...options
+}: { text: string } & SmoothTextOptions) {
+  const [visibleText] = useSmoothText(text, options);
+  return visibleText;
+}

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!" },
       ]),
     ]);
   });