@convex-dev/agent 0.7.1 → 0.7.3

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

@@ -9,15 +9,31 @@ 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 { runAbortCleanup } from "./streamText.js";
+import { runStreamCleanup } from "./streamText.js";
+import type { StreamingOptions } from "./streaming.js";
+import { errorToString } from "./utils.js";
 
 const schema = defineSchema({});
 type DataModel = DataModelFromSchemaDefinition<typeof schema>;
 const action = actionGeneric as ActionBuilder<DataModel, "public">;
 
 const FINAL_TEXT = "Hello from the model";
+const PROVIDER_FAILURE_TEXT = "Mock provider failure";
+const CLEANUP_FAILURE_TEXT = "finalizeMessage rejected";
+
+function hasKeys(
+  value: unknown,
+  keys: string[],
+): value is Record<string, unknown> {
+  return (
+    value !== null &&
+    typeof value === "object" &&
+    keys.every((key) => key in value)
+  );
+}
 
 const agent = new Agent(components.agent, {
   name: "stream-test",
@@ -34,6 +50,39 @@ const emptyAgent = new Agent(components.agent, {
   }),
 });
 
+const failingAgent = new Agent(components.agent, {
+  name: "failing-stream-test",
+  languageModel: mockModel({
+    content: [{ type: "text", text: "partial response" }],
+    fail: { error: PROVIDER_FAILURE_TEXT },
+  }),
+});
+
+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.
@@ -55,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 };
   },
 });
 
@@ -91,14 +139,291 @@ 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 }) => {
+    const providerErrors: string[] = [];
+    let aborts = 0;
+    const failingCtx = {
+      ...ctx,
+      runMutation: (async (reference, args) => {
+        if (hasKeys(args, ["messageId", "result"])) {
+          throw new Error(CLEANUP_FAILURE_TEXT);
+        }
+        return ctx.runMutation(reference, args);
+      }) as typeof ctx.runMutation,
+    };
+    let caught: string | undefined;
+    try {
+      await failingAgent.streamText(
+        failingCtx,
+        { threadId },
+        {
+          prompt: "Test",
+          onError: ({ error }) => {
+            providerErrors.push(errorToString(error));
+          },
+          onAbort: () => {
+            aborts += 1;
+          },
+        },
+        { saveStreamDeltas: { chunking: "word", throttleMs: 0 } },
+      );
+    } catch (error) {
+      caught = errorToString(error);
+    }
+    return { providerErrors, aborts, caught };
+  },
+});
+
+// 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);
@@ -147,28 +472,64 @@ describe("streamText with saveStreamDeltas.returnImmediately (issue #265)", () =
 });
 
 describe("streamText abort cleanup", () => {
-  test("attempts every cleanup and rethrows the first internal failure", async () => {
+  test("finishes durable cleanup before invoking onAbort", async () => {
     const calls: string[] = [];
-    const firstFailure = new Error("failed pending message cleanup");
+    let resolveStreamer!: () => void;
+    const syncFailure = new Error("synchronous pending message cleanup");
 
-    await expect(
-      runAbortCleanup({
-        failCall: async () => {
-          calls.push("call.fail");
-          throw firstFailure;
-        },
-        failStreamer: async () => {
+    const cleanup = runStreamCleanup({
+      failCall: () => {
+        calls.push("call.fail");
+        throw syncFailure;
+      },
+      failStreamer: () =>
+        new Promise<void>((resolve) => {
           calls.push("streamer.fail");
-          throw new Error("failed stream cleanup");
-        },
-        onAbort: () => {
-          calls.push("user.onAbort");
-        },
-      }),
-    ).rejects.toBe(firstFailure);
+          resolveStreamer = resolve;
+        }),
+      onAbort: () => {
+        calls.push("user.onAbort");
+      },
+    });
+
+    await Promise.resolve();
+    expect(calls).toEqual(["call.fail", "streamer.fail"]);
+    resolveStreamer();
+    await expect(cleanup).rejects.toBe(syncFailure);
 
     expect(calls).toEqual(["call.fail", "streamer.fail", "user.onAbort"]);
   });
+
+  test("surfaces a cleanup failure without hiding the provider error", async () => {
+    const t = initConvexTest(schema);
+    const threadId = await t.run(async (ctx) =>
+      createThread(ctx, components.agent, { userId: "u1" }),
+    );
+
+    const { providerErrors, aborts, caught } = await t.action(
+      testApi.streamTextCleanupFailure,
+      { threadId },
+    );
+
+    expect(providerErrors).toEqual([PROVIDER_FAILURE_TEXT]);
+    expect(aborts).toBe(0);
+    expect(caught).toBe(CLEANUP_FAILURE_TEXT);
+
+    const streaming = await t.run(async (ctx) =>
+      ctx.runQuery(components.agent.streams.list, {
+        threadId,
+        statuses: ["streaming"],
+      }),
+    );
+    const aborted = await t.run(async (ctx) =>
+      ctx.runQuery(components.agent.streams.list, {
+        threadId,
+        statuses: ["aborted"],
+      }),
+    );
+    expect(streaming).toHaveLength(0);
+    expect(aborted).toHaveLength(1);
+  });
 });
 
 describe("streamText with an empty final step (issue #274)", () => {
@@ -221,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"]);
+  });
+});