@convex-dev/agent 0.7.1 → 0.7.3

package/src/vercel/client/streamText.ts

@@ -9,6 +9,7 @@ import type { Context } from "@ai-sdk/provider-utils";
 import { streamText as streamTextAi } from "ai";
 import {
   compressUIMessageChunks,
+  DEFAULT_STREAMING_OPTIONS,
   DeltaStreamer,
   mergeTransforms,
   type StreamingOptions,
@@ -27,19 +28,22 @@ 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: {
+export async function runStreamCleanup(cleanup: {
   failCall: () => Promise<void>;
   failStreamer: () => Promise<void>;
   onAbort?: () => PromiseLike<void> | void;
 }): Promise<void> {
   const results = await Promise.allSettled([
-    cleanup.failCall(),
-    cleanup.failStreamer(),
+    Promise.resolve().then(() => cleanup.failCall()),
+    Promise.resolve().then(() => cleanup.failStreamer()),
   ]);
-  await cleanup.onAbort?.();
-  const failure = results.find((result) => result.status === "rejected");
-  if (failure) throw failure.reason;
+  const [abortResult] = await Promise.allSettled([
+    Promise.resolve().then(() => cleanup.onAbort?.()),
+  ]);
+  const failure = [...results, abortResult].find(
+    (result) => result.status === "rejected",
+  );
+  if (failure?.status === "rejected") throw failure.reason;
 }
 
 /**
@@ -103,15 +107,11 @@ export async function streamText<
       Tools,
       object,
       RUNTIME_CONTEXT
-    >(
-      ctx,
-      component,
-      streamTextArgs,
-      options,
-      "streamText",
-    );
+    >(ctx, component, streamTextArgs, options, "streamText");
 
   const steps: StepResult<Tools, RUNTIME_CONTEXT>[] = [];
+  let firstStreamError: string | undefined;
+  let streamCleanupFailure: { error: unknown } | undefined;
   let initialResponseMessages: ModelMessage[] = [];
   let initialResponseMessagesSaved = false;
   const responseMessagesForStep = (
@@ -135,6 +135,7 @@ export async function streamText<
   // When false (saveStreamDeltas.returnImmediately === true), we cannot
   // defer the final-step save to a post-await block — the function has
   // already returned by the time onStepFinish fires. See issue #265.
+  const savesMessages = options?.storageOptions?.saveMessages !== "none";
   const willAwaitStream =
     Boolean(threadId) &&
     (options.saveStreamDeltas === true ||
@@ -156,6 +157,10 @@ export async function streamText<
             materialize: (parts) =>
               materializeUIMessageChunkFiles(ctx, component, parts),
             abortSignal: args.abortSignal,
+            // The message save finishes the stream row atomically (issue
+            // #181) — but only when there is a save. With saveMessages set to
+            // "none" nothing does, so the streamer keeps finish ownership.
+            finishHandledExternally: savesMessages,
           },
           {
             threadId,
@@ -180,15 +185,26 @@ export async function streamText<
     ),
     onError: async (error) => {
       console.error("onError", error);
-      await call.fail(errorToString(error.error));
-      await streamer?.fail(errorToString(error.error));
+      const reason = (firstStreamError ??= errorToString(error.error));
+      try {
+        await runStreamCleanup({
+          failCall: () => call.fail(reason),
+          failStreamer: async () => streamer?.fail(reason),
+        });
+      } catch (cleanupError) {
+        streamCleanupFailure ??= { error: cleanupError };
+        console.error("Failed to clean up errored stream:", cleanupError);
+      }
       return streamTextArgs.onError?.(error);
     },
     onAbort: async (event) => {
+      const providerTriggeredAbort =
+        firstStreamError !== undefined && !args.abortSignal?.aborted;
+      if (providerTriggeredAbort) return;
       const reason = args.abortSignal?.reason
         ? errorToString(args.abortSignal.reason)
         : "streamText aborted";
-      await runAbortCleanup({
+      await runStreamCleanup({
         failCall: () => call.fail(reason),
         failStreamer: async () => streamer?.fail(reason),
         onAbort: () => streamTextArgs.onAbort?.(event),
@@ -215,10 +231,13 @@ export async function streamText<
       const createPendingMessage = await willContinue(steps, args.stopWhen);
       if (!createPendingMessage && streamer) {
         // Final step with streaming enabled.
-        streamer.markFinishedExternally();
         if (willAwaitStream) {
           // We're about to `await stream` below — defer the save so it
-          // happens atomically with stream finish (issue #181).
+          // happens atomically with stream finish (issue #181). Don't touch
+          // the streamer here: the stream-level `finish` chunk is emitted
+          // after this callback, so the row has to stay `streaming` and keep
+          // accepting parts until consumeStream reaches EOF, which is the only
+          // point where everything has actually been handed over.
           pendingFinalStep = {
             step,
             responseMessages: responseMessagesForStep(step),
@@ -226,14 +245,26 @@ export async function streamText<
         } else {
           // returnImmediately path: streamText is about to return without
           // awaiting consumption, so the deferred-save block below won't
-          // see this step. Save inline now (issue #265).
-          const finishStreamId = await streamer.getOrCreateStreamId();
-          await call.save(
-            { step, responseMessages: responseMessagesForStep(step) },
-            false,
-            finishStreamId,
-          );
-          initialResponseMessagesSaved = true;
+          // see this step. Save inline now (issue #265). Nothing awaits the
+          // stream here, so this is the last moment we can drain deltas, see
+          // flushAndStopAccepting for the window that leaves.
+          await streamer.flushAndStopAccepting();
+          const finishStreamId = await streamer.getOrCreateStreamId({
+            ifAborted: "returnUndefined",
+          });
+          if (finishStreamId) {
+            await call.save(
+              { step, responseMessages: responseMessagesForStep(step) },
+              false,
+              finishStreamId,
+            );
+            // The save finishes the row only when it stores messages. With
+            // saveMessages "none" nothing else will, so do it here.
+            if (!savesMessages) {
+              await streamer.finish();
+            }
+            initialResponseMessagesSaved = true;
+          }
         }
       } else {
         await call.save(
@@ -248,7 +279,13 @@ export async function streamText<
     typeof streamTextAi<Tools, RUNTIME_CONTEXT, OUTPUT>
   >[0]) as StreamTextResult<Tools, RUNTIME_CONTEXT, OUTPUT>;
   const stream = streamer?.consumeStream(
-    result.toUIMessageStream<AIUIMessage<Tools>>(),
+    result.toUIMessageStream<AIUIMessage<Tools>>({
+      sendSources:
+        typeof options.saveStreamDeltas === "object"
+          ? (options.saveStreamDeltas.sendSources ??
+            DEFAULT_STREAMING_OPTIONS.sendSources)
+          : DEFAULT_STREAMING_OPTIONS.sendSources,
+    }),
   );
   if (willAwaitStream) {
     try {
@@ -258,8 +295,11 @@ export async function streamText<
       // If the stream errored (e.g. onStepFinish threw), the DeltaStreamer's
       // finish() was never called, leaving the streaming message stuck in
       // "streaming" state. Clean it up by marking it as aborted.
-      await streamer?.fail(e instanceof Error ? e.message : String(e));
-      // Save the deferred final step if it was already generated but not yet persisted
+      try {
+        await streamer?.fail(errorToString(e));
+      } catch (cleanupError) {
+        streamCleanupFailure ??= { error: cleanupError };
+      }
       if (pendingFinalStep) {
         try {
           await call.save(pendingFinalStep, false);
@@ -272,10 +312,22 @@ export async function streamText<
     }
   }
 
+  if (streamCleanupFailure) throw streamCleanupFailure.error;
+
   // If we deferred the final step save, do it now with atomic stream finish.
   if (pendingFinalStep && streamer) {
-    const finishStreamId = await streamer.getOrCreateStreamId();
-    await call.save(pendingFinalStep, false, finishStreamId);
+    const finishStreamId = await streamer.getOrCreateStreamId({
+      ifAborted: "returnUndefined",
+    });
+    if (finishStreamId) {
+      await call.save(pendingFinalStep, false, finishStreamId);
+    }
+    pendingFinalStep = undefined;
+  } else if (willAwaitStream && streamer) {
+    // No final step was deferred (e.g. the generation produced none), so no
+    // save will finish the stream. The streamer doesn't finish itself, so do
+    // it here rather than leaving the row to time out.
+    await streamer.finish();
   }
   const metadata: GenerationOutputMetadata = {
     promptMessageId,

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

@@ -190,17 +190,16 @@ describe("HTTP Streaming Initiation", () => {
     });
   });
 
-  test("markFinishedExternally prevents consumeStream from calling finish", async () => {
+  test("finishHandledExternally prevents consumeStream from calling finish", async () => {
     await t.run(async (ctx) => {
       const streamer = new DeltaStreamer(
         components.agent,
         ctx,
-        { ...defaultTestOptions },
+        { ...defaultTestOptions, finishHandledExternally: true },
         { ...testMetadata, threadId },
       );
 
       await streamer.getStreamId();
-      streamer.markFinishedExternally();
 
       const result = streamText({
         model: mockModel({
@@ -217,6 +216,44 @@ describe("HTTP Streaming Initiation", () => {
         { threadId, statuses: ["streaming"] },
       );
       expect(streamingStreams).toHaveLength(1);
+
+      // ...but the parts still made it into deltas, since the row kept
+      // accepting until end-of-stream.
+      const deltas = await ctx.runQuery(components.agent.streams.listDeltas, {
+        threadId,
+        cursors: [{ streamId: streamer.streamId!, cursor: 0 }],
+      });
+      const types = deltas.flatMap((d) => d.parts).map((p) => p.type);
+      expect(types.at(0)).toBe("start");
+      expect(types.at(-1)).toBe("finish");
+    });
+  });
+
+  test("flushAndStopAccepting drops later parts but keeps the row streaming", async () => {
+    await t.run(async (ctx) => {
+      const streamer = new DeltaStreamer(
+        components.agent,
+        ctx,
+        { ...defaultTestOptions, finishHandledExternally: true },
+        { ...testMetadata, threadId },
+      );
+
+      await streamer.addParts([{ type: "start" }]);
+      await streamer.flushAndStopAccepting();
+      await streamer.addParts([{ type: "finish" }]);
+
+      const deltas = await ctx.runQuery(components.agent.streams.listDeltas, {
+        threadId,
+        cursors: [{ streamId: streamer.streamId!, cursor: 0 }],
+      });
+      const types = deltas.flatMap((d) => d.parts).map((p) => p.type);
+      expect(types).toEqual(["start"]);
+
+      const streamingStreams = await ctx.runQuery(
+        components.agent.streams.list,
+        { threadId, statuses: ["streaming"] },
+      );
+      expect(streamingStreams).toHaveLength(1);
     });
   });
 });

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

@@ -5,10 +5,7 @@ import type { GenericSchema, SchemaDefinition } from "convex/server";
 import { streamText } from "ai";
 import { components, initConvexTest } from "./setup.test.js";
 import { mockModel } from "./mockModel.js";
-import {
-  compressUIMessageChunks,
-  DeltaStreamer,
-} from "./streaming.js";
+import { compressUIMessageChunks, DeltaStreamer } from "./streaming.js";
 import { getParts } from "../deltas.js";
 import type { TestConvex } from "convex-test";
 
@@ -200,10 +197,71 @@ describe("DeltaStreamer", () => {
       expect(streamer.abortController.signal.aborted).toBe(true);
       await streamer.addParts(["ignored"]);
       expect(streamer.streamId).toBeUndefined();
+      await expect(
+        streamer.getOrCreateStreamId({ ifAborted: "returnUndefined" }),
+      ).resolves.toBeUndefined();
       await expect(streamer.getOrCreateStreamId()).rejects.toThrow(
         "Cannot create a stream after it has been aborted",
       );
+      await expect(streamer.getStreamId()).rejects.toThrow(
+        "Cannot create a stream after it has been aborted",
+      );
+    });
+  });
+
+  test("preserves the public throwing behavior after an existing stream aborts", async () => {
+    await t.run(async (ctx) => {
+      const streamer = new DeltaStreamer<string>(
+        components.agent,
+        ctx,
+        { ...defaultTestOptions },
+        { ...testMetadata, threadId },
+      );
+      const streamId = await streamer.getStreamId();
+
+      await streamer.fail("provider error");
+
+      expect(streamer.streamId).toBe(streamId);
+      await expect(streamer.getStreamId()).rejects.toThrow(
+        "Cannot create a stream after it has been aborted",
+      );
+      await expect(streamer.getOrCreateStreamId()).rejects.toThrow(
+        "Cannot create a stream after it has been aborted",
+      );
+      await expect(
+        streamer.getOrCreateStreamId({ ifAborted: "returnUndefined" }),
+      ).resolves.toBeUndefined();
+    });
+  });
+
+  test("does not return an ID when abort wins during stream creation", async () => {
+    let resolveCreate!: (streamId: string) => void;
+    const creatingStream = new Promise<string>((resolve) => {
+      resolveCreate = resolve;
+    });
+    const runMutation = vi
+      .fn()
+      .mockImplementationOnce(() => creatingStream)
+      .mockResolvedValueOnce(true);
+    const streamer = new DeltaStreamer<string>(
+      components.agent,
+      { runMutation } as unknown as MutationCtx,
+      { ...defaultTestOptions },
+      { ...testMetadata, threadId },
+    );
+
+    const getting = streamer.getOrCreateStreamId({
+      ifAborted: "returnUndefined",
     });
+    const failing = streamer.fail("provider error");
+    resolveCreate("stream-1");
+
+    await failing;
+    await expect(getting).resolves.toBeUndefined();
+    expect(streamer.streamId).toBe("stream-1");
+    await expect(streamer.getStreamId()).rejects.toThrow(
+      "Cannot create a stream after it has been aborted",
+    );
   });
 
   test("shares signal and fail cleanup while stream creation is in flight", async () => {
@@ -293,10 +351,13 @@ describe("DeltaStreamer", () => {
   });
 
   test("aborts the component stream when a delta write fails", async () => {
+    const deltaFailure = {
+      error: { code: "provider_disconnected", message: "Provider dropped" },
+    };
     const runMutation = vi
       .fn()
       .mockResolvedValueOnce("stream-1")
-      .mockRejectedValueOnce(new Error("delta failed"))
+      .mockRejectedValueOnce(deltaFailure)
       .mockResolvedValueOnce(undefined);
     let abortReason: string | undefined;
     const streamer = new DeltaStreamer<string>(
@@ -314,11 +375,14 @@ describe("DeltaStreamer", () => {
     await streamer.addParts(["A"]);
     await streamer.finish();
 
-    expect(abortReason).toBe("delta failed");
+    expect(abortReason).toBe("provider_disconnected: Provider dropped");
     expect(runMutation).toHaveBeenNthCalledWith(
       3,
       components.agent.streams.abort,
-      { streamId: "stream-1", reason: "delta failed" },
+      {
+        streamId: "stream-1",
+        reason: "provider_disconnected: Provider dropped",
+      },
     );
   });
 
@@ -428,4 +492,34 @@ describe("DeltaStreamer", () => {
     );
   });
   // TODO: test fetching partial stream data - syncStreams w/ cursors
+
+  test("does not drop a part parked in stream creation when the final step lands", async () => {
+    let resolveCreate!: (streamId: string) => void;
+    const creating = new Promise<string>((r) => (resolveCreate = r));
+    const sent: unknown[] = [];
+    const runMutation = vi
+      .fn()
+      .mockImplementationOnce(() => creating)
+      .mockImplementation((_ref: unknown, args: unknown) => {
+        sent.push(args);
+        return Promise.resolve(true);
+      });
+    const streamer = new DeltaStreamer<string>(
+      components.agent,
+      { runMutation } as unknown as MutationCtx,
+      { ...defaultTestOptions },
+      { ...testMetadata, threadId },
+    );
+
+    // A part arrives and parks in streams.create.
+    const adding = streamer.addParts(["A"]);
+    // The final step lands while creation is still in flight.
+    const finishing = streamer.flushAndStopAccepting();
+    resolveCreate("stream-1");
+    await Promise.all([adding, finishing]);
+
+    expect(sent).toHaveLength(1);
+    expect((sent[0] as { parts: string[] }).parts).toEqual(["A"]);
+  });
+
 });

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

@@ -0,0 +1,144 @@
+import { afterEach, beforeEach, expect, test, vi } from "vitest";
+import { createThread } from "../../client/threads.js";
+import { components, initConvexTest } from "./setup.test.js";
+import { DeltaStreamer } from "./streaming.js";
+
+beforeEach(() => {
+  vi.useFakeTimers();
+  vi.setSystemTime(new Date("2026-01-01T00:00:00Z"));
+});
+
+afterEach(() => {
+  vi.useRealTimers();
+});
+
+test("publishes a quiet tail after preparing an earlier batch takes time", async () => {
+  const t = initConvexTest();
+  let releasePreparation!: () => void;
+  let enteredPreparation!: () => void;
+  const preparing = new Promise<void>((resolve) => {
+    enteredPreparation = resolve;
+  });
+  const release = new Promise<void>((resolve) => {
+    releasePreparation = resolve;
+  });
+  await t.action(async (ctx) => {
+    const threadId = await createThread(ctx, components.agent, {});
+    const streamer = new DeltaStreamer<string>(
+      components.agent,
+      ctx,
+      {
+        throttleMs: 100,
+        compress: null,
+        abortSignal: undefined,
+        onAsyncAbort: async (reason) => {
+          throw new Error(reason);
+        },
+        materialize: async (parts) => {
+          if (parts.includes("first")) {
+            enteredPreparation();
+            await release;
+          }
+          return { parts, fileRefs: [] };
+        },
+      },
+      { threadId, order: 0, stepOrder: 0, format: undefined },
+    );
+    const savedParts = async () => {
+      const deltas = await ctx.runQuery(components.agent.streams.listDeltas, {
+        threadId,
+        cursors: [{ streamId: await streamer.getStreamId(), cursor: 0 }],
+      });
+      return deltas.flatMap((delta) => delta.parts);
+    };
+    try {
+      await streamer.addParts(["first"]);
+      await preparing;
+      await streamer.addParts(["tail"]);
+      await vi.advanceTimersByTimeAsync(250);
+      expect(await savedParts()).toEqual([]);
+      releasePreparation();
+      for (let elapsed = 0; elapsed < 200; elapsed++) {
+        await vi.advanceTimersByTimeAsync(1);
+        if ((await savedParts()).includes("first")) break;
+      }
+      expect(await savedParts()).toEqual(["first"]);
+      await vi.advanceTimersByTimeAsync(50);
+      expect(await savedParts()).toEqual(["first"]);
+      await vi.advanceTimersByTimeAsync(300);
+      expect.soft(await savedParts()).toEqual(["first", "tail"]);
+      const streams = await ctx.runQuery(components.agent.streams.list, {
+        threadId,
+        statuses: ["streaming"],
+      });
+      expect(streams).toHaveLength(1);
+    } finally {
+      releasePreparation();
+      await streamer.finish();
+    }
+    expect(await savedParts()).toEqual(["first", "tail"]);
+  });
+});
+
+test("a wake armed during a write respects the deadline that write set", async () => {
+  const t = initConvexTest();
+  let releasePreparation!: () => void;
+  let enteredPreparation!: () => void;
+  const preparing = new Promise<void>((resolve) => {
+    enteredPreparation = resolve;
+  });
+  const release = new Promise<void>((resolve) => {
+    releasePreparation = resolve;
+  });
+  await t.action(async (ctx) => {
+    const threadId = await createThread(ctx, components.agent, {});
+    const streamer = new DeltaStreamer<string>(
+      components.agent,
+      ctx,
+      {
+        throttleMs: 100,
+        compress: null,
+        abortSignal: undefined,
+        onAsyncAbort: async () => {},
+        materialize: async (parts) => {
+          if (parts.includes("first")) {
+            enteredPreparation();
+            await release;
+          }
+          return { parts, fileRefs: [] };
+        },
+      },
+      { threadId, order: 0, stepOrder: 0, format: undefined },
+    );
+    const streamId = await streamer.getStreamId();
+    const savedParts = async () => {
+      const deltas = await ctx.runQuery(components.agent.streams.listDeltas, {
+        threadId,
+        cursors: [{ streamId, cursor: 0 }],
+      });
+      return deltas.flatMap((delta) => delta.parts);
+    };
+
+    await streamer.addParts(["first"]);
+    await preparing;
+    // The tail is admitted while the first batch is still being prepared, so
+    // the wake it arms is measured against the previous write's deadline.
+    await vi.advanceTimersByTimeAsync(10);
+    await streamer.addParts(["tail"]);
+    releasePreparation();
+    for (let elapsed = 0; elapsed < 200; elapsed++) {
+      await vi.advanceTimersByTimeAsync(1);
+      if ((await savedParts()).includes("first")) break;
+    }
+    expect(await savedParts()).toEqual(["first"]);
+
+    // The first batch has now reset the window, so the tail is not due yet.
+    await vi.advanceTimersByTimeAsync(50);
+    expect(await savedParts()).toEqual(["first"]);
+
+    // It still has to arrive once the window it was re-armed against elapses,
+    // otherwise holding it back would pass this test by never publishing.
+    await vi.advanceTimersByTimeAsync(100);
+    expect(await savedParts()).toEqual(["first", "tail"]);
+  });
+});