@convex-dev/agent 0.6.3 → 0.7.0-alpha.0

package/src/{client → vercel/client}/streamText.ts

@@ -1,9 +1,11 @@
 import type {
+  ModelMessage,
   StepResult,
   StreamTextResult,
   ToolSet,
   UIMessage as AIUIMessage,
 } from "ai";
+import type { Context } from "@ai-sdk/provider-utils";
 import { streamText as streamTextAi } from "ai";
 import {
   compressUIMessageChunks,
@@ -14,16 +16,31 @@ import {
 import type {
   ActionCtx,
   AgentComponent,
-  AgentPrompt,
   GenerationOutputMetadata,
   Options,
-  Output,
+  StreamingTextArgs,
 } from "./types.js";
+import type { Output as AISDKOutput } from "ai";
 import { startGeneration } from "./start.js";
-import type { Agent } from "./index.js";
-import { getModelName, getProviderName } from "../shared.js";
+import type { Agent } from "../index.js";
+import { getModelName, getProviderName } from "../../shared.js";
 import { errorToString, willContinue } from "./utils.js";
 
+/** Finish every abort cleanup path before surfacing an internal failure. */
+export async function runAbortCleanup(cleanup: {
+  failCall: () => Promise<void>;
+  failStreamer: () => Promise<void>;
+  onAbort?: () => PromiseLike<void> | void;
+}): Promise<void> {
+  const results = await Promise.allSettled([
+    cleanup.failCall(),
+    cleanup.failStreamer(),
+  ]);
+  await cleanup.onAbort?.();
+  const failure = results.find((result) => result.status === "rejected");
+  if (failure) throw failure.reason;
+}
+
 /**
  * This behaves like {@link streamText} from the "ai" package except that
  * it add context based on the userId and threadId and saves the input and
@@ -32,8 +49,14 @@ import { errorToString, willContinue } from "./utils.js";
  * to a thread (and optionally userId).
  */
 export async function streamText<
-  TOOLS extends ToolSet,
-  OUTPUT extends Output<any, any, any> = never,
+  AgentTools extends ToolSet,
+  TOOLS extends ToolSet | undefined = undefined,
+  OUTPUT extends AISDKOutput.Output<any, any, any> = AISDKOutput.Output<
+    string,
+    string,
+    never
+  >,
+  RUNTIME_CONTEXT extends Context = Context,
 >(
   ctx: ActionCtx,
   component: AgentComponent,
@@ -41,17 +64,7 @@ export async function streamText<
    * The arguments to the streamText function, similar to the ai sdk's
    * {@link streamText} function, along with Agent prompt options.
    */
-  streamTextArgs: AgentPrompt &
-    Omit<
-      Parameters<typeof streamTextAi<TOOLS, OUTPUT>>[0],
-      "model" | "prompt" | "messages"
-    > & {
-      /**
-       * The tools to use for the tool calls. This will override tools specified
-       * in the Agent constructor or createThread / continueThread.
-       */
-      tools?: TOOLS;
-    },
+  streamTextArgs: StreamingTextArgs<AgentTools, TOOLS, OUTPUT, RUNTIME_CONTEXT>,
   /**
    * The {@link ContextOptions} and {@link StorageOptions}
    * options to use for fetching contextual messages and saving input/output messages.
@@ -73,17 +86,49 @@ export async function streamText<
     saveStreamDeltas?: boolean | StreamingOptions;
     agentForToolCtx?: Agent;
   },
-): Promise<StreamTextResult<TOOLS, OUTPUT> & GenerationOutputMetadata> {
+): Promise<
+  StreamTextResult<
+    TOOLS extends undefined ? AgentTools : TOOLS,
+    RUNTIME_CONTEXT,
+    OUTPUT
+  > &
+    GenerationOutputMetadata
+> {
+  type Tools = TOOLS extends undefined ? AgentTools : TOOLS;
   const { threadId } = options ?? {};
   const { args, userId, order, stepOrder, promptMessageId, ...call } =
-    await startGeneration(ctx, component, streamTextArgs, options);
+    await startGeneration<
+      StreamingTextArgs<AgentTools, TOOLS, OUTPUT, RUNTIME_CONTEXT>,
+      Tools,
+      object,
+      RUNTIME_CONTEXT
+    >(
+      ctx,
+      component,
+      streamTextArgs,
+      options,
+      "streamText",
+    );
 
-  const steps: StepResult<TOOLS>[] = [];
+  const steps: StepResult<Tools, RUNTIME_CONTEXT>[] = [];
+  let initialResponseMessages: ModelMessage[] = [];
+  let initialResponseMessagesSaved = false;
+  const responseMessagesForStep = (
+    step: StepResult<Tools, RUNTIME_CONTEXT>,
+  ) => [
+    ...(initialResponseMessagesSaved ? [] : initialResponseMessages),
+    ...step.response.messages,
+  ];
 
   // Track the final step for atomic save with stream finish (issue #181).
   // Only used when streamText awaits stream consumption itself; the
   // `returnImmediately` path saves inline instead (see onStepFinish below).
-  let pendingFinalStep: StepResult<TOOLS> | undefined;
+  let pendingFinalStep:
+    | {
+        step: StepResult<Tools, RUNTIME_CONTEXT>;
+        responseMessages: ModelMessage[];
+      }
+    | undefined;
 
   // Whether streamText will await stream consumption before returning.
   // When false (saveStreamDeltas.returnImmediately === true), we cannot
@@ -123,7 +168,7 @@ export async function streamText<
         )
       : undefined;
 
-  const result = streamTextAi({
+  const result = streamTextAi<Tools, RUNTIME_CONTEXT, OUTPUT>({
     ...args,
     abortSignal: streamer?.abortController.signal ?? args.abortSignal,
     experimental_transform: mergeTransforms(
@@ -136,7 +181,19 @@ export async function streamText<
       await streamer?.fail(errorToString(error.error));
       return streamTextArgs.onError?.(error);
     },
+    onAbort: async (event) => {
+      const reason = args.abortSignal?.reason
+        ? errorToString(args.abortSignal.reason)
+        : "streamText aborted";
+      await runAbortCleanup({
+        failCall: () => call.fail(reason),
+        failStreamer: async () => streamer?.fail(reason),
+        onAbort: () => streamTextArgs.onAbort?.(event),
+      });
+    },
     prepareStep: async (options) => {
+      if (options.stepNumber === 0)
+        initialResponseMessages = [...options.responseMessages];
       const result = await streamTextArgs.prepareStep?.(options);
       if (result) {
         const model = result.model ?? options.model;
@@ -150,7 +207,7 @@ export async function streamText<
       }
       return undefined;
     },
-    onStepFinish: async (step) => {
+    onStepEnd: async (step) => {
       steps.push(step);
       const createPendingMessage = await willContinue(steps, args.stopWhen);
       if (!createPendingMessage && streamer) {
@@ -159,22 +216,36 @@ export async function streamText<
         if (willAwaitStream) {
           // We're about to `await stream` below — defer the save so it
           // happens atomically with stream finish (issue #181).
-          pendingFinalStep = step;
+          pendingFinalStep = {
+            step,
+            responseMessages: responseMessagesForStep(step),
+          };
         } 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 }, false, finishStreamId);
+          await call.save(
+            { step, responseMessages: responseMessagesForStep(step) },
+            false,
+            finishStreamId,
+          );
+          initialResponseMessagesSaved = true;
         }
       } else {
-        await call.save({ step }, createPendingMessage);
+        await call.save(
+          { step, responseMessages: responseMessagesForStep(step) },
+          createPendingMessage,
+        );
+        initialResponseMessagesSaved = true;
       }
-      return args.onStepFinish?.(step);
+      return (streamTextArgs.onStepEnd ?? streamTextArgs.onStepFinish)?.(step);
     },
-  }) as StreamTextResult<TOOLS, OUTPUT>;
+  } as Parameters<
+    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>>(),
   );
   if (willAwaitStream) {
     try {
@@ -188,7 +259,7 @@ export async function streamText<
       // Save the deferred final step if it was already generated but not yet persisted
       if (pendingFinalStep) {
         try {
-          await call.save({ step: pendingFinalStep }, false);
+          await call.save(pendingFinalStep, false);
         } catch (saveError) {
           console.error("Failed to save deferred final step:", saveError);
         }
@@ -201,7 +272,7 @@ export async function streamText<
   // 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({ step: pendingFinalStep }, false, finishStreamId);
+    await call.save(pendingFinalStep, false, finishStreamId);
   }
   const metadata: GenerationOutputMetadata = {
     promptMessageId,

package/src/{client → vercel/client}/streaming.integration.test.ts

@@ -1,5 +1,5 @@
 import { beforeEach, describe, expect, test } from "vitest";
-import { createThread } from "./index.js";
+import { createThread } from "../../client/threads.js";
 import type { GenericSchema, SchemaDefinition } from "convex/server";
 import { streamText } from "ai";
 import { components, initConvexTest } from "./setup.test.js";
@@ -12,10 +12,9 @@ import {
 import {
   getParts,
   deriveUIMessagesFromDeltas,
-  deriveUIMessagesFromTextStreamParts,
 } from "../deltas.js";
 import type { TestConvex } from "convex-test";
-import type { StreamDelta, StreamMessage } from "../validators.js";
+import type { StreamDelta, StreamMessage } from "../../validators.js";
 import { dedupeMessages } from "../react/useUIMessages.js";
 
 const defaultTestOptions = {
@@ -152,9 +151,7 @@ describe("HTTP Streaming Initiation", () => {
 
       // Verify we can reconstruct the text from deltas
       const { parts } = getParts(deltas);
-      const textParts = parts.filter(
-        (p: any) => p.type === "text-delta",
-      );
+      const textParts = parts.filter((p: any) => p.type === "text-delta");
       expect(textParts.length).toBeGreaterThan(0);
     });
   });
@@ -265,10 +262,9 @@ describe("Stream Exclusion Logic", () => {
       await streamer2.addParts([{ type: "start" }]);
 
       // Default list: only streaming
-      const defaultStreams = await ctx.runQuery(
-        components.agent.streams.list,
-        { threadId },
-      );
+      const defaultStreams = await ctx.runQuery(components.agent.streams.list, {
+        threadId,
+      });
       expect(defaultStreams).toHaveLength(1);
       expect(defaultStreams[0].status).toBe("streaming");
       expect(defaultStreams[0].order).toBe(1);
@@ -325,10 +321,10 @@ describe("Stream Exclusion Logic", () => {
       expect(finishedStreams[0].status).toBe("finished");
 
       // Query for only aborted
-      const abortedStreams = await ctx.runQuery(
-        components.agent.streams.list,
-        { threadId, statuses: ["aborted"] },
-      );
+      const abortedStreams = await ctx.runQuery(components.agent.streams.list, {
+        threadId,
+        statuses: ["aborted"],
+      });
       expect(abortedStreams).toHaveLength(1);
       expect(abortedStreams[0].status).toBe("aborted");
 
@@ -487,10 +483,10 @@ describe("Delta Stream Consumption", () => {
       expect(laterParts.length).toBeLessThanOrEqual(allParts.length);
 
       // Fetching from the end cursor should yield nothing
-      const noDeltas = await ctx.runQuery(
-        components.agent.streams.listDeltas,
-        { threadId, cursors: [{ cursor: endCursor, streamId }] },
-      );
+      const noDeltas = await ctx.runQuery(components.agent.streams.listDeltas, {
+        threadId,
+        cursors: [{ cursor: endCursor, streamId }],
+      });
       expect(noDeltas).toHaveLength(0);
     });
   });
@@ -526,16 +522,13 @@ describe("Delta Stream Consumption", () => {
       const id2 = streamer2.streamId!;
 
       // Fetch deltas for both streams simultaneously
-      const deltas = await ctx.runQuery(
-        components.agent.streams.listDeltas,
-        {
-          threadId,
-          cursors: [
-            { cursor: 0, streamId: id1 },
-            { cursor: 0, streamId: id2 },
-          ],
-        },
-      );
+      const deltas = await ctx.runQuery(components.agent.streams.listDeltas, {
+        threadId,
+        cursors: [
+          { cursor: 0, streamId: id1 },
+          { cursor: 0, streamId: id2 },
+        ],
+      });
 
       // Should have deltas for both streams
       const s1Deltas = deltas.filter((d) => d.streamId === id1);
@@ -568,10 +561,10 @@ describe("Delta Stream Consumption", () => {
         threadId,
         statuses: ["finished"],
       });
-      const deltas = await ctx.runQuery(
-        components.agent.streams.listDeltas,
-        { threadId, cursors: [{ cursor: 0, streamId }] },
-      );
+      const deltas = await ctx.runQuery(components.agent.streams.listDeltas, {
+        threadId,
+        cursors: [{ cursor: 0, streamId }],
+      });
 
       // Derive UI messages
       const uiMessages = await deriveUIMessagesFromDeltas(
@@ -615,10 +608,10 @@ describe("Delta Stream Consumption", () => {
       await streamer.consumeStream(result.toUIMessageStream());
       const streamId = streamer.streamId!;
 
-      const deltas = await ctx.runQuery(
-        components.agent.streams.listDeltas,
-        { threadId, cursors: [{ cursor: 0, streamId }] },
-      );
+      const deltas = await ctx.runQuery(components.agent.streams.listDeltas, {
+        threadId,
+        cursors: [{ cursor: 0, streamId }],
+      });
       const { parts } = getParts(deltas);
 
       // Compressed: all text-deltas for one text section should be merged
@@ -677,77 +670,6 @@ describe("Delta Stream Consumption", () => {
     expect((parts[0] as { type: string }).type).toBe("new");
     expect(cursor).toBe(6);
   });
-
-  test("TextStreamPart format delta reconstruction with tool calls", () => {
-    const streamId = "s1";
-    const streamMessage: StreamMessage = {
-      streamId,
-      order: 1,
-      stepOrder: 0,
-      status: "streaming",
-    };
-    const deltas: StreamDelta[] = [
-      {
-        streamId,
-        start: 0,
-        end: 1,
-        parts: [{ type: "text-delta", id: "txt-0", text: "Let me call a tool. " }],
-      },
-      {
-        streamId,
-        start: 1,
-        end: 2,
-        parts: [
-          {
-            type: "tool-call",
-            toolCallId: "tc1",
-            toolName: "search",
-            input: { query: "hello" },
-          },
-        ],
-      },
-      {
-        streamId,
-        start: 2,
-        end: 3,
-        parts: [
-          {
-            type: "tool-result",
-            toolCallId: "tc1",
-            toolName: "search",
-            output: "Found 3 results",
-          },
-        ],
-      },
-      {
-        streamId,
-        start: 3,
-        end: 4,
-        parts: [
-          { type: "text-delta", id: "txt-1", text: "Here are the results." },
-        ],
-      },
-    ];
-
-    const [messages, , changed] = deriveUIMessagesFromTextStreamParts(
-      "thread1",
-      [streamMessage],
-      [],
-      deltas,
-    );
-
-    expect(messages).toHaveLength(1);
-    expect(changed).toBe(true);
-
-    const msg = messages[0];
-    expect(msg.text).toContain("Let me call a tool.");
-    expect(msg.text).toContain("Here are the results.");
-
-    const toolParts = msg.parts.filter((p: any) =>
-      p.type.startsWith("tool-"),
-    );
-    expect(toolParts.length).toBeGreaterThan(0);
-  });
 });
 
 // ============================================================================
@@ -879,18 +801,21 @@ describe("Fallback Behavior", () => {
       order: 0,
       stepOrder: 0,
       status: "streaming",
+      format: "UIMessageChunk",
     };
     const finishedMsg: StreamMessage = {
       streamId: "s2",
       order: 1,
       stepOrder: 0,
       status: "finished",
+      format: "UIMessageChunk",
     };
     const abortedMsg: StreamMessage = {
       streamId: "s3",
       order: 2,
       stepOrder: 0,
       status: "aborted",
+      format: "UIMessageChunk",
     };
 
     const msgs = await deriveUIMessagesFromDeltas(
@@ -1001,10 +926,10 @@ describe("Stream Lifecycle Integration", () => {
       expect(finished).toHaveLength(1);
 
       // 4. Derive UI messages from stored deltas
-      const deltas = await ctx.runQuery(
-        components.agent.streams.listDeltas,
-        { threadId, cursors: [{ cursor: 0, streamId }] },
-      );
+      const deltas = await ctx.runQuery(components.agent.streams.listDeltas, {
+        threadId,
+        cursors: [{ cursor: 0, streamId }],
+      });
       const uiMessages = await deriveUIMessagesFromDeltas(
         threadId,
         finished,
@@ -1056,10 +981,10 @@ describe("Stream Lifecycle Integration", () => {
       expect(aborted[0].status).toBe("aborted");
 
       // Even aborted streams have their deltas stored
-      const deltas = await ctx.runQuery(
-        components.agent.streams.listDeltas,
-        { threadId, cursors: [{ cursor: 0, streamId }] },
-      );
+      const deltas = await ctx.runQuery(components.agent.streams.listDeltas, {
+        threadId,
+        cursors: [{ cursor: 0, streamId }],
+      });
       expect(deltas.length).toBeGreaterThan(0);
     });
   });

package/src/{client → vercel/client}/streaming.test.ts

@@ -1,10 +1,14 @@
-import { beforeEach, describe, expect, test } from "vitest";
-import { createThread } from "./index.js";
+import { beforeEach, describe, expect, test, vi } from "vitest";
+import { createThread } from "../../client/threads.js";
+import type { MutationCtx } from "../../client/types.js";
 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";
 
@@ -181,6 +185,176 @@ describe("DeltaStreamer", () => {
       }
     });
   });
-  // TODO: test errors & aborted states
+  test("honors a signal that was aborted before construction", async () => {
+    const abortController = new AbortController();
+    abortController.abort();
+
+    await t.run(async (ctx) => {
+      const streamer = new DeltaStreamer<string>(
+        components.agent,
+        ctx,
+        { ...defaultTestOptions, abortSignal: abortController.signal },
+        { ...testMetadata, threadId },
+      );
+
+      expect(streamer.abortController.signal.aborted).toBe(true);
+      await streamer.addParts(["ignored"]);
+      expect(streamer.streamId).toBeUndefined();
+      await expect(streamer.getOrCreateStreamId()).rejects.toThrow(
+        "Cannot create a stream after it has been aborted",
+      );
+    });
+  });
+
+  test("shares signal and fail cleanup while stream creation is in flight", async () => {
+    let resolveCreate!: (streamId: string) => void;
+    const creatingStream = new Promise<string>((resolve) => {
+      resolveCreate = resolve;
+    });
+    let resolveAbort!: () => void;
+    const abortingStream = new Promise<void>((resolve) => {
+      resolveAbort = resolve;
+    });
+    const runMutation = vi
+      .fn()
+      .mockImplementationOnce(() => creatingStream)
+      .mockImplementationOnce(() => abortingStream);
+    const abortController = new AbortController();
+    const streamer = new DeltaStreamer<string>(
+      components.agent,
+      { runMutation } as unknown as MutationCtx,
+      { ...defaultTestOptions, abortSignal: abortController.signal },
+      { ...testMetadata, threadId },
+    );
+
+    const streamId = streamer.getStreamId();
+    abortController.abort();
+    const failing = streamer.fail("creation failed");
+    let failSettled = false;
+    void failing.then(() => {
+      failSettled = true;
+    });
+    resolveCreate("stream-1");
+    await expect(streamId).resolves.toBe("stream-1");
+    await vi.waitFor(() => expect(runMutation).toHaveBeenCalledTimes(2));
+
+    expect(failSettled).toBe(false);
+    resolveAbort();
+    await failing;
+
+    expect(runMutation).toHaveBeenNthCalledWith(
+      2,
+      components.agent.streams.abort,
+      { streamId: "stream-1", reason: "abortSignal" },
+    );
+  });
+
+  test("aborts the component stream when a delta write fails", async () => {
+    const runMutation = vi
+      .fn()
+      .mockResolvedValueOnce("stream-1")
+      .mockRejectedValueOnce(new Error("delta failed"))
+      .mockResolvedValueOnce(undefined);
+    let abortReason: string | undefined;
+    const streamer = new DeltaStreamer<string>(
+      components.agent,
+      { runMutation } as unknown as MutationCtx,
+      {
+        ...defaultTestOptions,
+        onAsyncAbort: async (reason) => {
+          abortReason = reason;
+        },
+      },
+      { ...testMetadata, threadId },
+    );
+
+    await streamer.addParts(["A"]);
+    await streamer.finish();
+
+    expect(abortReason).toBe("delta failed");
+    expect(runMutation).toHaveBeenNthCalledWith(
+      3,
+      components.agent.streams.abort,
+      { streamId: "stream-1", reason: "delta failed" },
+    );
+  });
+
+  test("surfaces pending-message cleanup failure after aborting the stream", async () => {
+    const pendingMessageFailure = new Error("pending message cleanup failed");
+    let resolveComponentAbort!: () => void;
+    const componentAborted = new Promise<void>((resolve) => {
+      resolveComponentAbort = resolve;
+    });
+    const runMutation = vi
+      .fn()
+      .mockResolvedValueOnce("stream-1")
+      .mockRejectedValueOnce(new Error("delta failed"))
+      .mockImplementationOnce(() => {
+        resolveComponentAbort();
+        return Promise.resolve();
+      });
+    const streamer = new DeltaStreamer<string>(
+      components.agent,
+      { runMutation } as unknown as MutationCtx,
+      {
+        ...defaultTestOptions,
+        onAsyncAbort: async () => {
+          throw pendingMessageFailure;
+        },
+      },
+      { ...testMetadata, threadId },
+    );
+    const stream = {
+      async *[Symbol.asyncIterator]() {
+        yield "A";
+        await componentAborted;
+      },
+    } as unknown as Parameters<typeof streamer.consumeStream>[0];
+
+    await expect(streamer.consumeStream(stream)).rejects.toBe(
+      pendingMessageFailure,
+    );
+
+    expect(runMutation).toHaveBeenNthCalledWith(
+      3,
+      components.agent.streams.abort,
+      { streamId: "stream-1", reason: "delta failed" },
+    );
+  });
+
+  test("finishes external abort cleanup when the active delta write fails", async () => {
+    const pendingMessageFailure = new Error("pending message cleanup failed");
+    let rejectDelta!: (error: Error) => void;
+    const deltaWrite = new Promise<never>((_, reject) => {
+      rejectDelta = reject;
+    });
+    const runMutation = vi
+      .fn()
+      .mockResolvedValueOnce("stream-1")
+      .mockImplementationOnce(() => deltaWrite)
+      .mockResolvedValueOnce(undefined);
+    const streamer = new DeltaStreamer<string>(
+      components.agent,
+      { runMutation } as unknown as MutationCtx,
+      {
+        ...defaultTestOptions,
+        onAsyncAbort: async () => {
+          throw pendingMessageFailure;
+        },
+      },
+      { ...testMetadata, threadId },
+    );
+
+    await streamer.addParts(["A"]);
+    const failing = streamer.fail("external abort");
+    rejectDelta(new Error("delta failed"));
+
+    await expect(failing).rejects.toBe(pendingMessageFailure);
+    expect(runMutation).toHaveBeenNthCalledWith(
+      3,
+      components.agent.streams.abort,
+      { streamId: "stream-1", reason: "external abort" },
+    );
+  });
   // TODO: test fetching partial stream data - syncStreams w/ cursors
 });