@convex-dev/agent 0.7.2 → 0.7.3

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

@@ -9,9 +9,11 @@ import {
   anyApi,
 } from "convex/server";
 import { v } from "convex/values";
+import type { LanguageModelV4Source } from "@ai-sdk/provider";
 import { components, initConvexTest } from "./setup.test.js";
 import { mockModel } from "./mockModel.js";
 import { runStreamCleanup } from "./streamText.js";
+import type { StreamingOptions } from "./streaming.js";
 import { errorToString } from "./utils.js";
 
 const schema = defineSchema({});
@@ -56,6 +58,31 @@ const failingAgent = new Agent(components.agent, {
   }),
 });
 
+const sourceParts: LanguageModelV4Source[] = [
+  {
+    type: "source",
+    sourceType: "url",
+    id: "source-url-1",
+    url: "https://example.com/reference",
+    title: "Reference",
+  },
+  {
+    type: "source",
+    sourceType: "document",
+    id: "source-document-1",
+    mediaType: "application/pdf",
+    title: "Document",
+    filename: "document.pdf",
+  },
+];
+
+const sourceAgent = new Agent(components.agent, {
+  name: "source-stream-test",
+  languageModel: mockModel({
+    content: [{ type: "text", text: FINAL_TEXT }, ...sourceParts],
+  }),
+});
+
 // Action that exercises streamText with saveStreamDeltas.returnImmediately=true.
 // It consumes the stream after streamText returns, simulating the HTTP response
 // path described in issue #265.
@@ -77,7 +104,6 @@ export const streamTextReturnImmediately = action({
     // Drain the stream the way an HTTP response would. This triggers
     // onStepFinish for every step, including the final one.
     await result.consumeStream();
-    return { ok: true };
   },
 });
 
@@ -113,6 +139,83 @@ export const streamTextEmptyReturnImmediately = action({
   },
 });
 
+export const streamTextThrottled = action({
+  args: { threadId: v.string() },
+  handler: async (ctx, { threadId }) => {
+    const result = await agent.streamText(
+      ctx,
+      { threadId },
+      { prompt: "Test" },
+      {
+        saveStreamDeltas: {
+          returnImmediately: true,
+          chunking: "word",
+          throttleMs: 60_000,
+        },
+      },
+    );
+    await result.consumeStream();
+    return { ok: true };
+  },
+});
+
+// Same as streamTextThrottled, but awaited: streamText consumes the stream
+// itself, so the terminal transition happens at end-of-stream.
+export const streamTextThrottledAwaited = action({
+  args: { threadId: v.string() },
+  handler: async (ctx, { threadId }) => {
+    await agent.streamText(
+      ctx,
+      { threadId },
+      { prompt: "Test" },
+      {
+        saveStreamDeltas: {
+          chunking: "word",
+          throttleMs: 60_000,
+        },
+      },
+    );
+    return { ok: true };
+  },
+});
+
+export const streamTextNoStorage = action({
+  args: { threadId: v.string() },
+  handler: async (ctx, { threadId }) => {
+    await agent.streamText(
+      ctx,
+      { threadId },
+      { prompt: "Test" },
+      {
+        saveStreamDeltas: { chunking: "word", throttleMs: 0 },
+        storageOptions: { saveMessages: "none" },
+      },
+    );
+    return { ok: true };
+  },
+});
+
+export const streamTextNoStorageImmediate = action({
+  args: { threadId: v.string() },
+  handler: async (ctx, { threadId }) => {
+    const r = await agent.streamText(
+      ctx,
+      { threadId },
+      { prompt: "Test" },
+      {
+        saveStreamDeltas: {
+          returnImmediately: true,
+          chunking: "word",
+          throttleMs: 0,
+        },
+        storageOptions: { saveMessages: "none" },
+      },
+    );
+    await r.consumeStream();
+    return { ok: true };
+  },
+});
+
 export const streamTextCleanupFailure = action({
   args: { threadId: v.string() },
   handler: async (ctx, { threadId }) => {
@@ -150,15 +253,177 @@ export const streamTextCleanupFailure = action({
   },
 });
 
+// A generation that someone aborts out of band while it is streaming, the way
+// a client cancelling a request would: list the streaming row and abort it.
+// The throttle holds every part after the first, so the only remaining delta
+// write is the one the finishing save drains, and the component refuses it.
+export const streamTextAbortedMidStream = action({
+  args: { threadId: v.string() },
+  handler: async (ctx, { threadId }) => {
+    const result = await agent.streamText(
+      ctx,
+      { threadId },
+      { prompt: "Test" },
+      {
+        saveStreamDeltas: {
+          returnImmediately: true,
+          chunking: "word",
+          throttleMs: 60_000,
+        },
+      },
+    );
+    for (let i = 0; i < 50; i++) {
+      const streaming = await ctx.runQuery(components.agent.streams.list, {
+        threadId,
+        statuses: ["streaming"],
+      });
+      if (streaming.length) {
+        await ctx.runMutation(components.agent.streams.abort, {
+          streamId: streaming[0].streamId,
+          reason: "external abort",
+        });
+        break;
+      }
+      await new Promise((resolve) => setTimeout(resolve, 5));
+    }
+    await result.consumeStream();
+  },
+});
+
+// An empty generation on the awaited path with nothing stored: no part ever
+// reaches the streamer, so no row exists when consumption ends.
+export const streamTextEmptyNoStorageAwaited = action({
+  args: { threadId: v.string() },
+  handler: async (ctx, { threadId }) => {
+    await emptyAgent.streamText(
+      ctx,
+      { threadId },
+      { prompt: "Test" },
+      {
+        saveStreamDeltas: { chunking: "word", throttleMs: 0 },
+        storageOptions: { saveMessages: "none" },
+      },
+    );
+    return { ok: true };
+  },
+});
+
+export const streamTextWithSources = action({
+  args: { threadId: v.string(), sendSources: v.optional(v.boolean()) },
+  handler: async (ctx, { threadId, sendSources }) => {
+    const saveStreamDeltas: StreamingOptions = {
+      chunking: "word",
+      throttleMs: 0,
+    };
+    if (sendSources !== undefined) {
+      saveStreamDeltas.sendSources = sendSources;
+    }
+    await sourceAgent.streamText(
+      ctx,
+      { threadId },
+      { prompt: "Test" },
+      { saveStreamDeltas },
+    );
+    return { ok: true };
+  },
+});
+
 const testApi: ApiFromModules<{
   fns: {
     streamTextReturnImmediately: typeof streamTextReturnImmediately;
+    streamTextThrottled: typeof streamTextThrottled;
+    streamTextThrottledAwaited: typeof streamTextThrottledAwaited;
+    streamTextAbortedMidStream: typeof streamTextAbortedMidStream;
+    streamTextNoStorage: typeof streamTextNoStorage;
+    streamTextNoStorageImmediate: typeof streamTextNoStorageImmediate;
+    streamTextEmptyNoStorageAwaited: typeof streamTextEmptyNoStorageAwaited;
     streamTextEmptyAwaited: typeof streamTextEmptyAwaited;
     streamTextEmptyReturnImmediately: typeof streamTextEmptyReturnImmediately;
     streamTextCleanupFailure: typeof streamTextCleanupFailure;
+    streamTextWithSources: typeof streamTextWithSources;
   };
 }>["fns"] = anyApi["streamText.test"] as any;
 
+describe("streamText source visibility", () => {
+  test.each([
+    { name: "omitted", sendSources: undefined },
+    { name: "enabled", sendSources: true },
+  ])(
+    "keeps awaited deltas healthy with sources $name",
+    async ({ sendSources }) => {
+      const t = initConvexTest(schema);
+      const threadId = await t.run(async (ctx) =>
+        createThread(ctx, components.agent, { userId: "u1" }),
+      );
+
+      await t.action(testApi.streamTextWithSources, {
+        threadId,
+        ...(sendSources === undefined ? {} : { sendSources }),
+      });
+
+      const streams = await t.run(async (ctx) =>
+        ctx.runQuery(components.agent.streams.list, {
+          threadId,
+          statuses: ["streaming", "finished", "aborted"],
+        }),
+      );
+      expect(streams).toEqual([
+        expect.objectContaining({ status: "finished" }),
+      ]);
+      const deltas = await t.run(async (ctx) =>
+        ctx.runQuery(components.agent.streams.listDeltas, {
+          threadId,
+          cursors: streams.map((stream) => ({
+            streamId: stream.streamId,
+            cursor: 0,
+          })),
+        }),
+      );
+      const parts = deltas.flatMap((delta) => delta.parts);
+      expect(
+        parts
+          .filter((part) => part.type === "text-delta")
+          .map((part) => part.delta)
+          .join(""),
+      ).toBe(FINAL_TEXT);
+      const streamedSources = parts.filter(
+        (part) => part.type === "source-url" || part.type === "source-document",
+      );
+      expect(streamedSources).toEqual(
+        sendSources
+          ? [
+              expect.objectContaining({
+                type: "source-url",
+                sourceId: "source-url-1",
+                url: "https://example.com/reference",
+                title: "Reference",
+              }),
+              expect.objectContaining({
+                type: "source-document",
+                sourceId: "source-document-1",
+                mediaType: "application/pdf",
+                title: "Document",
+                filename: "document.pdf",
+              }),
+            ]
+          : [],
+      );
+
+      const messages = await t.run(async (ctx) =>
+        sourceAgent.listMessages(ctx, {
+          threadId,
+          paginationOpts: { cursor: null, numItems: 50 },
+        }),
+      );
+      expect(
+        messages.page.filter(
+          (message) => message.message?.role === "assistant",
+        ),
+      ).toMatchObject([{ sources: sourceParts }]);
+    },
+  );
+});
+
 describe("streamText with saveStreamDeltas.returnImmediately (issue #265)", () => {
   test("persists the final assistant text to the messages table", async () => {
     const t = initConvexTest(schema);
@@ -317,3 +582,201 @@ describe("streamText with an empty final step (issue #274)", () => {
     },
   );
 });
+
+describe("saveStreamDeltas flushes buffered parts (issue #323)", () => {
+  test("deltas hold the full text when the generation outpaces the throttle", async () => {
+    const t = initConvexTest(schema);
+    const threadId = await t.run(async (ctx) =>
+      createThread(ctx, components.agent, { userId: "u1" }),
+    );
+
+    await t.action(testApi.streamTextThrottled, { threadId });
+    await t.finishAllScheduledFunctions(() => {});
+
+    const streams = await t.run(async (ctx) =>
+      ctx.runQuery(components.agent.streams.list, {
+        threadId,
+        statuses: ["streaming", "finished", "aborted"],
+      }),
+    );
+    const deltas = await t.run(async (ctx) =>
+      ctx.runQuery(components.agent.streams.listDeltas, {
+        threadId,
+        cursors: streams.map((s) => ({ streamId: s.streamId, cursor: 0 })),
+      }),
+    );
+
+    expect(streams).toHaveLength(1);
+    expect(streams[0].status).toBe("finished");
+
+    let cursor = 0;
+    for (const delta of deltas) {
+      expect(delta.start).toBe(cursor);
+      cursor = delta.end;
+    }
+
+    const parts = deltas.flatMap((d) => d.parts);
+    const types = parts.map((p) => p.type);
+    expect(types.at(0)).toBe("start");
+    expect(types).toContain("text-start");
+    expect(types).toContain("text-end");
+    // The stream-level "finish" chunk is emitted after the last step ends, so
+    // it cannot exist yet; the row's finished status carries that instead.
+    expect(types.at(-1)).toBe("finish-step");
+    expect(types).not.toContain("finish");
+    expect(
+      parts
+        .filter((p) => p.type === "text-delta")
+        .map((p) => (p as { delta?: string }).delta ?? "")
+        .join(""),
+    ).toBe(FINAL_TEXT);
+
+    const messages = await t.run(async (ctx) =>
+      agent.listMessages(ctx, {
+        threadId,
+        paginationOpts: { cursor: null, numItems: 50 },
+      }),
+    );
+    expect(
+      messages.page
+        .filter((m) => m.message?.role === "assistant")
+        .map((m) => m.text)
+        .join(""),
+    ).toBe(FINAL_TEXT);
+  });
+
+  test("an out of band abort fails the generation instead of saving it", async () => {
+    const t = initConvexTest(schema);
+    const threadId = await t.run(async (ctx) =>
+      createThread(ctx, components.agent, { userId: "u1" }),
+    );
+
+    await t.action(testApi.streamTextAbortedMidStream, { threadId });
+    await t.finishAllScheduledFunctions(() => {});
+
+    const streams = await t.run(async (ctx) =>
+      ctx.runQuery(components.agent.streams.list, {
+        threadId,
+        statuses: ["streaming", "finished", "aborted"],
+      }),
+    );
+    expect(streams).toHaveLength(1);
+    expect(streams[0].status).toBe("aborted");
+
+    // Nobody gets to save a successful message onto a row someone aborted.
+    const messages = await t.run(async (ctx) =>
+      agent.listMessages(ctx, {
+        threadId,
+        paginationOpts: { cursor: null, numItems: 50 },
+      }),
+    );
+    expect(
+      messages.page
+        .filter((m) => m.message?.role === "assistant")
+        .map((m) => m.status),
+    ).toEqual(["failed"]);
+  });
+
+  test("the awaited path captures the stream-level finish chunk", async () => {
+    const t = initConvexTest(schema);
+    const threadId = await t.run(async (ctx) =>
+      createThread(ctx, components.agent, { userId: "u1" }),
+    );
+
+    await t.action(testApi.streamTextThrottledAwaited, { threadId });
+    await t.finishAllScheduledFunctions(() => {});
+
+    const streams = await t.run(async (ctx) =>
+      ctx.runQuery(components.agent.streams.list, {
+        threadId,
+        statuses: ["streaming", "finished", "aborted"],
+      }),
+    );
+    const deltas = await t.run(async (ctx) =>
+      ctx.runQuery(components.agent.streams.listDeltas, {
+        threadId,
+        cursors: streams.map((s) => ({ streamId: s.streamId, cursor: 0 })),
+      }),
+    );
+
+    expect(streams).toHaveLength(1);
+    expect(streams[0].status).toBe("finished");
+
+    let cursor = 0;
+    for (const delta of deltas) {
+      expect(delta.start).toBe(cursor);
+      cursor = delta.end;
+    }
+
+    const parts = deltas.flatMap((d) => d.parts);
+    const types = parts.map((p) => p.type);
+    // Unlike the returnImmediately path, nothing stops accepting parts early
+    // here: consumeStream drains at EOF, so the trailing chunks the AI SDK
+    // emits after the last onStepEnd are persisted too.
+    expect(types.at(0)).toBe("start");
+    expect(types.at(-1)).toBe("finish");
+    expect(types).toContain("finish-step");
+    expect(
+      parts
+        .filter((p) => p.type === "text-delta")
+        .map((p) => (p as { delta?: string }).delta ?? "")
+        .join(""),
+    ).toBe(FINAL_TEXT);
+  });
+});
+
+describe("stream finish ownership without message storage", () => {
+  test("the row still terminates when saveMessages is none", async () => {
+    const t = initConvexTest(schema);
+    const threadId = await t.run(async (ctx) =>
+      createThread(ctx, components.agent, { userId: "u1" }),
+    );
+
+    await t.action(testApi.streamTextNoStorage, { threadId });
+    await t.finishAllScheduledFunctions(() => {});
+
+    const streams = await t.run(async (ctx) =>
+      ctx.runQuery(components.agent.streams.list, {
+        threadId,
+        statuses: ["streaming", "finished", "aborted"],
+      }),
+    );
+    expect(streams.map((s) => s.status)).toEqual(["finished"]);
+  });
+
+  test("leaves no row behind when the generation produces nothing", async () => {
+    const t = initConvexTest(schema);
+    const threadId = await t.run(async (ctx) =>
+      createThread(ctx, components.agent, { userId: "u1" }),
+    );
+
+    await t.action(testApi.streamTextEmptyNoStorageAwaited, { threadId });
+    await t.finishAllScheduledFunctions(() => {});
+
+    const streams = await t.run(async (ctx) =>
+      ctx.runQuery(components.agent.streams.list, {
+        threadId,
+        statuses: ["streaming", "finished", "aborted"],
+      }),
+    );
+    expect(streams.filter((s) => s.status === "streaming")).toEqual([]);
+  });
+
+  test("the row still terminates on the returnImmediately path", async () => {
+    const t = initConvexTest(schema);
+    const threadId = await t.run(async (ctx) =>
+      createThread(ctx, components.agent, { userId: "u1" }),
+    );
+
+    await t.action(testApi.streamTextNoStorageImmediate, { threadId });
+    await t.finishAllScheduledFunctions(() => {});
+
+    const streams = await t.run(async (ctx) =>
+      ctx.runQuery(components.agent.streams.list, {
+        threadId,
+        statuses: ["streaming", "finished", "aborted"],
+      }),
+    );
+    expect(streams.map((s) => s.status)).toEqual(["finished"]);
+  });
+});

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,
@@ -134,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 ||
@@ -155,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,
@@ -225,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),
@@ -236,7 +245,10 @@ 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).
+          // 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",
           });
@@ -246,6 +258,11 @@ export async function streamText<
               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;
           }
         }
@@ -262,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 {
@@ -300,6 +323,11 @@ export async function streamText<
       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

@@ -492,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"]);
+  });
+
 });