@convex-dev/agent 0.6.4 → 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";
@@ -9,9 +9,12 @@ import {
   DeltaStreamer,
   mergeTransforms,
 } from "./streaming.js";
-import { getParts, deriveUIMessagesFromDeltas } from "../deltas.js";
+import {
+  getParts,
+  deriveUIMessagesFromDeltas,
+} 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 = {

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
 });

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

@@ -17,7 +17,7 @@ import {
   type StreamArgs,
   type StreamDelta,
   type StreamMessage,
-} from "../validators.js";
+} from "../../validators.js";
 import type {
   ActionCtx,
   AgentComponent,
@@ -208,6 +208,7 @@ export class DeltaStreamer<T> {
   #nextParts: T[] = [];
   #latestWrite: number = 0;
   #ongoingWrite: Promise<void> | undefined;
+  #abortPromise: Promise<void> | undefined;
   #cursor: number = 0;
   public abortController: AbortController;
   // When true, the stream will be finished externally (e.g., atomically via addMessages)
@@ -243,34 +244,28 @@ export class DeltaStreamer<T> {
     this.#nextParts = [];
     this.abortController = new AbortController();
     if (config.abortSignal) {
-      config.abortSignal.addEventListener("abort", async () => {
-        try {
-          if (this.abortController.signal.aborted) {
-            return;
-          }
-          this.abortController.abort();
-          // Wait for in-flight stream creation before trying to abort it
-          if (this.#creatingStreamIdPromise) {
-            await this.#creatingStreamIdPromise;
-          }
-          if (this.streamId) {
-            await this.#ongoingWrite;
-            await this.ctx.runMutation(this.component.streams.abort, {
-              streamId: this.streamId,
-              reason: "abortSignal",
-            });
-          }
-        } catch {
-          // Best-effort cleanup — the stream will be garbage-collected
-          // by the 10-minute timeout if this fails.
-        }
-      });
+      const abortFromSignal = () => {
+        void this.#abort("abortSignal").catch(() => {
+          // Best-effort cleanup — the stream timeout is the fallback.
+        });
+      };
+      if (config.abortSignal.aborted) {
+        abortFromSignal();
+      } else {
+        config.abortSignal.addEventListener("abort", abortFromSignal, {
+          once: true,
+        });
+      }
     }
   }
 
   // Avoid race conditions by only creating once
   #creatingStreamIdPromise: Promise<string> | undefined;
   public async getStreamId() {
+    if (this.abortController.signal.aborted) {
+      await this.#abortPromise;
+      throw new Error("Cannot create a stream after it has been aborted");
+    }
     if (!this.streamId) {
       if (!this.#creatingStreamIdPromise) {
         this.#creatingStreamIdPromise = this.ctx.runMutation(
@@ -310,9 +305,11 @@ export class DeltaStreamer<T> {
     }
     // Skip finish if it will be handled externally (atomically with message save)
     // or if the stream was aborted (e.g., due to a failed delta write).
-    // Aborted streams are cleaned up via streams.abort (called by the abort
-    // signal handler), so we don't need to call finish() for them.
-    if (!this.#finishedExternally && !this.abortController.signal.aborted) {
+    // Abort cleanup owns the terminal component transition, so consumeStream
+    // must wait for it instead of also trying to finish the stream.
+    if (this.abortController.signal.aborted) {
+      await this.#waitForAbortCleanup();
+    } else if (!this.#finishedExternally) {
       await this.finish();
     }
   }
@@ -343,28 +340,27 @@ export class DeltaStreamer<T> {
       return;
     }
     this.#latestWrite = Date.now();
+    let success: boolean;
     try {
-      const success = await this.ctx.runMutation(
+      success = await this.ctx.runMutation(
         this.component.streams.addDelta,
         delta,
       );
-      if (!success) {
-        // An in-flight #sendDelta started before markFinishedExternally()
-        // will get `success === false` because the stream row is already
-        // "finished". That's a benign late-write miss, not a failure —
-        // don't convert it into an abort.
-        if (this.#finishedExternally) {
-          return;
-        }
-        await this.config.onAsyncAbort("async abort");
-        this.abortController.abort();
-        return;
-      }
     } catch (e) {
-      await this.config.onAsyncAbort(
+      await this.#abortDelta(
         e instanceof Error ? e.message : "unknown error",
       );
-      this.abortController.abort();
+      return;
+    }
+    if (!success) {
+      // An in-flight #sendDelta started before markFinishedExternally()
+      // will get `success === false` because the stream row is already
+      // "finished". That's a benign late-write miss, not a failure —
+      // don't convert it into an abort.
+      if (this.#finishedExternally) {
+        return;
+      }
+      await this.#abortDelta("async abort");
       return;
     }
     // Now that we've sent the delta, check if we need to send another one.
@@ -400,9 +396,17 @@ export class DeltaStreamer<T> {
     if (!this.streamId) {
       return;
     }
-    await this.#ongoingWrite;
+    try {
+      await this.#ongoingWrite;
+    } catch (error) {
+      if (this.abortController.signal.aborted) {
+        await this.#waitForAbortCleanup();
+      }
+      throw error;
+    }
     await this.#sendDelta(); // #sendDelta checks aborted internally
     if (this.abortController.signal.aborted) {
+      await this.#waitForAbortCleanup();
       return;
     }
     await this.ctx.runMutation(this.component.streams.finish, {
@@ -411,18 +415,76 @@ export class DeltaStreamer<T> {
   }
 
   public async fail(reason: string) {
-    if (this.abortController.signal.aborted) {
-      return;
+    await this.#abort(reason);
+  }
+
+  async #abortDelta(reason: string) {
+    let callbackFailure: { error: unknown } | undefined;
+    try {
+      await this.config.onAsyncAbort(reason);
+    } catch (error) {
+      callbackFailure = { error };
     }
-    this.abortController.abort();
-    if (!this.streamId) {
-      return;
+    if (!this.#abortPromise) {
+      try {
+        await this.#abort(reason, false);
+      } catch {
+        // The stream timeout is the bounded cleanup fallback.
+      }
     }
-    await this.#ongoingWrite;
-    await this.ctx.runMutation(this.component.streams.abort, {
-      streamId: this.streamId,
-      reason,
-    });
+    if (callbackFailure) throw callbackFailure.error;
+  }
+
+  #abort(reason: string, waitForOngoingWrite = true): Promise<void> {
+    if (!this.#abortPromise) {
+      this.abortController.abort();
+      this.#abortPromise = this.#abortCreatedStream(
+        reason,
+        waitForOngoingWrite,
+      );
+    }
+    return this.#abortPromise;
+  }
+
+  async #waitForAbortCleanup() {
+    let writeFailure: { error: unknown } | undefined;
+    try {
+      await this.#ongoingWrite;
+    } catch (error) {
+      writeFailure = { error };
+    }
+    try {
+      await this.#abortPromise;
+    } catch (error) {
+      writeFailure ??= { error };
+    }
+    if (writeFailure) throw writeFailure.error;
+  }
+
+  async #abortCreatedStream(reason: string, waitForOngoingWrite: boolean) {
+    if (this.#creatingStreamIdPromise) {
+      this.streamId ??= await this.#creatingStreamIdPromise;
+    }
+    if (!this.streamId) return;
+    let writeFailure: { error: unknown } | undefined;
+    if (waitForOngoingWrite) {
+      try {
+        await this.#ongoingWrite;
+      } catch (error) {
+        writeFailure = { error };
+      }
+    }
+    let abortFailure: { error: unknown } | undefined;
+    try {
+      await this.ctx.runMutation(this.component.streams.abort, {
+        streamId: this.streamId,
+        reason,
+      });
+    } catch (error) {
+      abortFailure = { error };
+    }
+    if (writeFailure) throw writeFailure.error;
+    if (abortFailure) throw abortFailure.error;
   }
 }