@convex-dev/agent 0.7.0 → 0.7.2

package/src/component/messages.ts

@@ -40,7 +40,12 @@ import {
   vVectorId,
 } from "./vector/tables.js";
 import { changeRefcount } from "./files.js";
-import { getStreamingMessagesWithMetadata, finishHandler } from "./streams.js";
+import {
+  getStreamingMessagesWithMetadata,
+  finishHandler,
+  releaseStreamFileOwnershipByIds,
+  abortStreamsAtOrder,
+} from "./streams.js";
 import { partial } from "convex-helpers/validators";
 
 function publicMessage(message: Doc<"messages">): MessageDoc {
@@ -61,21 +66,45 @@ export async function deleteMessage(
   }
 }
 
+/**
+ * Deleting a message strands any generation still writing to its order, which
+ * would otherwise only surface as a missing-parent failure when that generation
+ * finalizes. Aborting the stream lets the in-flight run stop on its own.
+ */
+async function abortStreamsForDeleted(
+  ctx: MutationCtx,
+  deleted: (Doc<"messages"> | null)[],
+) {
+  const seen = new Set<string>();
+  for (const message of deleted) {
+    if (!message) continue;
+    const key = `${message.threadId}:${message.order}`;
+    if (seen.has(key)) continue;
+    seen.add(key);
+    await abortStreamsAtOrder(ctx, {
+      threadId: message.threadId,
+      order: message.order,
+      reason: "Message deleted",
+    });
+  }
+}
+
 export const deleteByIds = mutation({
   args: { messageIds: v.array(v.id("messages")) },
   returns: v.array(v.id("messages")),
   handler: async (ctx, args) => {
-    const deletedMessageIds = await Promise.all(
+    const deleted = await Promise.all(
       args.messageIds.map(async (id) => {
         const message = await ctx.db.get("messages", id);
         if (message) {
           await deleteMessage(ctx, message);
-          return id;
+          return message;
         }
         return null;
       }),
     );
-    return deletedMessageIds.filter((id) => id !== null);
+    await abortStreamsForDeleted(ctx, deleted);
+    return deleted.filter((m) => m !== null).map((m) => m._id);
   },
 });
 
@@ -141,6 +170,7 @@ export const deleteByOrder = mutation({
       })
       .take(64);
     await Promise.all(messages.map((m) => deleteMessage(ctx, m)));
+    await abortStreamsForDeleted(ctx, messages);
     return {
       isDone: messages.length < 64,
       lastOrder: messages.at(-1)?.order,
@@ -153,6 +183,13 @@ const addMessagesArgs = {
   userId: v.optional(v.string()),
   threadId: v.id("threads"),
   promptMessageId: v.optional(v.id("messages")),
+  /**
+   * For saves that belong to a run anchored on promptMessageId: if that
+   * message is gone the run is obsolete, so abandon the save instead of
+   * throwing. A caller passing an id that never existed still gets an error.
+   */
+  abandonIfPromptMissing: v.optional(v.boolean()),
+  order: v.optional(v.union(v.number(), v.literal("next"))),
   agentName: v.optional(v.string()),
   messages: v.array(vMessageWithMetadataInternal),
   embeddings: v.optional(vMessageEmbeddingsWithDimension),
@@ -171,6 +208,15 @@ export const addMessages = mutation({
   handler: addMessagesHandler,
   returns: v.object({ messages: v.array(vMessageDoc) }),
 });
+
+function incrementMessagePosition(value: number, field: "order" | "stepOrder") {
+  assert(
+    Number.isSafeInteger(value) && value < Number.MAX_SAFE_INTEGER,
+    `${field} cannot be incremented past Number.MAX_SAFE_INTEGER`,
+  );
+  return value + 1;
+}
+
 async function addMessagesHandler(
   ctx: MutationCtx,
   args: ObjectType<typeof addMessagesArgs>,
@@ -189,12 +235,30 @@ async function addMessagesHandler(
     finishStreamId,
     messages,
     promptMessageId,
+    abandonIfPromptMissing,
+    order: requestedOrder,
     pendingMessageId,
     hideFromUserIdSearch,
     ...rest
   } = args;
   const promptMessage =
     promptMessageId && (await ctx.db.get("messages", promptMessageId));
+  assert(
+    requestedOrder === undefined ||
+      requestedOrder === "next" ||
+      (Number.isSafeInteger(requestedOrder) &&
+        requestedOrder >= 0 &&
+        requestedOrder < Number.MAX_SAFE_INTEGER),
+    "order must be a non-negative safe integer less than Number.MAX_SAFE_INTEGER",
+  );
+  assert(
+    requestedOrder === undefined || !promptMessageId,
+    "order and promptMessageId cannot both be provided",
+  );
+  assert(
+    requestedOrder === undefined || !pendingMessageId,
+    "order and pendingMessageId cannot both be provided",
+  );
   if (failPendingSteps) {
     assert(args.threadId, "threadId is required to fail pending steps");
     const pendingMessages = await ctx.db
@@ -224,7 +288,21 @@ async function addMessagesHandler(
   let order, stepOrder;
   let fail = false;
   let error: string | undefined;
-  if (promptMessageId) {
+  const startsAtNextOrder = requestedOrder === "next";
+  const explicitOrder =
+    typeof requestedOrder === "number" ? requestedOrder : undefined;
+  if (startsAtNextOrder) {
+    const maxMessage = await getMaxMessage(ctx, threadId);
+    order = incrementMessagePosition(maxMessage?.order ?? -1, "order");
+    stepOrder = -1;
+  } else if (explicitOrder !== undefined) {
+    order = explicitOrder;
+    const maxMessage = await getMaxMessage(ctx, threadId, order);
+    stepOrder = maxMessage?.stepOrder ?? -1;
+  } else if (promptMessageId) {
+    if (!promptMessage && abandonIfPromptMissing) {
+      return { messages: [] };
+    }
     assert(promptMessage, `Parent message ${promptMessageId} not found`);
     if (promptMessage.status === "failed") {
       fail = true;
@@ -285,8 +363,8 @@ async function addMessagesHandler(
       if (pendingMessage.status === "failed") {
         fail = true;
         error =
-          `Trying to update a message that failed: ${pendingMessageId}, ` +
-          `error: ${pendingMessage.error ?? error}`;
+          pendingMessage.error ??
+          `Trying to update a message that failed: ${pendingMessageId}`;
         messageDoc.status = "failed";
         messageDoc.error = error;
       }
@@ -305,20 +383,28 @@ async function addMessagesHandler(
       toReturn.push((await ctx.db.get("messages", pendingMessage._id))!);
       continue;
     }
-    if (message.message.role === "user") {
-      if (promptMessage && promptMessage.order === order) {
-        // see if there's a later message than the parent message order
+    if ((startsAtNextOrder || explicitOrder !== undefined) && i === 0) {
+      stepOrder = incrementMessagePosition(stepOrder, "stepOrder");
+    } else if (message.message.role === "user") {
+      if (
+        (explicitOrder !== undefined && order === explicitOrder) ||
+        (promptMessage && promptMessage.order === order)
+      ) {
+        // Avoid colliding with a later order when saving from an older one.
         const maxMessage = await getMaxMessage(ctx, threadId);
-        order = (maxMessage?.order ?? order) + 1;
+        order = incrementMessagePosition(
+          Math.max(maxMessage?.order ?? order, order),
+          "order",
+        );
       } else {
-        order++;
+        order = incrementMessagePosition(order, "order");
       }
       stepOrder = 0;
     } else {
       if (order < 0) {
         order = 0;
       }
-      stepOrder++;
+      stepOrder = incrementMessagePosition(stepOrder, "stepOrder");
     }
     const messageId = await ctx.db.insert("messages", {
       ...messageDoc,
@@ -334,7 +420,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 +498,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 +522,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();
+  });
+});

package/src/component/streams.ts

@@ -25,6 +25,7 @@ import {
   getPersistedUIMessageChunkParts,
   projectPersistedUIMessageChunks,
 } from "../streaming/materializePersistedUIMessageChunks.js";
+import { changeRefcount } from "./files.js";
 
 const SECOND = 1000;
 const MINUTE = 60 * SECOND;
@@ -35,9 +36,16 @@ const TIMEOUT_INTERVAL = 10 * MINUTE;
 const DELETE_STREAM_DELAY = MINUTE * 5; // 5 minutes
 
 const deltaValidator = schema.tables.streamDeltas.validator;
+const streamFileRefValidator = v.object({
+  url: v.string(),
+  fileId: v.id("files"),
+});
 
 export const addDelta = mutation({
-  args: deltaValidator,
+  args: {
+    ...deltaValidator.fields,
+    fileRefs: v.optional(v.array(streamFileRefValidator)),
+  },
   returns: v.boolean(),
   handler: async (ctx, args) => {
     const stream = await ctx.db.get("streamingMessages", args.streamId);
@@ -48,7 +56,30 @@ export const addDelta = mutation({
     if (stream.state.kind !== "streaming") {
       return false;
     }
-    await ctx.db.insert("streamDeltas", args);
+    const { fileRefs, ...delta } = args;
+    if (fileRefs?.length) {
+      const previous = stream.fileRefs ?? [];
+      // The persisted chunks refer to files by URL. Several URLs can still
+      // resolve to one durable file, but a URL must resolve to exactly one.
+      const refsByUrl = new Map(previous.map((ref) => [ref.url, ref] as const));
+      for (const ref of fileRefs) {
+        const existing = refsByUrl.get(ref.url);
+        if (existing && existing.fileId !== ref.fileId) {
+          throw new Error(`Stream file URL maps to multiple files: ${ref.url}`);
+        }
+        refsByUrl.set(ref.url, ref);
+      }
+      const next = [...refsByUrl.values()];
+      await changeRefcount(
+        ctx,
+        previous.map(({ fileId }) => fileId),
+        next.map(({ fileId }) => fileId),
+      );
+      await ctx.db.patch("streamingMessages", args.streamId, {
+        fileRefs: next,
+      });
+    }
+    await ctx.db.insert("streamDeltas", delta);
     await heartbeatStream(ctx, { streamId: args.streamId });
     return true;
   },
@@ -89,7 +120,10 @@ export const listDeltas = query({
 });
 
 export const create = mutation({
-  args: omit(schema.tables.streamingMessages.validator.fields, ["state"]),
+  args: omit(schema.tables.streamingMessages.validator.fields, [
+    "state",
+    "fileRefs",
+  ]),
   returns: v.id("streamingMessages"),
   handler: async (ctx, args) => {
     const state = { kind: "streaming" as const, lastHeartbeat: Date.now() };
@@ -163,24 +197,29 @@ function publicStreamMessage(m: Doc<"streamingMessages">): StreamMessage {
   };
 }
 
+export async function abortStreamsAtOrder(
+  ctx: MutationCtx,
+  args: { threadId: Id<"threads">; order: number; reason: string },
+) {
+  const streams = await ctx.db
+    .query("streamingMessages")
+    .withIndex("threadId_state_order_stepOrder", (q) =>
+      q
+        .eq("threadId", args.threadId)
+        .eq("state.kind", "streaming")
+        .eq("order", args.order),
+    )
+    .take(100);
+  for (const stream of streams) {
+    await abortById(ctx, { streamId: stream._id, reason: args.reason });
+  }
+  return streams.length > 0;
+}
+
 export const abortByOrder = mutation({
   args: { threadId: v.id("threads"), order: v.number(), reason: v.string() },
   returns: v.boolean(),
-  handler: async (ctx, args) => {
-    const streams = await ctx.db
-      .query("streamingMessages")
-      .withIndex("threadId_state_order_stepOrder", (q) =>
-        q
-          .eq("threadId", args.threadId)
-          .eq("state.kind", "streaming")
-          .eq("order", args.order),
-      )
-      .take(100);
-    for (const stream of streams) {
-      await abortById(ctx, { streamId: stream._id, reason: args.reason });
-    }
-    return streams.length > 0;
-  },
+  handler: abortStreamsAtOrder,
 });
 
 export const abort = mutation({
@@ -212,8 +251,9 @@ async function abortById(
     return false;
   }
   await cleanupTimeoutFn(ctx, stream);
+  const cleanupFnId = await scheduleStreamDeletion(ctx, args.streamId);
   await ctx.db.patch("streamingMessages", args.streamId, {
-    state: { kind: "aborted", reason: args.reason },
+    state: { kind: "aborted", reason: args.reason, cleanupFnId },
   });
   return true;
 }
@@ -264,11 +304,7 @@ export async function finishHandler(
     return;
   }
   await cleanupTimeoutFn(ctx, stream);
-  const cleanupFnId = await ctx.scheduler.runAfter(
-    DELETE_STREAM_DELAY,
-    api.streams.deleteStreamAsync,
-    { streamId: args.streamId },
-  );
+  const cleanupFnId = await scheduleStreamDeletion(ctx, args.streamId);
   await ctx.db.patch("streamingMessages", args.streamId, {
     state: { kind: "finished", endedAt: Date.now(), cleanupFnId },
   });
@@ -282,6 +318,29 @@ export const heartbeat = mutation({
   handler: heartbeatStream,
 });
 
+async function releaseStreamFileOwnership(
+  ctx: MutationCtx,
+  stream: Doc<"streamingMessages">,
+) {
+  if (!stream.fileRefs?.length) return;
+  await changeRefcount(
+    ctx,
+    stream.fileRefs.map(({ fileId }) => fileId),
+    [],
+  );
+}
+
+async function scheduleStreamDeletion(
+  ctx: MutationCtx,
+  streamId: Id<"streamingMessages">,
+) {
+  return ctx.scheduler.runAfter(
+    DELETE_STREAM_DELAY,
+    api.streams.deleteStreamAsync,
+    { streamId },
+  );
+}
+
 async function heartbeatStream(
   ctx: MutationCtx,
   args: { streamId: Id<"streamingMessages"> },
@@ -325,18 +384,24 @@ async function heartbeatStream(
 export const timeoutStream = internalMutation({
   args: { streamId: v.id("streamingMessages") },
   returns: v.null(),
-  handler: async (ctx, args) => {
-    const stream = await ctx.db.get("streamingMessages", args.streamId);
-    if (!stream || stream.state.kind !== "streaming") {
-      console.warn("Stream not found", args.streamId);
-      return;
-    }
-    await ctx.db.patch("streamingMessages", args.streamId, {
-      state: { kind: "aborted", reason: "timeout" },
-    });
-  },
+  handler: timeoutStreamHandler,
 });
 
+export async function timeoutStreamHandler(
+  ctx: MutationCtx,
+  args: { streamId: Id<"streamingMessages"> },
+) {
+  const stream = await ctx.db.get("streamingMessages", args.streamId);
+  if (!stream || stream.state.kind !== "streaming") {
+    console.warn("Stream not found", args.streamId);
+    return;
+  }
+  const cleanupFnId = await scheduleStreamDeletion(ctx, args.streamId);
+  await ctx.db.patch("streamingMessages", args.streamId, {
+    state: { kind: "aborted", reason: "timeout", cleanupFnId },
+  });
+}
+
 async function deletePageForStreamId(
   ctx: MutationCtx,
   args: { streamId: Id<"streamingMessages">; cursor?: string },
@@ -354,8 +419,12 @@ async function deletePageForStreamId(
   if (deltas.isDone) {
     const stream = await ctx.db.get("streamingMessages", args.streamId);
     if (stream) {
+      await releaseStreamFileOwnership(ctx, stream);
       await cleanupTimeoutFn(ctx, stream);
-      if (stream.state.kind === "finished" && stream.state.cleanupFnId) {
+      if (
+        (stream.state.kind === "finished" || stream.state.kind === "aborted") &&
+        stream.state.cleanupFnId
+      ) {
         const scheduledFunction = await ctx.db.system.get(
           "_scheduled_functions",
           stream.state.cleanupFnId,
@@ -370,6 +439,18 @@ async function deletePageForStreamId(
   return deltas;
 }
 
+export async function releaseStreamFileOwnershipByIds(
+  ctx: MutationCtx,
+  streamIds: Id<"streamingMessages">[],
+) {
+  for (const streamId of new Set(streamIds)) {
+    const stream = await ctx.db.get("streamingMessages", streamId);
+    if (!stream?.fileRefs?.length) continue;
+    await releaseStreamFileOwnership(ctx, stream);
+    await ctx.db.patch("streamingMessages", streamId, { fileRefs: undefined });
+  }
+}
+
 export async function deleteStreamsPageForThreadId(
   ctx: MutationCtx,
   args: { threadId: Id<"threads">; streamOrder?: number; deltaCursor?: string },
@@ -544,6 +625,7 @@ export async function getStreamingMessagesWithMetadata(
     streamId: Id<"streamingMessages">;
     reason: string;
   }>;
+  streamsToRelease: Id<"streamingMessages">[];
 }> {
   // See if there are any streaming messages for this order
   const streamingMessages = await getStreamingMessages(
@@ -570,8 +652,10 @@ export async function getStreamingMessagesWithMetadata(
             streamMessage,
             parts,
             metadata,
+            streamingMessage.fileRefs,
           ).slice(numToSkip),
           failure: undefined,
+          streamToRelease: streamingMessage._id,
         };
       } catch (error) {
         return {
@@ -580,6 +664,7 @@ export async function getStreamingMessagesWithMetadata(
             streamId: streamingMessage._id,
             reason: error instanceof Error ? error.message : String(error),
           },
+          streamToRelease: undefined,
         };
       }
     }),
@@ -589,5 +674,8 @@ export async function getStreamingMessagesWithMetadata(
     materializationFailures: materializedStreams.flatMap(({ failure }) =>
       failure ? [failure] : [],
     ),
+    streamsToRelease: materializedStreams.flatMap(({ streamToRelease }) =>
+      streamToRelease ? [streamToRelease] : [],
+    ),
   };
 }

package/src/component/vector/tables.ts

@@ -68,7 +68,7 @@ export type VectorSchema = SchemaDefinition<
 >;
 
 export const VectorDimensions = [
-  128, 256, 512, 768, 1024, 1408, 1536, 2048, 3072, 4096,
+  128, 256, 512, 768, 1024, 1408, 1536, 2048, 2560, 3072, 4096,
 ] as const;
 export function validateVectorDimension(
   dimension: number,