@convex-dev/agent 0.5.0-alpha.1 → 0.6.0-alpha.1

package/src/component/schema.ts

@@ -1,5 +1,5 @@
 import { defineSchema, defineTable } from "convex/server";
-import { type Infer, v } from "convex/values";
+import { v } from "convex/values";
 import {
   vThreadStatus,
   vMessage,
@@ -11,11 +11,9 @@ import {
   vProviderOptions,
   vProviderMetadata,
   vReasoningDetails,
-  vTextStreamPart,
 } from "../validators.js";
 import { typedV } from "convex-helpers/validators";
 import vectorTables, { vVectorId } from "./vector/tables.js";
-import { omit } from "convex-helpers";
 
 export const schema = defineSchema({
   threads: defineTable({
@@ -29,12 +27,8 @@ export const schema = defineSchema({
     order: /*DEPRECATED*/ v.optional(v.number()),
   })
     .index("userId", ["userId"])
-    .searchIndex("title", {
-      searchField: "title",
-      filterFields: ["userId"],
-    }),
+    .searchIndex("title", { searchField: "title", filterFields: ["userId"] }),
   messages: defineTable({
-    id: v.optional(v.string()), // external id, e.g. from Vercel AI SDK
     userId: v.optional(v.string()), // useful for searching across threads
     threadId: v.id("threads"),
     order: v.number(),
@@ -60,11 +54,13 @@ export const schema = defineSchema({
     usage: v.optional(vUsage),
     providerMetadata: v.optional(vProviderMetadata), // Received from model
     sources: v.optional(v.array(vSource)),
-    reasoning: v.optional(v.string()),
-    reasoningDetails: v.optional(vReasoningDetails),
     warnings: v.optional(v.array(vLanguageModelCallWarning)),
     finishReason: v.optional(vFinishReason),
+    // Likely deprecated soon
+    reasoning: v.optional(v.string()),
+    reasoningDetails: v.optional(vReasoningDetails),
     // DEPRECATED
+    id: v.optional(v.string()), // external id, e.g. from Vercel AI SDK
     parentMessageId: v.optional(v.id("messages")),
     stepId: v.optional(v.string()),
     files: v.optional(v.array(v.any())),
@@ -95,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(),
@@ -115,10 +117,7 @@ export const schema = defineSchema({
         endedAt: v.number(),
         cleanupFnId: v.optional(v.id("_scheduled_functions")),
       }),
-      v.object({
-        kind: v.literal("aborted"),
-        reason: v.string(),
-      }),
+      v.object({ kind: v.literal("aborted"), reason: v.string() }),
     ),
   })
     // There should only be one per "order" index
@@ -135,7 +134,7 @@ export const schema = defineSchema({
     // the indexes work like: 0 <first> 1 <second> 2 <third> 3 ...
     start: v.number(), // inclusive
     end: v.number(), // exclusive
-    parts: v.array(vTextStreamPart),
+    parts: v.array(v.any()),
   }).index("streamId_start_end", ["streamId", "start", "end"]),
 
   memories: defineTable({
@@ -150,7 +149,9 @@ export const schema = defineSchema({
 
   files: defineTable({
     storageId: v.string(),
-    mimeType: v.string(),
+    mediaType: v.optional(v.string()),
+    /** @deprecated Use `mediaType` instead. */
+    mimeType: v.optional(v.string()),
     filename: v.optional(v.string()),
     hash: v.string(),
     refcount: v.number(),
@@ -162,39 +163,12 @@ export const schema = defineSchema({
   // To authenticate playground usage
   // Delete a key to invalidate it
   // Provide a name to easily identify it / invalidate by name
-  apiKeys: defineTable({
-    name: v.optional(v.string()),
-  }).index("name", ["name"]),
+  apiKeys: defineTable({ name: v.optional(v.string()) }).index("name", [
+    "name",
+  ]),
 });
 
 export const vv = typedV(schema);
 export { vv as v };
 
-// Public
-export const vThreadDoc = v.object({
-  _id: v.string(),
-  _creationTime: v.number(),
-  userId: v.optional(v.string()), // Unset for anonymous
-  title: v.optional(v.string()),
-  summary: v.optional(v.string()),
-  status: vThreadStatus,
-});
-export type ThreadDoc = Infer<typeof vThreadDoc>;
-
-export const vMessageDoc = v.object({
-  _id: v.string(),
-  _creationTime: v.number(),
-  ...omit(schema.tables.messages.validator.fields, [
-    "parentMessageId",
-    "stepId",
-    "files",
-  ]),
-  // Overwrite all the types that have a v.id validator
-  // Outside of the component, they are strings
-  threadId: v.string(),
-  embeddingId: v.optional(v.string()),
-  fileIds: v.optional(v.array(v.string())),
-});
-export type MessageDoc = Infer<typeof vMessageDoc>;
-
 export default schema;

package/src/component/setup.test.ts

@@ -1,5 +1,12 @@
 /// <reference types="vite/client" />
+import { convexTest } from "convex-test";
 import { test } from "vitest";
+import schema from "./schema.js";
 export const modules = import.meta.glob("./**/*.*s");
 
+export function initConvexTest() {
+  const t = convexTest(schema, modules);
+  return t;
+}
+
 test("setup", () => {});

package/src/component/streams.ts

@@ -1,7 +1,9 @@
 import { omit, pick } from "convex-helpers";
 import { v } from "convex/values";
 import {
+  type MessageWithMetadataInternal,
   type StreamDelta,
+  type StreamMessage,
   vStreamDelta,
   vStreamMessage,
 } from "../validators.js";
@@ -18,6 +20,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 { deriveUIMessagesFromDeltas } from "../deltas.js";
+import { fromUIMessages } from "../UIMessages.js";
 
 const SECOND = 1000;
 const MINUTE = 60 * SECOND;
@@ -33,13 +38,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 +56,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 +90,8 @@ 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() };
+    // TODO: enforce order/stepOrder uniqueness?
     const streamId = await ctx.db.insert("streamingMessages", {
       ...args,
       state,
@@ -135,28 +138,29 @@ export const list = query({
       ["order", "stepOrder"],
     ).take(100);
 
-    return messages.map((m) => ({
-      streamId: m._id,
-      status: m.state.kind,
-      ...pick(m, [
-        "order",
-        "stepOrder",
-        "userId",
-        "agentName",
-        "model",
-        "provider",
-        "providerOptions",
-      ]),
-    }));
+    return messages.map((m) => publicStreamMessage(m));
   },
 });
 
+function publicStreamMessage(m: Doc<"streamingMessages">): StreamMessage {
+  return {
+    streamId: m._id,
+    status: m.state.kind,
+    ...pick(m, [
+      "format",
+      "order",
+      "stepOrder",
+      "userId",
+      "agentName",
+      "model",
+      "provider",
+      "providerOptions",
+    ]),
+  };
+}
+
 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 +173,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 +183,7 @@ export const abort = mutation({
   args: {
     streamId: v.id("streamingMessages"),
     reason: v.string(),
+    finalDelta: v.optional(deltaValidator),
   },
   returns: v.boolean(),
   handler: abortById,
@@ -189,16 +191,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 +226,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 +308,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 +322,7 @@ export const timeoutStream = internalMutation({
       return;
     }
     await ctx.db.patch(args.streamId, {
-      state: {
-        kind: "aborted",
-        reason: "timeout",
-      },
+      state: { kind: "aborted", reason: "timeout" },
     });
   },
 });
@@ -362,11 +377,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 +386,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({
@@ -407,7 +414,14 @@ export const deleteAllStreamsForThreadIdAsync = mutation({
     streamOrder: v.optional(v.number()),
     deltaCursor: v.optional(v.string()),
   }),
-  handler: async (ctx, args) => {
+  handler: async (
+    ctx,
+    args,
+  ): Promise<{
+    isDone: boolean;
+    streamOrder?: number;
+    deltaCursor?: string;
+  }> => {
     const result = await deleteStreamsPageForThreadId(ctx, args);
     if (!result.isDone) {
       await ctx.scheduler.runAfter(
@@ -474,3 +488,90 @@ 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(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,
+          [publicStreamMessage(streamingMessage)],
+          deltas,
+        );
+        // We don't save messages that have already been saved
+        const numToSkip = stepOrder - streamingMessage.stepOrder;
+        const messages = await Promise.all(
+          (await 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;
+}

package/src/component/users.test.ts

@@ -58,7 +58,6 @@ describe("users", () => {
 
     expect(secondPage.page).toHaveLength(2);
     // Should not have duplicate users
-    console.log(firstPage.page, secondPage.page);
     expect(
       firstPage.page.every((user) => !secondPage.page.includes(user)),
     ).toBe(true);

package/src/component/vector/index.ts

@@ -35,11 +35,9 @@ export const paginate = query({
     const tableName = getVectorTableName(args.vectorDimension);
     const vectors = await paginator(ctx.db, schema)
       .query(tableName)
-      // eslint-disable-next-line @typescript-eslint/no-explicit-any
       .withIndex("model_table_threadId" as any, (q) =>
         args.table
-          ? // eslint-disable-next-line @typescript-eslint/no-explicit-any
-            (q.eq("model", args.targetModel) as any).eq("table", args.table)
+          ? (q.eq("model", args.targetModel) as any).eq("table", args.table)
           : q.eq("model", args.targetModel),
       )
       .paginate({