Agent: Added UsagePerRequest, Example: Added an example for custom cost component.

USAGE_TRACKING.md

@@ -0,0 +1,77 @@
+## Usage Tracking with `call.saveStepUsage`
+
+The new `call.saveStepUsage` function provides a clean way to save step usage data, similar to how `call.save` works for messages.
+
+### How it works
+
+When you use the `start` function from an agent, the returned `call` object now includes a `saveStepUsage` function:
+
+```typescript
+const { args, call, ... } = await agent.start(ctx, generateTextArgs, options);
+
+// In your onStepFinish callback:
+onStepFinish: async (step) => {
+  steps.push(step);
+  await call.save({ step }, createPendingMessage);
+
+  // Save usage data automatically
+  await call.saveStepUsage(step);
+
+  return args.onStepFinish?.(step);
+}
+```
+
+### What it does
+
+The `call.saveStepUsage` function:
+1. Checks if the step has usage data and required IDs (threadId, userId)
+2. Gets the most recently saved message from `call.getSavedMessages()`
+3. Serializes the AI SDK usage data using the existing `serializeUsage` function
+4. Calls a `pricePerRequest.create` mutation (if available) with:
+   - `messageId`: ID of the associated message
+   - `userId`: User who triggered the generation
+   - `threadId`: Thread where the generation occurred  
+   - `usage`: Serialized usage data (promptTokens, completionTokens, etc.)
+   - `calculatedAt`: Current timestamp
+   - `model`: Model name used for generation
+   - `provider`: Provider name used for generation
+
+### Usage data structure
+
+The usage data is automatically serialized from AI SDK format to Convex format:
+
+```typescript
+// AI SDK format
+step.usage: {
+  inputTokens: 100,
+  outputTokens: 50,
+  totalTokens: 150,
+  reasoningTokens?: 20,
+  cachedInputTokens?: 10
+}
+
+// Serialized to Convex format
+usage: {
+  promptTokens: 100,      // mapped from inputTokens
+  completionTokens: 50,   // mapped from outputTokens  
+  totalTokens: 150,
+  reasoningTokens: 20,    // optional
+  cachedInputTokens: 10   // optional
+}
+```
+
+### Implementation in streamText and generateText
+
+The `streamText` and `generateText` methods now automatically call `call.saveStepUsage(step)` in their `onStepFinish` callbacks, so usage data is tracked by default when you use these methods.
+
+### Setting up the pricePerRequest component
+
+To use this feature, you need to add a `pricePerRequest` component with a `create` mutation to your Convex schema. See the example implementation in the codebase.
+
+### Benefits
+
+- **Consistent with existing patterns**: Uses the same approach as `call.save`
+- **Automatic serialization**: Handles AI SDK to Convex format conversion
+- **Safe**: Only calls the mutation if the component exists
+- **Flexible**: Easy to extend with cost calculation logic
+- **Integrated**: Works automatically with existing `streamText` and `generateText` methods

example/convex/README.md

@@ -0,0 +1,90 @@
+# Welcome to your Convex functions directory!
+
+Write your Convex functions here.
+See https://docs.convex.dev/functions for more.
+
+A query function that takes two arguments looks like:
+
+```ts
+// convex/myFunctions.ts
+import { query } from "./_generated/server";
+import { v } from "convex/values";
+
+export const myQueryFunction = query({
+  // Validators for arguments.
+  args: {
+    first: v.number(),
+    second: v.string(),
+  },
+
+  // Function implementation.
+  handler: async (ctx, args) => {
+    // Read the database as many times as you need here.
+    // See https://docs.convex.dev/database/reading-data.
+    const documents = await ctx.db.query("tablename").collect();
+
+    // Arguments passed from the client are properties of the args object.
+    console.log(args.first, args.second);
+
+    // Write arbitrary JavaScript here: filter, aggregate, build derived data,
+    // remove non-public properties, or create new objects.
+    return documents;
+  },
+});
+```
+
+Using this query function in a React component looks like:
+
+```ts
+const data = useQuery(api.myFunctions.myQueryFunction, {
+  first: 10,
+  second: "hello",
+});
+```
+
+A mutation function looks like:
+
+```ts
+// convex/myFunctions.ts
+import { mutation } from "./_generated/server";
+import { v } from "convex/values";
+
+export const myMutationFunction = mutation({
+  // Validators for arguments.
+  args: {
+    first: v.string(),
+    second: v.string(),
+  },
+
+  // Function implementation.
+  handler: async (ctx, args) => {
+    // Insert or modify documents in the database here.
+    // Mutations can also read from the database like queries.
+    // See https://docs.convex.dev/database/writing-data.
+    const message = { body: args.first, author: args.second };
+    const id = await ctx.db.insert("messages", message);
+
+    // Optionally, return a value from your mutation.
+    return await ctx.db.get(id);
+  },
+});
+```
+
+Using this mutation function in a React component looks like:
+
+```ts
+const mutation = useMutation(api.myFunctions.myMutationFunction);
+function handleButtonPress() {
+  // fire and forget, the most common way to use mutations
+  mutation({ first: "Hello!", second: "me" });
+  // OR
+  // use the result once the mutation has completed
+  mutation({ first: "Hello!", second: "me" }).then((result) =>
+    console.log(result),
+  );
+}
+```
+
+Use the Convex CLI to push your functions to a deployment. See everything
+the Convex CLI can do by running `npx convex -h` in your project root
+directory. To learn more, launch the docs with `npx convex docs`.

example/convex/_generated/api.d.ts

@@ -19,6 +19,11 @@ import type * as chat_streamAbort from "../chat/streamAbort.js";
 import type * as chat_streaming from "../chat/streaming.js";
 import type * as chat_streamingReasoning from "../chat/streamingReasoning.js";
 import type * as chat_withoutAgent from "../chat/withoutAgent.js";
+import type * as cost__generated_api from "../cost/_generated/api.js";
+import type * as cost__generated_server from "../cost/_generated/server.js";
+import type * as cost_cost from "../cost/cost.js";
+import type * as cost_crons from "../cost/crons.js";
+import type * as cost_pricing from "../cost/pricing.js";
 import type * as crons from "../crons.js";
 import type * as debugging_rawRequestResponseHandler from "../debugging/rawRequestResponseHandler.js";
 import type * as files_addFile from "../files/addFile.js";
@@ -74,6 +79,11 @@ declare const fullApi: ApiFromModules<{
   "chat/streaming": typeof chat_streaming;
   "chat/streamingReasoning": typeof chat_streamingReasoning;
   "chat/withoutAgent": typeof chat_withoutAgent;
+  "cost/_generated/api": typeof cost__generated_api;
+  "cost/_generated/server": typeof cost__generated_server;
+  "cost/cost": typeof cost_cost;
+  "cost/crons": typeof cost_crons;
+  "cost/pricing": typeof cost_pricing;
   crons: typeof crons;
   "debugging/rawRequestResponseHandler": typeof debugging_rawRequestResponseHandler;
   "files/addFile": typeof files_addFile;
@@ -2772,6 +2782,27 @@ export declare const components: {
         }
       >;
     };
+    usagePerRequest: {
+      addUsage: FunctionReference<
+        "mutation",
+        "internal",
+        {
+          messageId: string;
+          model: string;
+          provider: string;
+          threadId: string;
+          usage: {
+            cachedInputTokens?: number;
+            completionTokens: number;
+            promptTokens: number;
+            reasoningTokens?: number;
+            totalTokens: number;
+          };
+          userId?: string;
+        },
+        any
+      >;
+    };
     users: {
       deleteAllForUserId: FunctionReference<
         "action",
@@ -3277,6 +3308,43 @@ export declare const components: {
       getServerTime: FunctionReference<"mutation", "internal", {}, number>;
     };
   };
+  cost: {
+    cost: {
+      addCost: FunctionReference<
+        "action",
+        "internal",
+        {
+          messageId: string;
+          modelId: string;
+          threadId: string;
+          usage: {
+            cachedInputTokens?: number;
+            completionTokens: number;
+            promptTokens: number;
+            reasoningTokens?: number;
+            totalTokens: number;
+          };
+          userId?: string;
+        },
+        any
+      >;
+    };
+    pricing: {
+      getAllPricing: FunctionReference<"query", "internal", {}, any>;
+      getPricing: FunctionReference<
+        "query",
+        "internal",
+        { modelId: string },
+        any
+      >;
+      getPricingByProvider: FunctionReference<
+        "query",
+        "internal",
+        { providerId: string },
+        any
+      >;
+    };
+  };
   rag: {
     chunks: {
       insert: FunctionReference<

example/convex/chat/basic.ts

@@ -16,7 +16,28 @@ export const generateTextInAnAction = action({
   args: { prompt: v.string(), threadId: v.string() },
   handler: async (ctx, { prompt, threadId }) => {
     await authorizeThreadAccess(ctx, threadId);
-    const result = await agent.generateText(ctx, { threadId }, { prompt });
+    const result = await agent.generateText(
+      ctx,
+      { threadId },
+      {
+        prompt,
+        onStepFinish: async (step, args) => {
+          await ctx.runAction(components.cost.cost.addCost, {
+            messageId: args?.messageId!,
+            userId: args?.userId,
+            threadId: args?.threadId!,
+            modelId: step.response.modelId,
+            usage: {
+              reasoningTokens: step.usage.reasoningTokens,
+              cachedInputTokens: step.usage.cachedInputTokens,
+              completionTokens: step.usage.outputTokens as number,
+              promptTokens: step.usage.inputTokens as number,
+              totalTokens: step.usage.totalTokens as number,
+            },
+          });
+        },
+      },
+    );
     return result.text;
   },
 });
@@ -48,7 +69,28 @@ export const sendMessage = mutation({
 export const generateResponse = internalAction({
   args: { promptMessageId: v.string(), threadId: v.string() },
   handler: async (ctx, { promptMessageId, threadId }) => {
-    await agent.generateText(ctx, { threadId }, { promptMessageId });
+    await agent.generateText(
+      ctx,
+      { threadId },
+      {
+        promptMessageId,
+        onStepFinish: async (step, args) => {
+          await ctx.runAction(components.cost.cost.addCost, {
+            messageId: args?.messageId!,
+            userId: args?.userId,
+            threadId: args?.threadId!,
+            modelId: step.response.modelId,
+            usage: {
+              reasoningTokens: step.usage.reasoningTokens,
+              cachedInputTokens: step.usage.cachedInputTokens,
+              completionTokens: step.usage.outputTokens as number,
+              promptTokens: step.usage.inputTokens as number,
+              totalTokens: step.usage.totalTokens as number,
+            },
+          });
+        },
+      },
+    );
   },
 });
 

example/convex/chat/streaming.ts

@@ -27,7 +27,7 @@ export const streamOneShot = action({
   args: { prompt: v.string(), threadId: v.string() },
   handler: async (ctx, { prompt, threadId }) => {
     await authorizeThreadAccess(ctx, threadId);
-    await storyAgent.streamText(
+    const message = await storyAgent.streamText(
       ctx,
       { threadId },
       { prompt },
@@ -68,7 +68,24 @@ export const streamAsync = internalAction({
     const result = await storyAgent.streamText(
       ctx,
       { threadId },
-      { promptMessageId },
+      {
+        promptMessageId,
+        onStepFinish: async (step, args) => {
+          await ctx.runAction(components.cost.cost.addCost, {
+            messageId: args?.messageId!,
+            userId: args?.userId,
+            threadId: args?.threadId!,
+            modelId: step.response.modelId,
+            usage: {
+              reasoningTokens: step.usage.reasoningTokens,
+              cachedInputTokens: step.usage.cachedInputTokens,
+              completionTokens: step.usage.outputTokens as number,
+              promptTokens: step.usage.inputTokens as number,
+              totalTokens: step.usage.totalTokens as number,
+            },
+          });
+        },
+      },
       // more custom delta options (`true` uses defaults)
       { saveStreamDeltas: { chunking: "word", throttleMs: 100 } },
     );

example/convex/convex.config.ts

@@ -3,11 +3,13 @@ import agent from "@convex-dev/agent/convex.config";
 import workflow from "@convex-dev/workflow/convex.config";
 import rateLimiter from "@convex-dev/rate-limiter/convex.config";
 import rag from "@convex-dev/rag/convex.config";
+import cost from "./cost/convex.config.js";
 
 const app = defineApp();
 app.use(agent);
 app.use(workflow);
 app.use(rateLimiter);
+app.use(cost);
 app.use(rag);
 
 export default app;