@convex-dev/agent 0.7.0 → 0.7.1

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() };
@@ -212,8 +246,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 +299,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 +313,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 +379,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 +414,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 +434,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 +620,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 +647,10 @@ export async function getStreamingMessagesWithMetadata(
             streamMessage,
             parts,
             metadata,
+            streamingMessage.fileRefs,
           ).slice(numToSkip),
           failure: undefined,
+          streamToRelease: streamingMessage._id,
         };
       } catch (error) {
         return {
@@ -580,6 +659,7 @@ export async function getStreamingMessagesWithMetadata(
             streamId: streamingMessage._id,
             reason: error instanceof Error ? error.message : String(error),
           },
+          streamToRelease: undefined,
         };
       }
     }),
@@ -589,5 +669,8 @@ export async function getStreamingMessagesWithMetadata(
     materializationFailures: materializedStreams.flatMap(({ failure }) =>
       failure ? [failure] : [],
     ),
+    streamsToRelease: materializedStreams.flatMap(({ streamToRelease }) =>
+      streamToRelease ? [streamToRelease] : [],
+    ),
   };
 }

package/src/streaming/materializePersistedUIMessageChunks.test.ts

@@ -51,6 +51,39 @@ describe("projectPersistedUIMessageChunks", () => {
     expect(validate(vMessageWithMetadataInternal, actual[1])).toBe(true);
   });
 
+  it("attaches materialized canonical tool-result files during recovery", () => {
+    const url = "https://files.example/tool-result";
+    const actual = projectPersistedUIMessageChunks(
+      stream,
+      [
+        {
+          type: "tool-input-available",
+          toolCallId: "call-1",
+          toolName: "render",
+          input: {},
+        },
+        {
+          type: "tool-output-available",
+          toolCallId: "call-1",
+          output: {
+            type: "content",
+            value: [
+              {
+                type: "file",
+                data: { type: "url", url },
+                mediaType: "application/octet-stream",
+              },
+            ],
+          },
+        },
+      ],
+      { status: "success" },
+      [{ url, fileId: "file-1" }],
+    );
+
+    expect(actual[1]).toMatchObject({ fileIds: ["file-1"] });
+  });
+
   it("keeps persisted sources on the step that produced them", () => {
     const chunks = [
       { type: "start-step" },

package/src/streaming/materializePersistedUIMessageChunks.ts

@@ -33,6 +33,7 @@ export function projectPersistedUIMessageChunks(
   stream: StreamMessage,
   chunks: readonly unknown[],
   metadata: PersistedStreamMetadata,
+  fileRefs: readonly { url: string; fileId: string }[] = [],
 ): MessageWithMetadataInternal[] {
   if (stream.format !== "UIMessageChunk") {
     throw new Error(
@@ -60,6 +61,7 @@ export function projectPersistedUIMessageChunks(
     reduced.state.parts,
     stream,
     metadata,
+    fileRefs,
   );
 }
 
@@ -91,6 +93,7 @@ export function projectPersistedUIMessageChunkParts(
   parts: PersistedUIMessagePart[],
   stream: StreamMessage,
   metadata: PersistedStreamMetadata,
+  fileRefs: readonly { url: string; fileId: string }[] = [],
 ): MessageWithMetadataInternal[] {
   const blocks: PersistedUIMessagePart[][] = [];
   let block: PersistedUIMessagePart[] = [];
@@ -267,6 +270,7 @@ export function projectPersistedUIMessageChunkParts(
     const hasToolCall =
       message.role === "tool" ||
       content.some((part) => part.type === "tool-call");
+    const fileIds = referencedFileIds(message, fileRefs);
     return {
       message,
       status: metadata.status,
@@ -284,6 +288,9 @@ export function projectPersistedUIMessageChunkParts(
         )
         .map((part) => part.text)
         .join(" "),
+      ...(fileIds.length > 0
+        ? { fileIds: fileIds as MessageWithMetadataInternal["fileIds"] }
+        : {}),
       ...(metadata.error !== undefined ? { error: metadata.error } : {}),
     } satisfies MessageWithMetadataInternal;
   });
@@ -316,6 +323,36 @@ function projectSources(parts: PersistedUIMessagePart[]) {
     );
 }
 
+function referencedFileIds(
+  message: Message,
+  fileRefs: readonly { url: string; fileId: string }[],
+) {
+  if (typeof message.content === "string") return [];
+  const urls = new Set<string>();
+  for (const part of message.content) {
+    if (part.type === "file") {
+      if (typeof part.data === "string") urls.add(part.data);
+    } else if (part.type === "reasoning-file") {
+      if ("url" in part && part.url) urls.add(part.url);
+    } else if (part.type === "tool-result" && part.output?.type === "content") {
+      for (const outputPart of part.output.value) {
+        if (
+          outputPart.type === "file" &&
+          outputPart.data.type === "url" &&
+          typeof outputPart.data.url === "string"
+        ) {
+          urls.add(outputPart.data.url);
+        }
+      }
+    }
+  }
+  return [
+    ...new Set(
+      fileRefs.filter((ref) => urls.has(ref.url)).map((ref) => ref.fileId),
+    ),
+  ];
+}
+
 function toolResult(
   part: PersistedToolPart,
   mode: "normal" | "error-text" | "error-json" | "execution-denied",

package/src/vercel/UIMessages.ts

@@ -344,7 +344,7 @@ function createUserUIMessage<
 
   const parts: UIMessage<METADATA, DATA_PARTS, TOOLS>["parts"] = [];
   if (text && !nonStringContent.length) {
-    parts.push({ type: "text", text });
+    parts.push({ type: "text", text, ...partCommon });
   }
   for (const contentPart of nonStringContent) {
     switch (contentPart.type) {

package/src/vercel/client/files.test.ts

@@ -0,0 +1,56 @@
+/// <reference types="vite/client" />
+
+import { describe, expect, test } from "vitest";
+import { storeFile } from "./files.js";
+
+describe("storeFile", () => {
+  test("throws a clear error when a reused file is missing from storage", async () => {
+    const ctx = {
+      runAction: async () => null,
+      runMutation: async () => ({
+        fileId: "existing-file",
+        storageId: "existing-storage",
+      }),
+      storage: {
+        getUrl: async () => null,
+      },
+    } as unknown as Parameters<typeof storeFile>[0];
+    const component = {
+      files: { useExistingFile: {}, addFile: {} },
+    } as unknown as Parameters<typeof storeFile>[1];
+
+    await expect(storeFile(ctx, component, new Blob(["x"]))).rejects.toThrow(
+      "File not found in storage: existing-storage",
+    );
+  });
+
+  test("cleans its losing blob before an existing URL read fails", async () => {
+    const deleted: string[] = [];
+    let mutationCount = 0;
+    const ctx = {
+      runAction: async () => null,
+      runMutation: async () => {
+        mutationCount++;
+        return mutationCount === 1
+          ? null
+          : { fileId: "existing-file", storageId: "existing-storage" };
+      },
+      storage: {
+        store: async () => "new-storage",
+        getMetadata: async () => null,
+        getUrl: async () => null,
+        delete: async (storageId: string) => {
+          deleted.push(storageId);
+        },
+      },
+    } as unknown as Parameters<typeof storeFile>[0];
+    const component = {
+      files: { useExistingFile: {}, addFile: {} },
+    } as unknown as Parameters<typeof storeFile>[1];
+
+    await expect(storeFile(ctx, component, new Blob(["x"]))).rejects.toThrow(
+      "File not found in storage: existing-storage",
+    );
+    expect(deleted).toEqual(["new-storage"]);
+  });
+});

package/src/vercel/client/files.ts

@@ -67,9 +67,13 @@ export async function storeFile(
   const reused = await ctx.runMutation(component.files.useExistingFile, {
     hash,
     filename,
+    mediaType: blob.type || undefined,
   });
   if (reused) {
-    const url = (await ctx.storage.getUrl(reused.storageId))!;
+    const url = await ctx.storage.getUrl(reused.storageId);
+    if (!url) {
+      throw new Error(`File not found in storage: ${reused.storageId}`);
+    }
     return {
       ...getParts(url, blob.type, filename),
       file: {
@@ -82,35 +86,56 @@ export async function storeFile(
     };
   }
   const newStorageId = await ctx.storage.store(blob);
-  if (sha256) {
-    const metadata = await ctx.storage.getMetadata(newStorageId);
-    if (metadata?.sha256 !== sha256) {
-      throw new Error("Hash mismatch: " + metadata?.sha256 + " != " + sha256);
-    }
-  }
-  const { fileId, storageId } = await ctx.runMutation(component.files.addFile, {
-    storageId: newStorageId,
-    hash,
-    filename,
-    mediaType: blob.type,
-  });
-  const url = (await ctx.storage.getUrl(storageId as Id<"_storage">))!;
-  if (storageId !== newStorageId) {
-    // We're re-using another file's storageId
-    // Because we try to reuse the file above, this should be very very rare
-    // and only in the case of racing to check then store the file.
+  let newStorageRegistered = false;
+  let cleanupAttempted = false;
+  const cleanupNewStorage = async () => {
+    if (cleanupAttempted) return;
+    cleanupAttempted = true;
     await ctx.storage.delete(newStorageId);
-  }
-  return {
-    ...getParts(url, blob.type, filename),
-    file: {
-      url,
-      fileId,
-      storageId: storageId as Id<"_storage">,
-      hash,
-      filename,
-    },
   };
+  try {
+    if (sha256) {
+      const metadata = await ctx.storage.getMetadata(newStorageId);
+      if (metadata?.sha256 !== sha256) {
+        throw new Error("Hash mismatch: " + metadata?.sha256 + " != " + sha256);
+      }
+    }
+    const { fileId, storageId } = await ctx.runMutation(
+      component.files.addFile,
+      {
+        storageId: newStorageId,
+        hash,
+        filename,
+        mediaType: blob.type,
+      },
+    );
+    newStorageRegistered = storageId === newStorageId;
+    // A competing request can win after useExistingFile but before addFile.
+    // Delete our losing blob before the existing file's URL is read, so a
+    // failed getUrl cannot leave the raw object orphaned.
+    if (!newStorageRegistered) {
+      await cleanupNewStorage();
+    }
+    const url = await ctx.storage.getUrl(storageId as Id<"_storage">);
+    if (!url) {
+      throw new Error(`File not found in storage: ${storageId}`);
+    }
+    return {
+      ...getParts(url, blob.type, filename),
+      file: {
+        url,
+        fileId,
+        storageId: storageId as Id<"_storage">,
+        hash,
+        filename,
+      },
+    };
+  } catch (error) {
+    if (!newStorageRegistered && !cleanupAttempted) {
+      await cleanupNewStorage().catch(() => {});
+    }
+    throw error;
+  }
 }
 
 /**

package/src/vercel/client/streamText.ts

@@ -25,6 +25,7 @@ import { startGeneration } from "./start.js";
 import type { Agent } from "../index.js";
 import { getModelName, getProviderName } from "../../shared.js";
 import { errorToString, willContinue } from "./utils.js";
+import { materializeUIMessageChunkFiles } from "../fileMaterialization.js";
 
 /** Finish every abort cleanup path before surfacing an internal failure. */
 export async function runAbortCleanup(cleanup: {
@@ -152,6 +153,8 @@ export async function streamText<
                 : undefined,
             onAsyncAbort: call.fail,
             compress: compressUIMessageChunks,
+            materialize: (parts) =>
+              materializeUIMessageChunkFiles(ctx, component, parts),
             abortSignal: args.abortSignal,
           },
           {

package/src/vercel/client/streaming.test.ts

@@ -249,6 +249,49 @@ describe("DeltaStreamer", () => {
     );
   });
 
+  test("waits for signal cleanup when the source throws", async () => {
+    let resolveDelta!: (value: boolean) => void;
+    const deltaWrite = new Promise<boolean>((resolve) => {
+      resolveDelta = resolve;
+    });
+    const runMutation = vi
+      .fn()
+      .mockResolvedValueOnce("stream-1")
+      .mockImplementationOnce(() => deltaWrite)
+      .mockResolvedValueOnce(undefined);
+    const abortController = new AbortController();
+    const streamer = new DeltaStreamer<string>(
+      components.agent,
+      { runMutation } as unknown as MutationCtx,
+      { ...defaultTestOptions, abortSignal: abortController.signal },
+      { ...testMetadata, threadId },
+    );
+    const sourceError = new Error("provider aborted");
+    const source = {
+      async *[Symbol.asyncIterator]() {
+        yield "chunk";
+        abortController.abort();
+        throw sourceError;
+      },
+    } as unknown as Parameters<typeof streamer.consumeStream>[0];
+
+    const consuming = streamer.consumeStream(source);
+    let settled = false;
+    void consuming.catch(() => {
+      settled = true;
+    });
+    await vi.waitFor(() => expect(runMutation).toHaveBeenCalledTimes(2));
+    expect(settled).toBe(false);
+
+    resolveDelta(true);
+    await expect(consuming).rejects.toBe(sourceError);
+    expect(runMutation).toHaveBeenNthCalledWith(
+      3,
+      components.agent.streams.abort,
+      { streamId: "stream-1", reason: "abortSignal" },
+    );
+  });
+
   test("aborts the component stream when a delta write fails", async () => {
     const runMutation = vi
       .fn()
@@ -279,6 +322,34 @@ describe("DeltaStreamer", () => {
     );
   });
 
+  test("aborts the component stream when file materialization fails", async () => {
+    const materializationFailure = new Error("storage failed");
+    const runMutation = vi
+      .fn()
+      .mockResolvedValueOnce("stream-1")
+      .mockResolvedValueOnce(undefined);
+    const streamer = new DeltaStreamer<string>(
+      components.agent,
+      { runMutation } as unknown as MutationCtx,
+      {
+        ...defaultTestOptions,
+        onAsyncAbort: async () => {},
+        materialize: async () => {
+          throw materializationFailure;
+        },
+      },
+      { ...testMetadata, threadId },
+    );
+
+    await streamer.addParts(["A"]);
+    await expect(streamer.finish()).resolves.toBeUndefined();
+    expect(runMutation).toHaveBeenNthCalledWith(
+      2,
+      components.agent.streams.abort,
+      { streamId: "stream-1", reason: "storage failed" },
+    );
+  });
+
   test("surfaces pending-message cleanup failure after aborting the stream", async () => {
     const pendingMessageFailure = new Error("pending message cleanup failed");
     let resolveComponentAbort!: () => void;