@convex-dev/agent 0.1.6 → 0.1.7-alpha.1

package/src/component/files.test.ts

@@ -0,0 +1,130 @@
+/// <reference types="vite/client" />
+
+import { convexTest } from "convex-test";
+import { describe, expect, test } from "vitest";
+import { api } from "./_generated/api.js";
+import schema from "./schema.js";
+import { modules } from "./setup.test.js";
+import type { Doc } from "./_generated/dataModel.js";
+import type { PaginationResult } from "convex/server";
+
+describe("files", () => {
+  test("addFile increments refcount and does not create a new entry", async () => {
+    const t = convexTest(schema, modules);
+    const storageId = "storage-1";
+    const hash = "hash-1";
+    const filename = "file.txt";
+    // Add the file for the first time
+    const { fileId, storageId: returnedStorageId } = await t.mutation(
+      api.files.addFile,
+      {
+        storageId,
+        hash,
+        filename,
+        mimeType: "text/plain",
+      }
+    );
+    expect(fileId).toBeTruthy();
+    expect(returnedStorageId).toBe(storageId);
+    // Add the same file again
+    const { fileId: fileId2 } = await t.mutation(api.files.addFile, {
+      storageId,
+      hash,
+      filename,
+      mimeType: "text/plain",
+    });
+    expect(fileId2).toBe(fileId);
+    // Add the same file with a different filename (should create a new entry)
+    const { fileId: fileId3 } = await t.mutation(api.files.addFile, {
+      storageId,
+      hash,
+      filename: "other.txt",
+      mimeType: "text/plain",
+    });
+    expect(fileId3).not.toBe(fileId);
+    // Add the same file with undefined filename (should create a new entry)
+    const { fileId: fileId4 } = await t.mutation(api.files.addFile, {
+      storageId,
+      hash,
+      filename: undefined,
+      mimeType: "text/plain",
+    });
+    expect(fileId4).not.toBe(fileId);
+  });
+
+  test("useExistingFile only matches files with the same hash and filename", async () => {
+    const t = convexTest(schema, modules);
+    const storageId = "storage-2";
+    const hash = "hash-2";
+    const filename = "file2.txt";
+    // Add a file
+    const { fileId } = await t.mutation(api.files.addFile, {
+      storageId,
+      hash,
+      filename,
+      mimeType: "text/plain",
+    });
+    // Should match
+    const fileId2 = await t.mutation(api.files.useExistingFile, {
+      hash,
+      filename,
+    });
+    expect(fileId2?.fileId).toBe(fileId);
+    // Should not match with different filename
+    const fileId3 = await t.mutation(api.files.useExistingFile, {
+      hash,
+      filename: "other2.txt",
+    });
+    expect(fileId3).toBeNull();
+    // Should not match with undefined filename
+    const fileId4 = await t.mutation(api.files.useExistingFile, {
+      hash,
+    });
+    expect(fileId4).toBeNull();
+  });
+
+  test("getFilesToDelete paginates through files with refcount 0 one at a time", async () => {
+    const t = convexTest(schema, modules);
+    // Add 3 files with refcount 0
+    const files = [];
+    for (let i = 0; i < 3; i++) {
+      const { fileId } = await t.mutation(api.files.addFile, {
+        storageId: `storage-del-${i}`,
+        hash: `hash-del-${i}`,
+        filename: `file-del-${i}.txt`,
+        mimeType: "text/plain",
+      });
+      // Manually set refcount to 0
+      await t.run(async (ctx) => {
+        await ctx.db.patch(fileId, { refcount: 0 });
+      });
+      files.push(fileId);
+    }
+    // Paginate through files to delete one at a time
+    let cursor: string | null = null;
+    const seen: string[] = [];
+    for (let i = 0; i < 3; i++) {
+      const { page, continueCursor, isDone }: PaginationResult<Doc<"files">> =
+        await t.query(api.files.getFilesToDelete, {
+          paginationOpts: {
+            numItems: 1,
+            cursor,
+          },
+        });
+      expect(page.length).toBe(1);
+      seen.push(page[0]._id);
+      cursor = continueCursor;
+      expect(isDone).toBe(false);
+    }
+    const { page, isDone } = await t.query(api.files.getFilesToDelete, {
+      paginationOpts: {
+        numItems: 1,
+        cursor,
+      },
+    });
+    expect(page.length).toBe(0);
+    expect(isDone).toBe(true);
+    // All fileIds should be seen
+    expect(seen.sort()).toEqual(files.sort());
+  });
+});

package/src/component/files.ts

@@ -2,71 +2,99 @@ import { paginator } from "convex-helpers/server/pagination";
 import type { Id } from "./_generated/dataModel.js";
 import { mutation, type MutationCtx, query } from "./_generated/server.js";
 import { schema, v } from "./schema.js";
+import { paginationOptsValidator } from "convex/server";
+import type { Infer } from "convex/values";
+
+const addFileArgs = v.object({
+  storageId: v.string(),
+  hash: v.string(),
+  filename: v.optional(v.string()),
+  mimeType: v.string(),
+});
 
 export const addFile = mutation({
-  args: {
-    storageId: v.string(),
-    hash: v.string(),
-  },
+  args: addFileArgs,
   handler: addFileHandler,
   returns: {
     fileId: v.id("files"),
-    storageIdUnused: v.boolean(),
+    storageId: v.string(),
   },
 });
 
 export async function addFileHandler(
   ctx: MutationCtx,
-  args: { storageId: string; hash: string }
+  args: Infer<typeof addFileArgs>
 ) {
   const existingFile = await ctx.db
     .query("files")
     .withIndex("hash", (q) => q.eq("hash", args.hash))
+    .filter((q) => q.eq(q.field("filename"), args.filename))
     .first();
   if (existingFile) {
     // increment the refcount
     await ctx.db.patch(existingFile._id, {
       refcount: existingFile.refcount + 1,
+      lastTouchedAt: Date.now(),
     });
     return {
       fileId: existingFile._id,
-      storageIdUnused: existingFile.storageId !== args.storageId,
+      storageId: existingFile.storageId,
     };
   }
   const fileId = await ctx.db.insert("files", {
-    storageId: args.storageId,
-    hash: args.hash,
-    refcount: 1,
+    ...args,
+    // We start out with it unused - when it's saved in a message we increment.
+    refcount: 0,
+    lastTouchedAt: Date.now(),
   });
   return {
     fileId,
-    storageIdUnused: false,
+    storageId: args.storageId,
   };
 }
 
+export const get = query({
+  args: {
+    fileId: v.id("files"),
+  },
+  returns: v.union(v.null(), v.doc("files")),
+  handler: async (ctx, args) => {
+    return ctx.db.get(args.fileId);
+  },
+});
+
 /**
  * If you plan to have the same file added over and over without a reference to
  * the fileId, you can use this query to get the fileId of the existing file.
- * And if it's
+ * Note: this will not increment the refcount. only saving messages does that.
+ * It will only match if the filename is the same (or both are undefined).
  */
 export const useExistingFile = mutation({
   args: {
     hash: v.string(),
+    filename: v.optional(v.string()),
   },
   handler: async (ctx, args) => {
     const file = await ctx.db
       .query("files")
       .withIndex("hash", (q) => q.eq("hash", args.hash))
+      .filter((q) => q.eq(q.field("filename"), args.filename))
       .first();
     if (!file) {
       return null;
     }
     await ctx.db.patch(file._id, {
-      refcount: file.refcount + 1,
+      lastTouchedAt: Date.now(),
     });
-    return file._id;
+    return { fileId: file._id, storageId: file.storageId };
   },
-  returns: v.union(v.id("files"), v.null()),
+  returns: v.union(
+    v.null(),
+    v.object({
+      fileId: v.id("files"),
+      storageId: v.string(),
+    })
+  ),
 });
 
 export const copyFile = mutation({
@@ -87,31 +115,59 @@ export async function copyFileHandler(
   }
   await ctx.db.patch(args.fileId, {
     refcount: file.refcount + 1,
+    lastTouchedAt: Date.now(),
   });
 }
 
+/**
+ * Get files that are unused and can be deleted.
+ * This is useful for cleaning up files that are no longer needed.
+ * Note: recently added files that have not been saved yet will show up here.
+ * You can inspect the `lastTouchedAt` field to see how recently it was used.
+ * I'd recommend not deleting anything touched in the last 24 hours.
+ */
 export const getFilesToDelete = query({
   args: {
-    cursor: v.optional(v.string()),
-    limit: v.optional(v.number()),
+    paginationOpts: paginationOptsValidator,
   },
   handler: async (ctx, args) => {
     const files = await paginator(ctx.db, schema)
       .query("files")
       .withIndex("refcount", (q) => q.eq("refcount", 0))
-      .paginate({
-        numItems: args.limit ?? 100,
-        cursor: args.cursor ?? null,
-      });
-    return {
-      files: files.page,
-      continueCursor: files.continueCursor,
-      isDone: files.isDone,
-    };
+      .paginate(args.paginationOpts);
+    return files;
   },
   returns: v.object({
-    files: v.array(v.doc("files")),
+    page: v.array(v.doc("files")),
     continueCursor: v.string(),
     isDone: v.boolean(),
   }),
 });
+
+export const deleteFiles = mutation({
+  args: {
+    fileIds: v.array(v.id("files")),
+    force: v.optional(v.boolean()),
+  },
+  handler: async (ctx, args) => {
+    await Promise.all(
+      args.fileIds.map(async (fileId) => {
+        const file = await ctx.db.get(fileId);
+        if (!file) {
+          console.error(`File ${fileId} not found when deleting, skipping...`);
+          return;
+        }
+        if (file.refcount && file.refcount > 0) {
+          if (!args.force) {
+            console.error(
+              `File ${fileId} has refcount ${file.refcount} > 0, skipping...`
+            );
+            return;
+          }
+        }
+        await ctx.db.delete(fileId);
+      })
+    );
+  },
+  returns: v.null(),
+});

package/src/component/messages.test.ts

@@ -129,7 +129,8 @@ describe("agent", () => {
     const { messages: messages2 } = await t.mutation(api.messages.addMessages, {
       threadId: thread._id as Id<"threads">,
       messages: [{ message: { role: "user", content: "hello" } }],
-      promptMessageId: messages.at(-1)!._id,
+      agentName: "test",
+      promptMessageId: messages.at(-1)!._id as Id<"messages">,
     });
     const maxMessage2 = await t.run(async (ctx) => {
       return await getMaxMessage(ctx, thread._id as Id<"threads">);

package/src/component/messages.ts

@@ -11,10 +11,9 @@ import {
 import {
   vMessageEmbeddings,
   vMessageStatus,
-  vMessageWithMetadata,
+  vMessageWithMetadataInternal,
   vPaginationResult,
   vSearchOptions,
-  vStepWithMessages,
 } from "../validators.js";
 import { api, internal } from "./_generated/api.js";
 import type { Doc, Id } from "./_generated/dataModel.js";
@@ -51,7 +50,7 @@ export const getThread = _getThread;
 export const updateThread = _updateThread;
 
 function publicMessage(message: Doc<"messages">): MessageDoc {
-  return omit(message, ["parentMessageId", "stepId"]);
+  return omit(message, ["parentMessageId", "stepId", "files"]);
 }
 
 export async function deleteMessage(
@@ -62,7 +61,7 @@ export async function deleteMessage(
   if (messageDoc.embeddingId) {
     await ctx.db.delete(messageDoc.embeddingId);
   }
-  for (const { fileId } of messageDoc.files ?? []) {
+  for (const fileId of messageDoc.fileIds ?? []) {
     if (!fileId) continue;
     const file = await ctx.db.get(fileId);
     if (file) {
@@ -78,10 +77,9 @@ export const messageStatuses = vMessageDoc.fields.status.members.map(
 const addMessagesArgs = {
   userId: v.optional(v.string()),
   threadId: v.id("threads"),
-  stepId: v.optional(v.id("steps")),
   promptMessageId: v.optional(v.id("messages")),
   agentName: v.optional(v.string()),
-  messages: v.array(vMessageWithMetadata),
+  messages: v.array(vMessageWithMetadataInternal),
   embeddings: v.optional(vMessageEmbeddings),
   pending: v.optional(v.boolean()),
   failPendingSteps: v.optional(v.boolean()),
@@ -91,7 +89,6 @@ export const addMessages = mutation({
   handler: addMessagesHandler,
   returns: v.object({
     messages: v.array(vMessageDoc),
-    pending: v.optional(vMessageDoc),
   }),
 });
 async function addMessagesHandler(
@@ -147,55 +144,53 @@ async function addMessagesHandler(
     stepOrder = -1;
   }
   const toReturn: Doc<"messages">[] = [];
-  if (messages.length > 0) {
-    if (embeddings) {
-      assert(
-        embeddings.vectors.length === messages.length,
-        "embeddings.vectors.length must match messages.length"
-      );
-    }
-    for (let i = 0; i < messages.length; i++) {
-      const message = messages[i];
-      let embeddingId: VectorTableId | undefined;
-      if (embeddings && embeddings.vectors[i]) {
-        embeddingId = await insertVector(ctx, embeddings.dimension, {
-          vector: embeddings.vectors[i]!,
-          model: embeddings.model,
-          table: "messages",
-          userId,
-          threadId,
-        });
-      }
-      stepOrder++;
-      const messageId = await ctx.db.insert("messages", {
-        ...rest,
-        ...message,
-        embeddingId,
-        parentMessageId: promptMessageId,
+  if (embeddings) {
+    assert(
+      embeddings.vectors.length === messages.length,
+      "embeddings.vectors.length must match messages.length"
+    );
+  }
+  for (let i = 0; i < messages.length; i++) {
+    const message = messages[i];
+    let embeddingId: VectorTableId | undefined;
+    if (embeddings && embeddings.vectors[i]) {
+      embeddingId = await insertVector(ctx, embeddings.dimension, {
+        vector: embeddings.vectors[i]!,
+        model: embeddings.model,
+        table: "messages",
         userId,
-        order,
-        tool: isTool(message.message),
-        text: extractText(message.message),
-        status: fail ? "failed" : pending ? "pending" : "success",
-        error: fail ? "Parent message failed" : undefined,
-        stepOrder,
+        threadId,
+      });
+    }
+    stepOrder++;
+    const messageId = await ctx.db.insert("messages", {
+      ...rest,
+      ...message,
+      embeddingId,
+      parentMessageId: promptMessageId,
+      userId,
+      order,
+      tool: isTool(message.message),
+      text: extractText(message.message),
+      status: fail ? "failed" : pending ? "pending" : "success",
+      error: fail ? "Parent message failed" : undefined,
+      stepOrder,
+    });
+    // Let's just not set the id field and have it set only in explicit cases.
+    // if (!message.id) {
+    //   await ctx.db.patch(messageId, {
+    //     id: messageId,
+    //   });
+    // }
+    for (const fileId of message.fileIds ?? []) {
+      if (!fileId) continue;
+      await ctx.db.patch(fileId, {
+        refcount: (await ctx.db.get(fileId))!.refcount + 1,
       });
-      // Let's just not set the id field and have it set only in explicit cases.
-      // if (!message.id) {
-      //   await ctx.db.patch(messageId, {
-      //     id: messageId,
-      //   });
-      // }
-      for (const { fileId } of message.files ?? []) {
-        if (!fileId) continue;
-        await ctx.db.patch(fileId, {
-          refcount: (await ctx.db.get(fileId))!.refcount + 1,
-        });
-      }
-      toReturn.push((await ctx.db.get(messageId))!);
     }
+    toReturn.push((await ctx.db.get(messageId))!);
   }
-  return { messages: toReturn };
+  return { messages: toReturn.map(publicMessage) };
 }
 
 // exported for tests
@@ -235,68 +230,6 @@ function orderedMessagesStream(
   );
 }
 
-const addStepArgs = {
-  userId: v.optional(v.string()),
-  threadId: v.id("threads"),
-  promptMessageId: v.id("messages"),
-  step: vStepWithMessages,
-  failPendingSteps: v.optional(v.boolean()),
-};
-
-export const addStep = mutation({
-  args: addStepArgs,
-  returns: v.array(vMessageDoc),
-  handler: addStepHandler,
-});
-async function addStepHandler(
-  ctx: MutationCtx,
-  args: ObjectType<typeof addStepArgs>
-) {
-  const parentMessage = await ctx.db.get(args.promptMessageId);
-  assert(parentMessage, `Message ${args.promptMessageId} not found`);
-  const order = parentMessage.order;
-  // TODO: only fetch the last one if we aren't failing pending steps
-  let steps = await ctx.db
-    .query("steps")
-    .withIndex("parentMessageId_order_stepOrder", (q) =>
-      // TODO: fetch pending, and commit later
-      q.eq("parentMessageId", args.promptMessageId)
-    )
-    .collect();
-  if (args.failPendingSteps) {
-    for (const step of steps) {
-      if (step.status === "pending") {
-        await ctx.db.patch(step._id, { status: "failed" });
-      }
-    }
-    steps = steps.filter((s) => s.status !== "failed");
-  }
-  const { step, messages } = args.step;
-  const stepId = await ctx.db.insert("steps", {
-    threadId: args.threadId,
-    parentMessageId: args.promptMessageId,
-    order,
-    stepOrder: (steps.at(-1)?.stepOrder ?? -1) + 1,
-    status: step.finishReason === "stop" ? "success" : "pending",
-    step,
-  });
-  const added = await addMessagesHandler(ctx, {
-    userId: args.userId,
-    threadId: args.threadId,
-    stepId,
-    promptMessageId: args.promptMessageId,
-    agentName: parentMessage.agentName,
-    messages,
-    pending: step.finishReason === "stop" ? false : true,
-    failPendingSteps: false,
-  });
-  // We don't commit if the parent is still pending.
-  if (step.finishReason === "stop") {
-    await commitMessageHandler(ctx, { messageId: args.promptMessageId });
-  }
-  return added.messages.map(publicMessage);
-}
-
 export const rollbackMessage = mutation({
   args: {
     messageId: v.id("messages"),
@@ -318,18 +251,6 @@ export const rollbackMessage = mutation({
       }
     }
 
-    const steps = await ctx.db
-      .query("steps")
-      .withIndex("parentMessageId_order_stepOrder", (q) =>
-        // TODO: fetch pending, and commit later
-        q.eq("parentMessageId", messageId)
-      )
-      .collect();
-    for (const step of steps) {
-      if (step.status === "pending") {
-        await ctx.db.patch(step._id, { status: "failed" });
-      }
-    }
     await ctx.db.patch(messageId, {
       status: "failed",
       error: error,
@@ -351,17 +272,6 @@ async function commitMessageHandler(
   const message = await ctx.db.get(messageId);
   assert(message, `Message ${messageId} not found`);
 
-  const allSteps = await ctx.db
-    .query("steps")
-    .withIndex("parentMessageId_order_stepOrder", (q) =>
-      q.eq("parentMessageId", messageId)
-    )
-    .collect();
-  for (const step of allSteps) {
-    if (step.status === "pending") {
-      await ctx.db.patch(step._id, { status: "success" });
-    }
-  }
   const order = message.order!;
   const messages = await mergedStream(
     [true, false].map((tool) =>

package/src/component/schema.ts

@@ -4,7 +4,6 @@ import {
   vThreadStatus,
   vMessage,
   vMessageStatus,
-  vStep,
   vUsage,
   vSource,
   vLanguageModelV1CallWarning,
@@ -12,8 +11,6 @@ import {
   vProviderOptions,
   vProviderMetadata,
   vReasoningDetails,
-  vFile,
-  vFileWithStringId,
   vTextStreamPart,
 } from "../validators.js";
 import { typedV } from "convex-helpers/validators";
@@ -35,10 +32,10 @@ export const schema = defineSchema({
     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"),
-    // Repeats until a non-tool message.
     order: v.number(),
     stepOrder: v.number(),
     embeddingId: v.optional(vVectorId),
+    fileIds: v.optional(v.array(v.id("files"))),
     error: v.optional(v.string()),
     status: vMessageStatus,
 
@@ -53,7 +50,6 @@ export const schema = defineSchema({
     // Convenience fields extracted from the message
     tool: v.boolean(), // either tool call (assistant) or tool result (tool)
     text: v.optional(v.string()),
-    files: v.optional(v.array(vFile)),
 
     // Result metadata
     usage: v.optional(vUsage),
@@ -65,7 +61,8 @@ export const schema = defineSchema({
     finishReason: v.optional(vFinishReason),
     // DEPRECATED
     parentMessageId: v.optional(v.id("messages")),
-    stepId: v.optional(v.id("steps")),
+    stepId: v.optional(v.string()),
+    files: v.optional(v.array(v.any())),
   })
     // Allows finding successful visible messages in order
     // Also surface pending messages separately to e.g. stream
@@ -85,27 +82,6 @@ export const schema = defineSchema({
     // Allows finding messages by vector embedding id
     .index("embeddingId", ["embeddingId"]),
 
-  steps: defineTable({
-    threadId: v.id("threads"),
-    // Could be different from the order if we fail.
-    parentMessageId: v.id("messages"),
-    order: v.number(), // parent message order
-    stepOrder: v.number(), // step order
-    step: vStep,
-    status: vMessageStatus,
-  })
-    .index("status_threadId_order_stepOrder", [
-      "status",
-      "threadId",
-      "order",
-      "stepOrder",
-    ])
-    .index("parentMessageId_order_stepOrder", [
-      "parentMessageId",
-      "order",
-      "stepOrder",
-    ]),
-
   // Status: if it's done, it's deleted, then deltas are vacuumed
   streamingMessages: defineTable({
     // extra metadata?
@@ -168,8 +144,11 @@ export const schema = defineSchema({
 
   files: defineTable({
     storageId: v.string(),
+    mimeType: v.string(),
+    filename: v.optional(v.string()),
     hash: v.string(),
     refcount: v.number(),
+    lastTouchedAt: v.number(),
   })
     .index("hash", ["hash"])
     .index("refcount", ["refcount"]),
@@ -202,12 +181,13 @@ export const vMessageDoc = v.object({
   ...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()),
-  files: v.optional(v.array(vFileWithStringId)),
+  fileIds: v.optional(v.array(v.string())),
 });
 export type MessageDoc = Infer<typeof vMessageDoc>;