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

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

@@ -1,36 +1,87 @@
 import {
-  stepCountIs,
-  type CallSettings,
+  isStepCount,
   type GenerateObjectResult,
   type IdGenerator,
   type LanguageModel,
+  type Instructions,
   type ModelMessage,
   type StepResult,
   type StopCondition,
   type ToolSet,
 } from "ai";
+import type { Context } from "@ai-sdk/provider-utils";
 import {
   serializeResponseMessages,
   serializeObjectResult,
 } from "../mapping.js";
 import { embedMessages, fetchContextWithPrompt } from "./search.js";
-import type { ActionCtx, AgentComponent, Config, Options } from "./types.js";
-import type { Message, MessageDoc } from "../validators.js";
+import type {
+  ActionCtx,
+  AgentCallSettings,
+  AgentComponent,
+  Config,
+  Options,
+} from "./types.js";
+import type { Message, MessageDoc } from "../../validators.js";
 import {
   getModelName,
   getProviderName,
   type ModelOrMetadata,
-} from "../shared.js";
+} from "../../shared.js";
 import { wrapTools, type ToolCtx } from "./createTool.js";
-import type { Agent } from "./index.js";
+import type { Agent } from "../index.js";
 import { assert, omit } from "convex-helpers";
 import { saveInputMessages } from "./saveInputMessages.js";
 import type { GenericActionCtx, GenericDataModel } from "convex/server";
 
+export function resolveUsageModel(
+  toSave:
+    | { step: { model?: ModelOrMetadata } }
+    | { object: unknown },
+  activeModel: ModelOrMetadata,
+): ModelOrMetadata {
+  return "step" in toSave ? (toSave.step.model ?? activeModel) : activeModel;
+}
+
+type RawRequestResponseInclude = Record<string, boolean>;
+
+/**
+ * The raw handler promises request and, for non-streaming calls, response
+ * bodies. Preserve explicit caller choices while opting in to those bodies.
+ */
+function rawRequestResponseInclude(
+  args: {
+    experimental_include?: RawRequestResponseInclude;
+    include?: RawRequestResponseInclude;
+  },
+  enabled: boolean,
+  operation:
+    | "generateText"
+    | "streamText"
+    | "generateObject"
+    | "streamObject"
+    | undefined,
+): RawRequestResponseInclude | undefined {
+  if (!enabled) return undefined;
+
+  const requested = {
+    ...args.experimental_include,
+    ...args.include,
+  };
+  return {
+    ...requested,
+    requestBody: requested.requestBody ?? true,
+    ...((operation === "generateText" || operation === "generateObject")
+      ? { responseBody: requested.responseBody ?? true }
+      : {}),
+  };
+}
+
 export async function startGeneration<
   T,
   Tools extends ToolSet = ToolSet,
   CustomCtx extends object = object,
+  RUNTIME_CONTEXT extends Context = Context,
 >(
   ctx: ActionCtx & CustomCtx,
   component: AgentComponent,
@@ -74,13 +125,20 @@ export async function startGeneration<
      *   the promptMessageId message, if provided.
      */
     messages?: (ModelMessage | Message)[];
+    instructions?: Instructions;
+    /** @deprecated Use instructions. */
+    system?: Instructions;
+    allowSystemInMessages?: boolean;
     /**
      * The abort signal to be passed to the LLM call. If triggered, it will
      * mark the pending message as failed. If the generation is asynchronously
      * aborted, it will trigger this signal when detected.
      */
     abortSignal?: AbortSignal;
-    stopWhen?: StopCondition<Tools> | Array<StopCondition<Tools>>;
+    runtimeContext?: RUNTIME_CONTEXT;
+    stopWhen?:
+      | StopCondition<Tools, RUNTIME_CONTEXT>
+      | Array<StopCondition<Tools, RUNTIME_CONTEXT>>;
     _internal?: { generateId?: IdGenerator };
   },
   {
@@ -94,14 +152,20 @@ export async function startGeneration<
       agentName: string;
       agentForToolCtx?: Agent;
     },
+  _operation?:
+    | "generateText"
+    | "streamText"
+    | "generateObject"
+    | "streamObject",
 ): Promise<{
   args: T & {
-    system?: string;
+    instructions?: Instructions;
     model: LanguageModel;
     messages: ModelMessage[];
     prompt?: never;
     tools?: Tools;
-  } & CallSettings;
+    runtimeContext?: RUNTIME_CONTEXT;
+  } & AgentCallSettings;
   order: number;
   stepOrder: number;
   userId: string | undefined;
@@ -109,7 +173,10 @@ export async function startGeneration<
   updateModel: (model: ModelOrMetadata | undefined) => void;
   save: <TOOLS extends ToolSet>(
     toSave:
-      | { step: StepResult<TOOLS> }
+      | {
+          step: StepResult<TOOLS, RUNTIME_CONTEXT>;
+          responseMessages?: ModelMessage[];
+        }
       | { object: GenerateObjectResult<unknown> },
     createPendingMessage?: boolean,
     finishStreamId?: string,
@@ -132,6 +199,15 @@ export async function startGeneration<
     prompt: args.prompt,
     promptMessageId: args.promptMessageId,
   });
+  const allowSystemInMessages = args.allowSystemInMessages ?? true;
+  if (
+    !allowSystemInMessages &&
+    context.messages.some((message) => message.role === "system")
+  ) {
+    throw new Error(
+      "System messages in assembled message history require allowSystemInMessages: true. Use instructions for top-level system guidance.",
+    );
+  }
 
   const saveMessages = opts.storageOptions?.saveMessages ?? "promptAndOutput";
   const { promptMessageId, pendingMessage, savedMessages } =
@@ -159,13 +235,22 @@ export async function startGeneration<
   assert(model, "model is required");
   let activeModel: ModelOrMetadata = model;
 
-  const fail = async (reason: string) => {
-    if (pendingMessageId) {
-      await ctx.runMutation(component.messages.finalizeMessage, {
-        messageId: pendingMessageId,
-        result: { status: "failed", error: reason },
-      });
+  // Both the caller's AbortSignal listener and AI SDK's onAbort can report
+  // the same cancellation. Share the one finalization instead of racing two
+  // mutations for the pending message.
+  let pendingMessageFailure: Promise<void> | undefined;
+  const fail = (reason: string): Promise<void> => {
+    if (!pendingMessageId) return Promise.resolve();
+    if (!pendingMessageFailure) {
+      const messageId = pendingMessageId;
+      pendingMessageFailure = ctx
+        .runMutation(component.messages.finalizeMessage, {
+          messageId,
+          result: { status: "failed", error: reason },
+        })
+        .then(() => undefined);
     }
+    return pendingMessageFailure;
   };
   if (args.abortSignal) {
     const abortSignal = args.abortSignal;
@@ -185,34 +270,52 @@ export async function startGeneration<
     agent: opts.agentForToolCtx,
   } satisfies ToolCtx;
   const tools = wrapTools(toolCtx, args.tools) as Tools;
+  const argsWithoutSystem = omit(
+    args as typeof args & { system?: Instructions },
+    ["system"],
+  );
+  const {
+    promptMessageId: _promptMessageId,
+    messages: _messages,
+    prompt: _prompt,
+    instructions: _instructions,
+    onStepFinish: _onStepFinish,
+    ...aiCallArgs
+  } = argsWithoutSystem as typeof argsWithoutSystem & {
+    onStepFinish?: unknown;
+  };
+  const include = rawRequestResponseInclude(
+    aiCallArgs as {
+      experimental_include?: RawRequestResponseInclude;
+      include?: RawRequestResponseInclude;
+    },
+    Boolean(opts.rawRequestResponseHandler),
+    _operation,
+  );
   const aiArgs = {
     ...opts.callSettings,
     providerOptions: opts.providerOptions,
-    ...omit(args, ["promptMessageId", "messages", "prompt"]),
+    ...aiCallArgs,
     model,
     messages: context.messages,
+    instructions: args.instructions ?? args.system,
+    allowSystemInMessages,
     stopWhen:
-      args.stopWhen ?? (opts.maxSteps ? stepCountIs(opts.maxSteps) : undefined),
+      args.stopWhen ?? (opts.maxSteps ? isStepCount(opts.maxSteps) : undefined),
     tools,
-  } as T & {
+    ...(include ? { include } : {}),
+  } as unknown as T & {
     model: LanguageModel;
     messages: ModelMessage[];
     prompt?: never;
     tools?: Tools;
     _internal?: { generateId?: IdGenerator };
-  } & CallSettings;
+  } & AgentCallSettings;
   // NOTE: We intentionally do NOT override _internal.generateId here.
   // The AI SDK uses generateId() for many internal IDs (approval IDs,
   // tool execution IDs, message IDs, etc.) and they must be unique.
   // The pending message is linked via the explicit `pendingMessageId`
   // parameter passed to addMessages in the save closure.
-  // Track how many response messages we've already saved across steps.
-  // step.response.messages is cumulative — each step appends to it.
-  // We need to know which messages are new in each step to serialize
-  // only the new ones (important for tool approval flows where the SDK
-  // may add extra messages like approval tool-results).
-  let previousResponseMessageCount = 0;
-
   return {
     args: aiArgs,
     order: order ?? 0,
@@ -228,7 +331,10 @@ export async function startGeneration<
     fail,
     save: async <TOOLS extends ToolSet>(
       toSave:
-        | { step: StepResult<TOOLS> }
+        | {
+            step: StepResult<TOOLS, RUNTIME_CONTEXT>;
+            responseMessages?: ModelMessage[];
+          }
         | { object: GenerateObjectResult<unknown> },
       createPendingMessage?: boolean,
       /**
@@ -247,17 +353,20 @@ export async function startGeneration<
             activeModel,
           );
         } else {
-          const allResponseMessages = toSave.step.response.messages;
-          const newResponseMessages = allResponseMessages.slice(
-            previousResponseMessageCount,
-          );
-          previousResponseMessageCount = allResponseMessages.length;
+          const newResponseMessages =
+            toSave.responseMessages ?? toSave.step.response.messages;
+          // Even an empty completed step needs a durable assistant result so
+          // the pending message can be finalized at the storage boundary.
+          const responseMessagesToSave: ModelMessage[] =
+            newResponseMessages.length > 0
+              ? newResponseMessages
+              : [{ role: "assistant", content: [] }];
           serialized = await serializeResponseMessages(
             ctx,
             component,
             toSave.step,
             activeModel,
-            newResponseMessages,
+            responseMessagesToSave,
           );
         }
         const embeddings = await embedMessages(
@@ -294,6 +403,7 @@ export async function startGeneration<
             );
           } else {
             pendingMessageId = lastMessage._id;
+            pendingMessageFailure = undefined;
             savedMessages.push(...saved.messages.slice(0, -1));
           }
         } else {
@@ -312,12 +422,13 @@ export async function startGeneration<
         });
       }
       if (opts.usageHandler && output.usage) {
+        const usageModel = resolveUsageModel(toSave, activeModel);
         await opts.usageHandler(ctx, {
           userId,
           threadId,
           agentName: opts.agentName,
-          model: getModelName(activeModel),
-          provider: getProviderName(activeModel),
+          model: getModelName(usageModel),
+          provider: getProviderName(usageModel),
           usage: output.usage,
           providerMetadata: output.providerMetadata,
         });

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

@@ -0,0 +1,223 @@
+import { describe, expect, test } from "vitest";
+import { Agent, createThread } from "../index.js";
+import {
+  defineSchema,
+  type DataModelFromSchemaDefinition,
+  type ApiFromModules,
+  type ActionBuilder,
+  actionGeneric,
+  anyApi,
+} from "convex/server";
+import { v } from "convex/values";
+import { components, initConvexTest } from "./setup.test.js";
+import { mockModel } from "./mockModel.js";
+import { runAbortCleanup } from "./streamText.js";
+
+const schema = defineSchema({});
+type DataModel = DataModelFromSchemaDefinition<typeof schema>;
+const action = actionGeneric as ActionBuilder<DataModel, "public">;
+
+const FINAL_TEXT = "Hello from the model";
+
+const agent = new Agent(components.agent, {
+  name: "stream-test",
+  languageModel: mockModel({
+    content: [{ type: "text", text: FINAL_TEXT }],
+  }),
+});
+
+const emptyAgent = new Agent(components.agent, {
+  name: "empty-stream-test",
+  languageModel: mockModel({
+    content: [],
+    providerMetadata: { mock: { emptyResponse: true } },
+  }),
+});
+
+// Action that exercises streamText with saveStreamDeltas.returnImmediately=true.
+// It consumes the stream after streamText returns, simulating the HTTP response
+// path described in issue #265.
+export const streamTextReturnImmediately = action({
+  args: { threadId: v.string() },
+  handler: async (ctx, { threadId }) => {
+    const result = await agent.streamText(
+      ctx,
+      { threadId },
+      { prompt: "Test" },
+      {
+        saveStreamDeltas: {
+          returnImmediately: true,
+          chunking: "word",
+          throttleMs: 0,
+        },
+      },
+    );
+    // 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 };
+  },
+});
+
+export const streamTextEmptyAwaited = action({
+  args: { threadId: v.string() },
+  handler: async (ctx, { threadId }) => {
+    await emptyAgent.streamText(
+      ctx,
+      { threadId },
+      { prompt: "Test" },
+      { saveStreamDeltas: true },
+    );
+    return { ok: true };
+  },
+});
+
+export const streamTextEmptyReturnImmediately = action({
+  args: { threadId: v.string() },
+  handler: async (ctx, { threadId }) => {
+    const result = await emptyAgent.streamText(
+      ctx,
+      { threadId },
+      { prompt: "Test" },
+      {
+        saveStreamDeltas: {
+          returnImmediately: true,
+          throttleMs: 0,
+        },
+      },
+    );
+    await result.consumeStream();
+    return { ok: true };
+  },
+});
+
+const testApi: ApiFromModules<{
+  fns: {
+    streamTextReturnImmediately: typeof streamTextReturnImmediately;
+    streamTextEmptyAwaited: typeof streamTextEmptyAwaited;
+    streamTextEmptyReturnImmediately: typeof streamTextEmptyReturnImmediately;
+  };
+}>["fns"] = anyApi["streamText.test"] as any;
+
+describe("streamText with saveStreamDeltas.returnImmediately (issue #265)", () => {
+  test("persists the final assistant text to the messages table", async () => {
+    const t = initConvexTest(schema);
+    const threadId = await t.run(async (ctx) =>
+      createThread(ctx, components.agent, { userId: "u1" }),
+    );
+
+    await t.action(testApi.streamTextReturnImmediately, { threadId });
+
+    // Allow any background work scheduled by consumeStream to settle.
+    await t.finishAllScheduledFunctions(() => {});
+
+    const messages = await t.run(async (ctx) =>
+      agent.listMessages(ctx, {
+        threadId,
+        paginationOpts: { cursor: null, numItems: 50 },
+      }),
+    );
+
+    const assistantTextMessages = messages.page.filter(
+      (m) =>
+        m.message?.role === "assistant" &&
+        typeof m.text === "string" &&
+        m.text.length > 0,
+    );
+    expect(
+      assistantTextMessages.length,
+      "expected at least one persisted assistant message with text",
+    ).toBeGreaterThan(0);
+
+    const combined = assistantTextMessages.map((m) => m.text).join("");
+    expect(combined).toContain(FINAL_TEXT);
+
+    // The stream should be marked finished, not stuck in "streaming".
+    const stillStreaming = await t.run(async (ctx) =>
+      ctx.runQuery(components.agent.streams.list, {
+        threadId,
+        statuses: ["streaming"],
+      }),
+    );
+    expect(
+      stillStreaming,
+      "stream should not be stuck in 'streaming' status",
+    ).toHaveLength(0);
+  });
+});
+
+describe("streamText abort cleanup", () => {
+  test("attempts every cleanup and rethrows the first internal failure", async () => {
+    const calls: string[] = [];
+    const firstFailure = new Error("failed pending message cleanup");
+
+    await expect(
+      runAbortCleanup({
+        failCall: async () => {
+          calls.push("call.fail");
+          throw firstFailure;
+        },
+        failStreamer: async () => {
+          calls.push("streamer.fail");
+          throw new Error("failed stream cleanup");
+        },
+        onAbort: () => {
+          calls.push("user.onAbort");
+        },
+      }),
+    ).rejects.toBe(firstFailure);
+
+    expect(calls).toEqual(["call.fail", "streamer.fail", "user.onAbort"]);
+  });
+});
+
+describe("streamText with an empty final step (issue #274)", () => {
+  test.each([
+    ["awaited", testApi.streamTextEmptyAwaited],
+    ["returnImmediately", testApi.streamTextEmptyReturnImmediately],
+  ])(
+    "finalizes the pending assistant message in the %s path",
+    async (_, fn) => {
+      const t = initConvexTest(schema);
+      const threadId = await t.run(async (ctx) =>
+        createThread(ctx, components.agent, { userId: "u1" }),
+      );
+
+      await t.action(fn, { threadId });
+      await t.finishAllScheduledFunctions(() => {});
+
+      const messages = await t.run(async (ctx) =>
+        emptyAgent.listMessages(ctx, {
+          threadId,
+          paginationOpts: { cursor: null, numItems: 50 },
+        }),
+      );
+
+      expect(
+        messages.page.filter((message) => message.status === "pending"),
+      ).toHaveLength(0);
+      expect(messages.page).toContainEqual(
+        expect.objectContaining({
+          status: "success",
+          message: { role: "assistant", content: [] },
+          model: "mock-model-id",
+          provider: "mock-provider",
+          providerMetadata: { mock: { emptyResponse: true } },
+          usage: expect.objectContaining({
+            promptTokens: 3,
+            completionTokens: 10,
+            totalTokens: 13,
+          }),
+        }),
+      );
+
+      const stillStreaming = await t.run(async (ctx) =>
+        ctx.runQuery(components.agent.streams.list, {
+          threadId,
+          statuses: ["streaming"],
+        }),
+      );
+      expect(stillStreaming).toHaveLength(0);
+    },
+  );
+});