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

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

@@ -1,4 +1,9 @@
-import type { ToolResultOutput } from "@ai-sdk/provider-utils";
+import type { JSONObject } from "@ai-sdk/provider";
+import type {
+  Context,
+  Experimental_SandboxSession,
+  ToolResultOutput,
+} from "@ai-sdk/provider-utils";
 import type {
   FlexibleSchema,
   ModelMessage,
@@ -8,8 +13,8 @@ import type {
 } from "ai";
 import { tool } from "ai";
 import type { GenericActionCtx, GenericDataModel } from "convex/server";
-import type { ProviderOptions } from "../validators.js";
-import type { Agent } from "./index.js";
+import type { ProviderOptions } from "../../validators.js";
+import type { Agent } from "../index.js";
 
 const MIGRATION_URL = "node_modules/@convex-dev/agent/MIGRATION.md";
 const warnedDeprecations = new Set<string>();
@@ -28,12 +33,16 @@ export type ToolCtx<DataModel extends GenericDataModel = GenericDataModel> =
     messageId?: string;
   };
 
+export type AgentToolExecutionOptions<TOOL_CONTEXT extends Context = Context> =
+  ToolExecutionOptions<TOOL_CONTEXT>;
+
 /**
  * Function that is called to determine if the tool needs approval before it can be executed.
  */
 export type ToolNeedsApprovalFunctionCtx<
   INPUT,
   Ctx extends ToolCtx = ToolCtx,
+  TOOL_CONTEXT extends Context = Context,
 > = (
   ctx: Ctx,
   input: INPUT,
@@ -52,7 +61,7 @@ export type ToolNeedsApprovalFunctionCtx<
      *
      * Experimental (can break in patch releases).
      */
-    experimental_context?: unknown;
+    context: TOOL_CONTEXT;
   },
 ) => boolean | PromiseLike<boolean>;
 
@@ -60,10 +69,11 @@ export type ToolExecuteFunctionCtx<
   INPUT,
   OUTPUT,
   Ctx extends ToolCtx = ToolCtx,
+  TOOL_CONTEXT extends Context = Context,
 > = (
   ctx: Ctx,
   input: INPUT,
-  options: ToolExecutionOptions,
+  options: AgentToolExecutionOptions<TOOL_CONTEXT>,
 ) => AsyncIterable<OUTPUT> | PromiseLike<OUTPUT>;
 
 type NeverOptional<N, T> = 0 extends 1 & N
@@ -83,6 +93,7 @@ export type ToolOutputPropertiesCtx<
   INPUT,
   OUTPUT,
   Ctx extends ToolCtx = ToolCtx,
+  TOOL_CONTEXT extends Context = Context,
 > = NeverOptional<
   OUTPUT,
   {
@@ -93,7 +104,7 @@ export type ToolOutputPropertiesCtx<
      * @param input - The input of the tool call.
      * @param options.abortSignal - A signal that can be used to abort the tool call.
      */
-    execute?: ToolExecuteFunctionCtx<INPUT, OUTPUT, Ctx>;
+    execute?: ToolExecuteFunctionCtx<INPUT, OUTPUT, Ctx, TOOL_CONTEXT>;
     outputSchema?: FlexibleSchema<OUTPUT>;
     /**
      * @deprecated Removed in v0.6.0. Use `execute` instead.
@@ -130,17 +141,27 @@ export type ToolInputProperties<INPUT> = {
  * @param tool The tool. See https://sdk.vercel.ai/docs/ai-sdk-core/tools-and-tool-calling
  * Currently contains deprecated parameters `args` and `handler` to maintain backwards compatibility
  * but these will be removed in the future. Use `inputSchema` and `execute` instead, respectively.
- * 
+ *
  * @returns A tool to be used with the AI SDK.
  */
-export function createTool<INPUT, OUTPUT, Ctx extends ToolCtx = ToolCtx>(
+export function createTool<
+  INPUT,
+  OUTPUT,
+  Ctx extends ToolCtx = ToolCtx,
+  TOOL_CONTEXT extends Context = Context,
+>(
   def: {
     /**
      * An optional description of what the tool does.
      * Will be used by the language model to decide whether to use the tool.
      * Not used for provider-defined tools.
      */
-    description?: string;
+    description?:
+      | string
+      | ((options: {
+          context: NoInfer<TOOL_CONTEXT>;
+          experimental_sandbox?: Experimental_SandboxSession;
+        }) => string);
     /**
      * An optional title of the tool.
      */
@@ -151,6 +172,8 @@ export function createTool<INPUT, OUTPUT, Ctx extends ToolCtx = ToolCtx>(
      * functionality that can be fully encapsulated in the provider.
      */
     providerOptions?: ProviderOptions;
+    metadata?: JSONObject;
+    contextSchema?: FlexibleSchema<TOOL_CONTEXT>;
   } & ToolInputProperties<INPUT> & {
       /**
        * An optional list of input examples that show the language
@@ -166,7 +189,8 @@ export function createTool<INPUT, OUTPUT, Ctx extends ToolCtx = ToolCtx>(
         | boolean
         | ToolNeedsApprovalFunctionCtx<
             [INPUT] extends [never] ? unknown : INPUT,
-            Ctx
+            Ctx,
+            TOOL_CONTEXT
           >;
       /**
        * Strict mode setting for the tool.
@@ -186,7 +210,7 @@ export function createTool<INPUT, OUTPUT, Ctx extends ToolCtx = ToolCtx>(
        */
       onInputStart?: (
         ctx: Ctx,
-        options: ToolExecutionOptions,
+        options: AgentToolExecutionOptions<TOOL_CONTEXT>,
       ) => void | PromiseLike<void>;
       /**
        * Optional function that is called when an argument streaming delta is available.
@@ -194,7 +218,9 @@ export function createTool<INPUT, OUTPUT, Ctx extends ToolCtx = ToolCtx>(
        */
       onInputDelta?: (
         ctx: Ctx,
-        options: { inputTextDelta: string } & ToolExecutionOptions,
+        options: {
+          inputTextDelta: string;
+        } & AgentToolExecutionOptions<TOOL_CONTEXT>,
       ) => void | PromiseLike<void>;
       /**
        * Optional function that is called when a tool call can be started,
@@ -204,9 +230,9 @@ export function createTool<INPUT, OUTPUT, Ctx extends ToolCtx = ToolCtx>(
         ctx: Ctx,
         options: {
           input: [INPUT] extends [never] ? unknown : INPUT;
-        } & ToolExecutionOptions,
+        } & AgentToolExecutionOptions<TOOL_CONTEXT>,
       ) => void | PromiseLike<void>;
-    } & ToolOutputPropertiesCtx<INPUT, OUTPUT, Ctx> & {
+    } & ToolOutputPropertiesCtx<INPUT, OUTPUT, Ctx, TOOL_CONTEXT> & {
       /**
        * Optional conversion function that maps the tool result to an output that can be used by the language model.
        *
@@ -234,7 +260,7 @@ export function createTool<INPUT, OUTPUT, Ctx extends ToolCtx = ToolCtx>(
         },
       ) => ToolResultOutput | PromiseLike<ToolResultOutput>;
     },
-): Tool<INPUT, OUTPUT> {
+): Tool<INPUT, OUTPUT, TOOL_CONTEXT> {
   // Runtime backwards compat - types will show errors but runtime still works
   const inputSchema = def.inputSchema ?? (def as any).args;
   if (!inputSchema)
@@ -260,16 +286,18 @@ export function createTool<INPUT, OUTPUT, Ctx extends ToolCtx = ToolCtx>(
         " handler function, define an outputSchema, or both",
     );
 
-  const t = tool<INPUT, OUTPUT>({
+  const t = tool<INPUT, OUTPUT, TOOL_CONTEXT>({
     type: "function",
     __acceptsCtx: true,
     ctx: def.ctx,
     description: def.description,
     title: def.title,
     providerOptions: def.providerOptions,
+    metadata: def.metadata,
     inputSchema,
+    contextSchema: def.contextSchema,
     inputExamples: def.inputExamples,
-    needsApproval(this: Tool<INPUT, OUTPUT>, input, options) {
+    needsApproval(this: Tool<INPUT, OUTPUT, TOOL_CONTEXT>, input, options) {
       const needsApproval = def.needsApproval;
       if (!needsApproval || typeof needsApproval === "boolean")
         return Boolean(needsApproval);
@@ -287,9 +315,9 @@ export function createTool<INPUT, OUTPUT, Ctx extends ToolCtx = ToolCtx>(
     ...(executeHandler
       ? {
           execute(
-            this: Tool<INPUT, OUTPUT>,
+            this: Tool<INPUT, OUTPUT, TOOL_CONTEXT>,
             input: INPUT,
-            options: ToolExecutionOptions,
+            options: ToolExecutionOptions<TOOL_CONTEXT>,
           ) {
             if (!getCtx(this)) {
               throw new Error(
@@ -306,25 +334,37 @@ export function createTool<INPUT, OUTPUT, Ctx extends ToolCtx = ToolCtx>(
   });
   if (def.onInputStart) {
     const origOnInputStart = def.onInputStart;
-    t.onInputStart = function (this: Tool<INPUT, OUTPUT>, options) {
+    t.onInputStart = function (
+      this: Tool<INPUT, OUTPUT, TOOL_CONTEXT>,
+      options,
+    ) {
       return origOnInputStart.call(this, getCtx(this), options);
     };
   }
   if (def.onInputDelta) {
     const origOnInputDelta = def.onInputDelta;
-    t.onInputDelta = function (this: Tool<INPUT, OUTPUT>, options) {
+    t.onInputDelta = function (
+      this: Tool<INPUT, OUTPUT, TOOL_CONTEXT>,
+      options,
+    ) {
       return origOnInputDelta.call(this, getCtx(this), options);
     };
   }
   if (def.onInputAvailable) {
     const origOnInputAvailable = def.onInputAvailable;
-    t.onInputAvailable = function (this: Tool<INPUT, OUTPUT>, options) {
+    t.onInputAvailable = function (
+      this: Tool<INPUT, OUTPUT, TOOL_CONTEXT>,
+      options,
+    ) {
       return origOnInputAvailable.call(this, getCtx(this), options);
     };
   }
   if (def.toModelOutput) {
     const origToModelOutput = def.toModelOutput;
-    t.toModelOutput = function (this: Tool<INPUT, OUTPUT>, options) {
+    t.toModelOutput = function (
+      this: Tool<INPUT, OUTPUT, TOOL_CONTEXT>,
+      options,
+    ) {
       return origToModelOutput.call(this, getCtx(this), options);
     };
   }

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

@@ -6,30 +6,32 @@ import {
   type ApiFromModules,
   type GenericActionCtx,
   type GenericDataModel,
+  type GenericMutationCtx,
   type GenericQueryCtx,
 } from "convex/server";
 import { v } from "convex/values";
 import {
-  createThread as createThread_,
-  listMessages as listMessages_,
-  toModelMessage,
   vContextOptions,
   vMessage,
   vMessageDoc,
   vPaginationResult,
   vStorageOptions,
   vThreadDoc,
-  type Agent,
-  type AgentComponent,
   vStreamArgs,
-  syncStreams,
-  vStreamMessagesReturnValue,
-  isTool,
-  extractText,
   type MessageDoc,
-} from "./index.js";
-import { serializeNewMessagesInStep } from "../mapping.js";
-import { getModelName, getProviderName } from "../shared.js";
+} from "../../validators.js";
+import { createThread as createThread_ } from "../../client/threads.js";
+import {
+  extractText,
+  getModelName,
+  getProviderName,
+  isTool,
+} from "../../shared.js";
+import { serializeResponseMessages, toModelMessage } from "../mapping.js";
+import type { Agent } from "../index.js";
+import { listMessages as listMessages_ } from "./messages.js";
+import { syncStreams, vStreamMessagesReturnValue } from "./streaming.js";
+import type { AgentComponent } from "./types.js";
 
 export type PlaygroundAPI = ApiFromModules<{
   playground: ReturnType<typeof definePlaygroundAPI>;
@@ -64,7 +66,10 @@ export function definePlaygroundAPI<DataModel extends GenericDataModel>(
     }
   }
 
-  async function validateApiKey(ctx: RunQueryCtx, apiKey: string) {
+  async function validateApiKey(
+    ctx: QueryCtx | MutationCtx | ActionCtx,
+    apiKey: string,
+  ) {
     await ctx.runQuery(component.apiKeys.validate, { apiKey });
   }
 
@@ -276,18 +281,20 @@ export function definePlaygroundAPI<DataModel extends GenericDataModel>(
         },
         { contextOptions, storageOptions, saveStreamDeltas: true },
       );
-      const outputMessages = await Promise.all(
-        (await steps).map(async (step) => {
-          const { messages } = await serializeNewMessagesInStep(
-            ctx,
-            component,
-            step,
-            {
-              model: getModelName(agent.options.languageModel),
-              provider: getProviderName(agent.options.languageModel),
-            },
-          );
-          return messages.map((messageWithMetadata, i) => {
+      const outputMessages: MessageDoc[][] = [];
+      for (const step of await steps) {
+        const { messages } = await serializeResponseMessages(
+          ctx,
+          component,
+          step,
+          {
+            model: getModelName(agent.options.languageModel),
+            provider: getProviderName(agent.options.languageModel),
+          },
+          step.response.messages,
+        );
+        outputMessages.push(
+          messages.map((messageWithMetadata, i) => {
             return {
               ...messageWithMetadata,
               tool: isTool(messageWithMetadata.message),
@@ -300,9 +307,9 @@ export function definePlaygroundAPI<DataModel extends GenericDataModel>(
               order: 0,
               stepOrder: i + 1,
             } satisfies MessageDoc;
-          });
-        }),
-      );
+          }),
+        );
+      }
       return { text: await text, messages: outputMessages.flat() };
     },
     returns: v.object({ text: v.string(), messages: v.array(vMessageDoc) }),
@@ -367,4 +374,12 @@ export function definePlaygroundAPI<DataModel extends GenericDataModel>(
   };
 }
 
-type RunQueryCtx = { runQuery: GenericQueryCtx<GenericDataModel>["runQuery"] };
+type QueryCtx = Pick<GenericQueryCtx<GenericDataModel>, "runQuery">;
+type MutationCtx = Pick<
+  GenericMutationCtx<GenericDataModel>,
+  "runQuery" | "runMutation"
+>;
+type ActionCtx = Pick<
+  GenericActionCtx<GenericDataModel>,
+  "runQuery" | "runMutation" | "runAction"
+>;

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

@@ -5,14 +5,14 @@ import type {
   ModelMessage,
   UserContent,
 } from "ai";
-import type { Id } from "../component/_generated/dataModel.js";
+import type { Id } from "../../component/_generated/dataModel.js";
 import type {
   ActionCtx,
   AgentComponent,
   MutationCtx,
   QueryCtx,
 } from "./types.js";
-import type { Message } from "../validators.js";
+import type { Message } from "../../validators.js";
 import { assert } from "convex-helpers";
 import type { StorageReader } from "convex/server";
 

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

@@ -2,9 +2,10 @@ import { describe, expect, test } from "vitest";
 import {
   Agent,
   createThread,
+  createTool,
   filterOutOrphanedToolMessages,
   type MessageDoc,
-} from "./index.js";
+} from "../index.js";
 import type { DataModelFromSchemaDefinition } from "convex/server";
 import {
   anyApi,
@@ -23,7 +24,7 @@ import { defineSchema } from "convex/server";
 import { stepCountIs } from "ai";
 import { components, initConvexTest } from "./setup.test.js";
 import { z } from "zod/v4";
-import { mockModel } from "./mockModel.js";
+import { MockLanguageModel, mockModel } from "./mockModel.js";
 
 const schema = defineSchema({});
 type DataModel = DataModelFromSchemaDefinition<typeof schema>;
@@ -34,12 +35,50 @@ const action = actionGeneric as ActionBuilder<DataModel, "public">;
 
 const TEST_TEXT = JSON.stringify({ hello: "world" });
 
+const agentModel = new MockLanguageModel({
+  content: [{ type: "text", text: TEST_TEXT }],
+});
 const agent = new Agent(components.agent, {
   name: "test",
   instructions: "You are a test agent",
+  languageModel: agentModel,
+});
+
+let capturedRawBodies: { request?: unknown; response?: unknown } | undefined;
+const rawBodyAgent = new Agent(components.agent, {
+  name: "raw-body-test",
   languageModel: mockModel({
-    content: [{ type: "text", text: TEST_TEXT }],
+    doGenerate: async () => ({
+      content: [{ type: "text", text: "raw" }],
+      finishReason: { unified: "stop", raw: undefined },
+      usage: {
+        inputTokens: { total: 1, noCache: 1, cacheRead: 0, cacheWrite: 0 },
+        outputTokens: { total: 1, text: 1, reasoning: 0 },
+      },
+      request: { body: { prompt: "request-body" } },
+      response: { body: { text: "response-body" }, headers: {} },
+      warnings: [],
+    }),
   }),
+  rawRequestResponseHandler: async (_ctx, event) => {
+    capturedRawBodies = {
+      request: event.request.body,
+      response: event.response.body,
+    };
+  },
+});
+
+export const captureRawBodies = action({
+  args: {},
+  handler: async (ctx) => {
+    capturedRawBodies = undefined;
+    await rawBodyAgent.generateText(
+      ctx,
+      { userId: "raw-body-user" },
+      { prompt: "raw" },
+    );
+    return capturedRawBodies;
+  },
 });
 
 export const testQuery = query({
@@ -62,6 +101,77 @@ export const createThreadManually = mutation({
   },
 });
 
+const saveStepAgent = new Agent(components.agent, {
+  name: "save-step-test",
+  instructions: "test",
+  tools: {
+    echo: createTool({
+      description: "Echo a value",
+      inputSchema: z.object({ value: z.string() }),
+      execute: async (_ctx, input) => `echo:${input.value}`,
+    }),
+  },
+  languageModel: mockModel({
+    contentSteps: [
+      [
+        {
+          type: "tool-call",
+          toolCallId: "ss-1",
+          toolName: "echo",
+          input: JSON.stringify({ value: "hi" }),
+        },
+      ],
+      [{ type: "text", text: "done" }],
+    ],
+  }),
+  stopWhen: stepCountIs(5),
+});
+
+export const replayStepsViaSaveStep = action({
+  args: { withPreviousStep: v.boolean() },
+  handler: async (ctx, args) => {
+    const { thread } = await saveStepAgent.createThread(ctx, {
+      userId: "ss-gen",
+    });
+    const genResult = await thread.generateText({ prompt: "echo hi" });
+    const steps = genResult.steps;
+
+    const { threadId } = await saveStepAgent.createThread(ctx, {
+      userId: "ss-replay",
+    });
+    const { messageId: promptMessageId } = await saveStepAgent.saveMessage(
+      ctx,
+      {
+        threadId,
+        message: { role: "user", content: "echo hi" },
+        skipEmbeddings: true,
+      },
+    );
+    let previousStep: (typeof steps)[number] | undefined;
+    for (const step of steps) {
+      await saveStepAgent.saveStep(ctx, {
+        threadId,
+        promptMessageId,
+        step,
+        previousStep: args.withPreviousStep ? previousStep : undefined,
+      });
+      previousStep = step;
+    }
+
+    const replayed = await saveStepAgent.listMessages(ctx, {
+      threadId,
+      paginationOpts: { cursor: null, numItems: 50 },
+      statuses: ["success", "pending", "failed"],
+    });
+    const contentTypes = replayed.page.flatMap((m) =>
+      Array.isArray(m.message?.content)
+        ? m.message!.content.map((c: { type?: string }) => c.type ?? "text")
+        : ["text"],
+    );
+    return { stepCount: steps.length, contentTypes };
+  },
+});
+
 export const createThreadMutation = agent.createThreadMutation();
 export const generateObjectAction = agent.asObjectAction({
   schema: z.object({ hello: z.string().describe("A string for testing") }),
@@ -162,6 +272,8 @@ const testApi: ApiFromModules<{
     generateTextAction: typeof generateTextAction;
     generateObjectAction: typeof generateObjectAction;
     saveMessageMutation: typeof saveMessageMutation;
+    captureRawBodies: typeof captureRawBodies;
+    replayStepsViaSaveStep: typeof replayStepsViaSaveStep;
   };
 }>["fns"] = anyApi["index.test"] as any;
 
@@ -177,6 +289,26 @@ describe("Agent thick client", () => {
     expect(result).toBeDefined();
     expect(result).toMatch(TEST_TEXT);
   });
+  test.each([true, false])(
+    "saveStep persists each SDK 7 step once (previousStep: %s)",
+    async (withPreviousStep) => {
+      const t = initConvexTest(schema);
+      const res = await t.action(testApi.replayStepsViaSaveStep, {
+        withPreviousStep,
+      });
+      expect(res.stepCount).toBe(2);
+      const toolCalls = res.contentTypes.filter(
+        (t) => t === "tool-call",
+      ).length;
+      const toolResults = res.contentTypes.filter(
+        (t) => t === "tool-result",
+      ).length;
+      expect({ toolCalls, toolResults }).toEqual({
+        toolCalls: 1,
+        toolResults: 1,
+      });
+    },
+  );
 });
 
 describe("filterOutOrphanedToolMessages", () => {
@@ -244,6 +376,14 @@ describe("filterOutOrphanedToolMessages", () => {
 });
 
 describe("Agent option variations and normal behavior", () => {
+  test("raw handler opts into retained SDK 7 request and response bodies", async () => {
+    const t = initConvexTest(schema);
+    await expect(t.action(testApi.captureRawBodies, {})).resolves.toEqual({
+      request: { prompt: "request-body" },
+      response: { text: "response-body" },
+    });
+  });
+
   test("Agent can be constructed with minimal options", () => {
     const a = new Agent(components.agent, {
       name: "minimal",
@@ -378,6 +518,25 @@ describe("Agent-generated mutations/actions/queries", () => {
     expect(objResult.object).toBeDefined();
   });
 
+  test("asTextAction maps its system wire field to AI SDK instructions", async () => {
+    const t = initConvexTest(schema);
+    const before = agentModel.doGenerateCalls.length;
+
+    await t.action(testApi.generateTextAction, {
+      userId: "8",
+      system: "Action-scoped instructions",
+      prompt: "Say hi",
+    });
+
+    const call = agentModel.doGenerateCalls.at(-1);
+    expect(agentModel.doGenerateCalls).toHaveLength(before + 1);
+    expect(call?.prompt[0]).toEqual({
+      role: "system",
+      content: "Action-scoped instructions",
+    });
+    expect(call).not.toHaveProperty("system");
+  });
+
   test("asSaveMessagesMutation works via t.mutation", async () => {
     const t = initConvexTest(schema);
     const threadId = await t.run(async (ctx) =>