@convex-dev/agent 0.7.0 → 0.7.1

package/src/component/files.test.ts

@@ -1,7 +1,7 @@
 /// <reference types="vite/client" />
 
 import { convexTest } from "convex-test";
-import { describe, expect, test } from "vitest";
+import { afterEach, describe, expect, test, vi } from "vitest";
 import { api } from "./_generated/api.js";
 import schema from "./schema.js";
 import { modules } from "./setup.test.js";
@@ -9,7 +9,8 @@ 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 () => {
+  afterEach(() => vi.useRealTimers());
+  test("addFile registers and reuses a file without acquiring ownership", async () => {
     const t = convexTest(schema, modules);
     const storageId = "storage-1";
     const hash = "hash-1";
@@ -34,6 +35,9 @@ describe("files", () => {
       mimeType: "text/plain",
     });
     expect(fileId2).toBe(fileId);
+    await expect(t.query(api.files.get, { fileId })).resolves.toMatchObject({
+      refcount: 0,
+    });
     // Add the same file with a different filename (should create a new entry)
     const { fileId: fileId3 } = await t.mutation(api.files.addFile, {
       storageId,
@@ -83,7 +87,31 @@ describe("files", () => {
     expect(fileId4).toBeNull();
   });
 
+  test("reuses a legacy file row with no media type", async () => {
+    const t = convexTest(schema, modules);
+    const fileId = await t.run((ctx) =>
+      ctx.db.insert("files", {
+        storageId: "legacy-storage",
+        hash: "legacy-hash",
+        filename: "legacy.txt",
+        refcount: 0,
+        lastTouchedAt: Date.now(),
+      }),
+    );
+
+    await expect(
+      t.mutation(api.files.addFile, {
+        storageId: "new-storage",
+        hash: "legacy-hash",
+        filename: "legacy.txt",
+        mediaType: "text/plain",
+      }),
+    ).resolves.toEqual({ fileId, storageId: "legacy-storage" });
+  });
+
   test("getFilesToDelete paginates through files with refcount 0 one at a time", async () => {
+    vi.useFakeTimers();
+    vi.setSystemTime(new Date("2026-01-01T00:00:00.000Z"));
     const t = convexTest(schema, modules);
     // Add 3 files with refcount 0
     const files = [];
@@ -94,12 +122,9 @@ describe("files", () => {
         filename: `file-del-${i}.txt`,
         mimeType: "text/plain",
       });
-      // Manually set refcount to 0
-      await t.run(async (ctx) => {
-        await ctx.db.patch("files", fileId, { refcount: 0 });
-      });
       files.push(fileId);
     }
+    vi.advanceTimersByTime(24 * 60 * 60 * 1000);
     // Paginate through files to delete one at a time
     let cursor: string | null = null;
     const seen: string[] = [];
@@ -127,4 +152,54 @@ describe("files", () => {
     // All fileIds should be seen
     expect(seen.sort()).toEqual(files.sort());
   });
+
+  test("does not expose or delete freshly registered files", async () => {
+    const t = convexTest(schema, modules);
+    const error = vi.spyOn(console, "error").mockImplementation(() => {});
+    const { fileId } = await t.mutation(api.files.addFile, {
+      storageId: "storage-fresh",
+      hash: "hash-fresh",
+      filename: "fresh.txt",
+      mimeType: "text/plain",
+    });
+
+    await expect(
+      t.query(api.files.getFilesToDelete, {
+        paginationOpts: { numItems: 10, cursor: null },
+      }),
+    ).resolves.toMatchObject({ page: [] });
+    await expect(
+      t.mutation(api.files.deleteFiles, { fileIds: [fileId] }),
+    ).resolves.toEqual([]);
+    await expect(t.query(api.files.get, { fileId })).resolves.not.toBeNull();
+    error.mockRestore();
+  });
+
+  test("refreshes a cache hit before the cleanup grace period expires", async () => {
+    vi.useFakeTimers();
+    vi.setSystemTime(new Date("2026-01-01T00:00:00.000Z"));
+    const t = convexTest(schema, modules);
+    const { fileId } = await t.mutation(api.files.addFile, {
+      storageId: "storage-cache-hit",
+      hash: "cache-hit-hash",
+      filename: "cache-hit.txt",
+      mediaType: "text/plain",
+    });
+
+    vi.advanceTimersByTime(24 * 60 * 60 * 1000 - 60 * 1000);
+    await expect(
+      t.mutation(api.files.useExistingFile, {
+        hash: "cache-hit-hash",
+        filename: "cache-hit.txt",
+        mediaType: "text/plain",
+      }),
+    ).resolves.toMatchObject({ fileId });
+
+    vi.advanceTimersByTime(60 * 1000);
+    await expect(
+      t.query(api.files.getFilesToDelete, {
+        paginationOpts: { numItems: 10, cursor: null },
+      }),
+    ).resolves.toMatchObject({ page: [] });
+  });
 });

package/src/component/files.ts

@@ -1,10 +1,12 @@
 import { paginator } from "convex-helpers/server/pagination";
-import type { Id } from "./_generated/dataModel.js";
+import type { Doc, 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 FILE_CLEANUP_GRACE_MS = 24 * 60 * 60 * 1000;
+
 const addFileArgs = v.object({
   storageId: v.string(),
   hash: v.string(),
@@ -28,18 +30,21 @@ export async function addFileHandler(
   args: Infer<typeof addFileArgs>,
 ) {
   // Support both mediaType (preferred) and mimeType (deprecated)
-  const mediaType = args.mediaType ?? args.mimeType;
+  const mediaType = normalizeMediaType(args.mediaType ?? args.mimeType);
 
-  const existingFile = await ctx.db
+  const existingFiles = await ctx.db
     .query("files")
     .withIndex("hash", (q) => q.eq("hash", args.hash))
     // eslint-disable-next-line @convex-dev/no-filter-in-query -- We do not expect many files with the same hash and different filenames
     .filter((q) => q.eq(q.field("filename"), args.filename))
-    .first();
+    .collect();
+  const existingFile = existingFiles.find((file) =>
+    sameMediaType(file, mediaType),
+  );
   if (existingFile) {
-    // increment the refcount
+    // Registration is not ownership. Persisted messages and in-flight streams
+    // acquire ownership explicitly.
     await ctx.db.patch("files", existingFile._id, {
-      refcount: existingFile.refcount + 1,
       lastTouchedAt: Date.now(),
     });
     return {
@@ -83,14 +88,18 @@ export const useExistingFile = mutation({
   args: {
     hash: v.string(),
     filename: v.optional(v.string()),
+    mediaType: v.optional(v.string()),
   },
   handler: async (ctx, args) => {
-    const file = await ctx.db
+    const files = await ctx.db
       .query("files")
       .withIndex("hash", (q) => q.eq("hash", args.hash))
       // eslint-disable-next-line @convex-dev/no-filter-in-query -- We do not expect many files with the same hash and different filenames
       .filter((q) => q.eq(q.field("filename"), args.filename))
-      .first();
+      .collect();
+    const file = files.find((candidate) =>
+      sameMediaType(candidate, normalizeMediaType(args.mediaType)),
+    );
     if (!file) {
       return null;
     }
@@ -108,36 +117,51 @@ export const useExistingFile = mutation({
   ),
 });
 
+function normalizeMediaType(value: string | undefined) {
+  const normalized = value?.trim().toLowerCase();
+  return normalized || undefined;
+}
+
+function sameMediaType(
+  file: Pick<Doc<"files">, "mediaType" | "mimeType">,
+  expected: string | undefined,
+) {
+  const actual = normalizeMediaType(file.mediaType ?? file.mimeType);
+  // File rows from before media-type tracking have no value to distinguish.
+  // Preserve the prior hash-and-filename reuse behavior for those rows.
+  return expected === undefined || actual === undefined || actual === expected;
+}
+
+/** Transfer ownership between two durable records. */
 export async function changeRefcount(
   ctx: MutationCtx,
-  prev: Id<"files">[],
+  previous: Id<"files">[],
   next: Id<"files">[],
 ) {
-  const prevSet = new Set(prev);
+  const previousSet = new Set(previous);
   const nextSet = new Set(next);
-  for (const fileId of prevSet) {
-    if (!nextSet.has(fileId)) {
-      const file = await ctx.db.get("files", fileId);
-      if (file) {
-        await ctx.db.patch("files", fileId, {
-          refcount: file.refcount - 1,
-        });
-      } else {
+  for (const fileId of new Set([...previousSet, ...nextSet])) {
+    const increment = Number(!previousSet.has(fileId) && nextSet.has(fileId));
+    const decrement = Number(previousSet.has(fileId) && !nextSet.has(fileId));
+    if (increment === 0 && decrement === 0) continue;
+    const file = await ctx.db.get("files", fileId);
+    if (!file) {
+      if (increment === 0) {
         console.error(`File ${fileId} not found when decrementing refcount`);
+        continue;
       }
+      throw new Error(`File ${fileId} not found when incrementing refcount`);
     }
-  }
-  for (const fileId of nextSet) {
-    if (!prevSet.has(fileId)) {
-      const file = await ctx.db.get("files", fileId);
-      if (file) {
-        await ctx.db.patch("files", fileId, {
-          refcount: file.refcount + 1,
-        });
-      } else {
-        throw new Error(`File ${fileId} not found when incrementing refcount`);
-      }
+    if (file.refcount < decrement) {
+      throw new Error(
+        `File ${fileId} refcount underflow: ${file.refcount} - ${decrement}`,
+      );
     }
+    const delta = increment - decrement;
+    await ctx.db.patch("files", fileId, {
+      refcount: file.refcount + delta,
+      lastTouchedAt: Date.now(),
+    });
   }
 }
 
@@ -153,22 +177,15 @@ export async function copyFileHandler(
   ctx: MutationCtx,
   args: { fileId: Id<"files"> },
 ) {
-  const file = await ctx.db.get("files", args.fileId);
-  if (!file) {
-    throw new Error("File not found");
-  }
-  await ctx.db.patch("files", args.fileId, {
-    refcount: file.refcount + 1,
-    lastTouchedAt: Date.now(),
-  });
+  await changeRefcount(ctx, [], [args.fileId]);
 }
 
 /**
  * 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.
+ * Files remain protected for 24 hours after registration or their last
+ * ownership change, so a file cannot be removed between registration and the
+ * transaction that saves its message or stream reference.
  */
 export const getFilesToDelete = query({
   args: {
@@ -177,7 +194,11 @@ export const getFilesToDelete = query({
   handler: async (ctx, args) => {
     const files = await paginator(ctx.db, schema)
       .query("files")
-      .withIndex("refcount", (q) => q.eq("refcount", 0))
+      .withIndex("refcount_lastTouchedAt", (q) =>
+        q
+          .eq("refcount", 0)
+          .lte("lastTouchedAt", Date.now() - FILE_CLEANUP_GRACE_MS),
+      )
       .paginate(args.paginationOpts);
     return files;
   },
@@ -210,6 +231,13 @@ export const deleteFiles = mutation({
             return null;
           }
         }
+        if (
+          !args.force &&
+          file.lastTouchedAt > Date.now() - FILE_CLEANUP_GRACE_MS
+        ) {
+          console.error(`File ${fileId} is still within the cleanup grace period, skipping...`);
+          return null;
+        }
         await ctx.db.delete("files", fileId);
         return fileId;
       }),

package/src/component/messages.test.ts

@@ -5,6 +5,7 @@ import { describe, expect, test, vi } from "vitest";
 import { api } from "./_generated/api.js";
 import type { Id } from "./_generated/dataModel.js";
 import { getMaxMessage } from "./messages.js";
+import { timeoutStreamHandler } from "./streams.js";
 import schema from "./schema.js";
 import { initConvexTest, modules } from "./setup.test.js";
 
@@ -924,4 +925,95 @@ describe("agent", () => {
     expect(state.messages[0].embeddingId).toBeUndefined();
     expect(state.embedding).toBeNull();
   });
+
+  test("transfers recovered file ownership and releases timed-out streams", async () => {
+    const t = initConvexTest();
+    const thread = await t.mutation(api.threads.createThread, {
+      userId: "stream-file-recovery",
+    });
+    const { messages } = await t.mutation(api.messages.addMessages, {
+      threadId: thread._id as Id<"threads">,
+      messages: [
+        { message: { role: "assistant", content: [] }, status: "pending" },
+      ],
+    });
+    const pending = messages[0]!;
+    const { fileId: recoveredFileId } = await t.mutation(api.files.addFile, {
+      storageId: "recovered-storage",
+      hash: "recovered-hash",
+      filename: "recovered.txt",
+    });
+    await t.mutation(api.files.copyFile, { fileId: recoveredFileId });
+    const recoveredStreamId = await t.run((ctx) =>
+      ctx.db.insert("streamingMessages", {
+        threadId: thread._id as Id<"threads">,
+        order: pending.order,
+        stepOrder: pending.stepOrder,
+        format: "UIMessageChunk",
+        state: { kind: "aborted", reason: "interrupted" },
+        fileRefs: [
+          { url: "https://files.example/recovered", fileId: recoveredFileId },
+        ],
+      }),
+    );
+    await t.run((ctx) =>
+      ctx.db.insert("streamDeltas", {
+        streamId: recoveredStreamId,
+        start: 0,
+        end: 3,
+        parts: [
+          { type: "start" },
+          {
+            type: "file",
+            url: "https://files.example/recovered",
+            mediaType: "text/plain",
+          },
+          { type: "finish" },
+        ],
+      }),
+    );
+    await t.mutation(api.messages.finalizeMessage, {
+      messageId: pending._id as Id<"messages">,
+      result: { status: "success" },
+    });
+    await expect(
+      t.query(api.files.get, { fileId: recoveredFileId }),
+    ).resolves.toMatchObject({
+      refcount: 1,
+    });
+    expect(
+      (await t.run((ctx) => ctx.db.get("streamingMessages", recoveredStreamId)))
+        ?.fileRefs,
+    ).toBeUndefined();
+
+    const { fileId: timedOutFileId } = await t.mutation(api.files.addFile, {
+      storageId: "timeout-storage",
+      hash: "timeout-hash",
+      filename: "timeout.txt",
+    });
+    await t.mutation(api.files.copyFile, { fileId: timedOutFileId });
+    const timedOutStreamId = await t.run((ctx) =>
+      ctx.db.insert("streamingMessages", {
+        threadId: thread._id as Id<"threads">,
+        order: pending.order + 1,
+        stepOrder: 0,
+        format: "UIMessageChunk",
+        state: { kind: "streaming", lastHeartbeat: Date.now() },
+        fileRefs: [
+          { url: "https://files.example/timeout", fileId: timedOutFileId },
+        ],
+      }),
+    );
+    await t.run((ctx) =>
+      timeoutStreamHandler(ctx, { streamId: timedOutStreamId }),
+    );
+    await t.mutation(api.streams.deleteStreamSync, {
+      streamId: timedOutStreamId,
+    });
+    await expect(
+      t.query(api.files.get, { fileId: timedOutFileId }),
+    ).resolves.toMatchObject({
+      refcount: 0,
+    });
+  });
 });

package/src/component/messages.ts

@@ -40,7 +40,11 @@ import {
   vVectorId,
 } from "./vector/tables.js";
 import { changeRefcount } from "./files.js";
-import { getStreamingMessagesWithMetadata, finishHandler } from "./streams.js";
+import {
+  getStreamingMessagesWithMetadata,
+  finishHandler,
+  releaseStreamFileOwnershipByIds,
+} from "./streams.js";
 import { partial } from "convex-helpers/validators";
 
 function publicMessage(message: Doc<"messages">): MessageDoc {
@@ -334,7 +338,10 @@ async function addMessagesHandler(
   // Atomically finish the stream if requested, preventing UI flickering
   // from separate mutations for message save and stream finish (issue #181).
   if (finishStreamId) {
-    await finishHandler(ctx, { streamId: finishStreamId });
+    await finishHandler(ctx, {
+      streamId: finishStreamId,
+    });
+    await releaseStreamFileOwnershipByIds(ctx, [finishStreamId]);
   }
   return { messages: toReturn.map(publicMessage) };
 }
@@ -409,7 +416,7 @@ export const finalizeMessage = mutation({
     }
     // See if we can add any in-progress data
     if (!message.message?.content.length) {
-      const { messages, materializationFailures } =
+      const { messages, materializationFailures, streamsToRelease } =
         await getStreamingMessagesWithMetadata(ctx, message, result);
       if (materializationFailures.length > 0) {
         console.error(
@@ -433,6 +440,7 @@ export const finalizeMessage = mutation({
           userId: message.userId,
           embeddings: undefined,
         });
+        await releaseStreamFileOwnershipByIds(ctx, streamsToRelease);
         return;
       }
     }

package/src/component/schema.ts

@@ -100,6 +100,15 @@ export const schema = defineSchema({
 
     threadId: v.id("threads"),
     order: v.number(),
+    /** Internal ownership sidecar for files materialized in stream deltas. */
+    fileRefs: v.optional(
+      v.array(
+        v.object({
+          url: v.string(),
+          fileId: v.id("files"),
+        }),
+      ),
+    ),
     /**
      * The step order of the first message in the stream.
      * If the stream ends up with both a tool call and a tool result,
@@ -117,7 +126,11 @@ 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(),
+        cleanupFnId: v.optional(v.id("_scheduled_functions")),
+      }),
     ),
   })
     // There should only be one per "order" index
@@ -158,7 +171,8 @@ export const schema = defineSchema({
     lastTouchedAt: v.number(),
   })
     .index("hash", ["hash"])
-    .index("refcount", ["refcount"]),
+    .index("refcount", ["refcount"])
+    .index("refcount_lastTouchedAt", ["refcount", "lastTouchedAt"]),
   ...vectorTables,
   // To authenticate playground usage
   // Delete a key to invalidate it

package/src/component/streams.test.ts

@@ -0,0 +1,109 @@
+/// <reference types="vite/client" />
+
+import { describe, expect, test } from "vitest";
+import { api } from "./_generated/api.js";
+import type { Id } from "./_generated/dataModel.js";
+import { initConvexTest } from "./setup.test.js";
+
+async function seedStream(t: ReturnType<typeof initConvexTest>) {
+  const thread = await t.mutation(api.threads.createThread, {
+    userId: "stream-files",
+  });
+  const threadId = thread._id as Id<"threads">;
+  const streamId = await t.mutation(api.streams.create, {
+    threadId,
+    order: 0,
+    stepOrder: 0,
+    format: "UIMessageChunk",
+  });
+  const { fileId } = await t.mutation(api.files.addFile, {
+    storageId: "stream-storage",
+    hash: "stream-hash",
+    filename: "stream.txt",
+  });
+  return { threadId, streamId, fileId };
+}
+
+describe("streams", () => {
+  test("stream file ownership flows from addDelta to the final message", async () => {
+    const t = initConvexTest();
+    const { threadId, streamId, fileId } = await seedStream(t);
+    const url = "https://files.example/stream";
+
+    await t.mutation(api.streams.addDelta, {
+      streamId,
+      start: 0,
+      end: 1,
+      parts: [{ type: "start" }],
+      fileRefs: [{ url, fileId }],
+    });
+    await expect(t.query(api.files.get, { fileId })).resolves.toMatchObject({
+      refcount: 1,
+    });
+
+    // Repeating the same URL/file pair keeps a single reference.
+    await t.mutation(api.streams.addDelta, {
+      streamId,
+      start: 1,
+      end: 2,
+      parts: [{ type: "text-delta", id: "text-1", delta: "hi" }],
+      fileRefs: [{ url, fileId }],
+    });
+    await expect(t.query(api.files.get, { fileId })).resolves.toMatchObject({
+      refcount: 1,
+    });
+
+    // A durable file can be referenced by more than one persisted URL.
+    await t.mutation(api.streams.addDelta, {
+      streamId,
+      start: 2,
+      end: 3,
+      parts: [{ type: "text-delta", id: "text-1", delta: "there" }],
+      fileRefs: [{ url: `${url}-alternate`, fileId }],
+    });
+    await expect(t.query(api.files.get, { fileId })).resolves.toMatchObject({
+      refcount: 1,
+    });
+
+    const { fileId: otherFileId } = await t.mutation(api.files.addFile, {
+      storageId: "other-storage",
+      hash: "other-hash",
+      filename: "other.txt",
+    });
+    await expect(
+      t.mutation(api.streams.addDelta, {
+        streamId,
+        start: 3,
+        end: 4,
+        parts: [{ type: "finish" }],
+        fileRefs: [{ url, fileId: otherFileId }],
+      }),
+    ).rejects.toThrow("Stream file URL maps to multiple files");
+    await expect(
+      t.query(api.files.get, { fileId: otherFileId }),
+    ).resolves.toMatchObject({ refcount: 0 });
+    expect(
+      (await t.run((ctx) => ctx.db.get("streamingMessages", streamId)))
+        ?.fileRefs,
+    ).toEqual([{ url, fileId }, { url: `${url}-alternate`, fileId }]);
+
+    await t.mutation(api.messages.addMessages, {
+      threadId,
+      messages: [
+        {
+          message: { role: "assistant", content: "done" },
+          fileIds: [fileId],
+        },
+      ],
+      finishStreamId: streamId,
+    });
+    await expect(t.query(api.files.get, { fileId })).resolves.toMatchObject({
+      refcount: 1,
+    });
+    const stream = await t.run((ctx) =>
+      ctx.db.get("streamingMessages", streamId),
+    );
+    expect(stream?.state.kind).toBe("finished");
+    expect(stream?.fileRefs).toBeUndefined();
+  });
+});