diff --git a/apps/docs/api-reference-generator/balances/trackTokens.mdx b/apps/docs/api-reference-generator/balances/trackTokens.mdx index 9fed45b3e..f523f8d11 100644 --- a/apps/docs/api-reference-generator/balances/trackTokens.mdx +++ b/apps/docs/api-reference-generator/balances/trackTokens.mdx @@ -1,6 +1,6 @@ --- title: "Track Token Usage" -openapi: "openapi POST /v1/balances.trackTokens" +openapi: "openapi POST /v1/balances.track_tokens" --- import { DynamicParamField } from "/components/dynamic-param-field.jsx"; diff --git a/apps/docs/mintlify/api-reference/balances/trackTokens.mdx b/apps/docs/mintlify/api-reference/balances/trackTokens.mdx index 28f811123..a258f5df3 100644 --- a/apps/docs/mintlify/api-reference/balances/trackTokens.mdx +++ b/apps/docs/mintlify/api-reference/balances/trackTokens.mdx @@ -1,6 +1,6 @@ --- title: "Track Token Usage" -openapi: "openapi POST /v1/balances.trackTokens" +openapi: "openapi POST /v1/balances.track_tokens" --- import { DynamicParamField } from "/components/dynamic-param-field.jsx"; @@ -66,7 +66,7 @@ await autumn.balances.trackTokens({ The ID of the customer. - + The AI model in `provider/model` format, matching keys from [Models.dev](https://models.dev) (e.g., `anthropic/claude-opus-4-6`, `openai/gpt-4o`, `openrouter/anthropic/claude-opus-4.6`). diff --git a/apps/docs/mintlify/documentation/customers/tracking-usage.mdx b/apps/docs/mintlify/documentation/customers/tracking-usage.mdx index 90a65b734..0df65ea15 100644 --- a/apps/docs/mintlify/documentation/customers/tracking-usage.mdx +++ b/apps/docs/mintlify/documentation/customers/tracking-usage.mdx @@ -171,7 +171,7 @@ await autumn.balances.track_tokens( ``` ```bash cURL -curl -X POST "https://api.useautumn.com/v1/balances.trackTokens" \ +curl -X POST "https://api.useautumn.com/v1/balances.track_tokens" \ -H "Authorization: Bearer am_sk_test_1234" \ -H "Content-Type: application/json" \ -d '{ diff --git a/apps/docs/mintlify/documentation/modelling-pricing/credit-systems.mdx b/apps/docs/mintlify/documentation/modelling-pricing/credit-systems.mdx index b4382ec75..f0fa1d9f9 100644 --- a/apps/docs/mintlify/documentation/modelling-pricing/credit-systems.mdx +++ b/apps/docs/mintlify/documentation/modelling-pricing/credit-systems.mdx @@ -368,7 +368,7 @@ await autumn.balances.track_tokens( ``` ```bash cURL -curl -X POST "https://api.useautumn.com/v1/balances.trackTokens" \ +curl -X POST "https://api.useautumn.com/v1/balances.track_tokens" \ -H "Authorization: Bearer am_sk_test_1234" \ -H "Content-Type: application/json" \ -d '{ diff --git a/packages/openapi/v2.3/contracts/balancesContract.ts b/packages/openapi/v2.3/contracts/balancesContract.ts index 22f1bc0a0..c36d2b030 100644 --- a/packages/openapi/v2.3/contracts/balancesContract.ts +++ b/packages/openapi/v2.3/contracts/balancesContract.ts @@ -8,12 +8,14 @@ import { FinalizeLockParamsV0Schema, TrackParamsSchema, TrackResponseV3Schema, + TrackTokensParamsSchema, UpdateBalanceParamsV0Schema, } from "@autumn/shared"; import { oc } from "@orpc/contract"; import { balancesCheckJsDoc, balancesTrackJsDoc, + balancesTrackTokensJsDoc, } from "../jsDocs/balancesJsDocs"; type SpecWithResponses = { @@ -129,6 +131,63 @@ export const balancesTrackContract = oc }), ); +export const balancesTrackTokensContract = oc + .route({ + method: "POST", + path: "/v1/balances.track_tokens", + operationId: "trackTokens", + description: balancesTrackTokensJsDoc, + spec: (spec) => + withAcceptedResponse( + spec, + "trackTokens", + "Accepted. Autumn is experiencing degraded service from a downstream provider, so the token usage event was accepted for replay and will be tracked as soon as the service is restored.", + ), + }) + .input( + TrackTokensParamsSchema.meta({ + title: "TrackTokensParams", + examples: [ + { + customer_id: "cus_123", + feature_id: "ai_credits", + model_id: "anthropic/claude-sonnet-4-20250514", + input_tokens: 1000, + output_tokens: 500, + }, + ], + }), + ) + .output( + TrackResponseV3Schema.meta({ + examples: [ + { + customer_id: "cus_123", + value: 0.006, + balance: { + ...API_BALANCE_V1_EXAMPLE, + feature_id: "ai_credits", + granted: 10, + remaining: 9.994, + usage: 0.006, + }, + deductions: [ + { + balance_id: "cus_ent_3DdSDoyFmoA9Neecl2a2Gc507X2", + feature_id: "ai_credits", + plan_id: "pro", + reset: { + interval: "month", + resets_at: 1781288736881, + }, + value: 0.006, + }, + ], + }, + ], + }), + ); + export const balancesCreateContract = oc .route({ method: "POST", diff --git a/packages/openapi/v2.3/contracts/index.ts b/packages/openapi/v2.3/contracts/index.ts index a509f1c7e..18a5933a8 100644 --- a/packages/openapi/v2.3/contracts/index.ts +++ b/packages/openapi/v2.3/contracts/index.ts @@ -5,6 +5,7 @@ import { balancesDeleteContract, balancesFinalizeContract, balancesTrackContract, + balancesTrackTokensContract, balancesUpdateContract, } from "./balancesContract.js"; import { @@ -95,6 +96,7 @@ export const v2_3ContractRouter = oc.router({ balancesFinalize: balancesFinalizeContract, balancesCheck: balancesCheckContract, balancesTrack: balancesTrackContract, + balancesTrackTokens: balancesTrackTokensContract, // Events eventsList: eventsListContract, diff --git a/packages/openapi/v2.3/jsDocs/balancesJsDocs.ts b/packages/openapi/v2.3/jsDocs/balancesJsDocs.ts index 4cbe6bc9c..12ec2c81a 100644 --- a/packages/openapi/v2.3/jsDocs/balancesJsDocs.ts +++ b/packages/openapi/v2.3/jsDocs/balancesJsDocs.ts @@ -1,4 +1,8 @@ -import { ExtCheckParamsSchema, TrackParamsSchema } from "@autumn/shared"; +import { + ExtCheckParamsSchema, + TrackParamsSchema, + TrackTokensParamsSchema, +} from "@autumn/shared"; import { createJSDocDescription, example } from "../../utils/jsDocs/index.js"; export const balancesCheckJsDoc = createJSDocDescription({ @@ -58,3 +62,26 @@ export const balancesTrackJsDoc = createJSDocDescription({ returns: "The usage value recorded, with either a single updated balance or a map of updated balances. If Autumn is experiencing degraded service from a downstream provider, the API may return 202 after accepting the event for replay so it can be tracked as soon as the service is restored.", }); + +export const balancesTrackTokensJsDoc = createJSDocDescription({ + description: + "Records AI token usage for a customer and returns the updated AI credit balance.", + whenToUse: + "Use this after an LLM request when you have input and output token counts. Autumn converts token usage to a dollar amount using the configured model pricing and markup, then tracks that value against the customer's AI credit system.", + body: TrackTokensParamsSchema, + examples: [ + example({ + description: "Track one LLM response", + values: { + customerId: "cus_123", + featureId: "ai_credits", + modelId: "anthropic/claude-sonnet-4-20250514", + inputTokens: 1000, + outputTokens: 500, + }, + }), + ], + methodName: "trackTokens", + returns: + "The dollar value recorded and the updated AI credit system balance. If Autumn is experiencing degraded service from a downstream provider, the API may return 202 after accepting the token usage event for replay so it can be tracked as soon as the service is restored.", +}); diff --git a/packages/openapi/v2.3/openapi2.3.ts b/packages/openapi/v2.3/openapi2.3.ts index 68014d488..f591ca6cc 100644 --- a/packages/openapi/v2.3/openapi2.3.ts +++ b/packages/openapi/v2.3/openapi2.3.ts @@ -23,6 +23,7 @@ import { SetupPaymentResponseV1Schema, TrackParamsSchema, TrackResponseV3Schema, + TrackTokensParamsSchema, UpdateBalanceParamsV0Schema, UpdateSubscriptionV1ParamsSchema, } from "@autumn/shared"; @@ -64,6 +65,7 @@ async function generateOpenApiDocument(): Promise> { registerInternalSchemas(UpdateBalanceParamsV0Schema); registerInternalSchemas(CheckParamsSchema); registerInternalSchemas(TrackParamsSchema); + registerInternalSchemas(TrackTokensParamsSchema); registerInternalSchemas(BillingResponseSchema); registerInternalSchemas(AttachPreviewResponseSchema); registerInternalSchemas(PreviewUpdateSubscriptionResponseSchema); diff --git a/server/src/internal/balances/handlers/handleTrackTokens.ts b/server/src/internal/balances/handlers/handleTrackTokens.ts index 024baf635..43207020f 100644 --- a/server/src/internal/balances/handlers/handleTrackTokens.ts +++ b/server/src/internal/balances/handlers/handleTrackTokens.ts @@ -1,7 +1,11 @@ +import { + AffectedResource, + Scopes, + TrackTokensParamsSchema, +} from "@autumn/shared"; import { createRoute } from "@/honoMiddlewares/routeHandler.js"; -import { runTrackV2 } from "@/internal/balances/track/runTrackV2.js"; +import { runTrackWithRollout } from "@/internal/balances/track/runTrackWithRollout.js"; import { getTokenTrackParams } from "@/internal/balances/track/utils/getTokenTrackParams.js"; -import { AffectedResource, Scopes, TrackTokensParamsSchema } from "@autumn/shared"; export const handleTrackTokens = createRoute({ scopes: [Scopes.Balances.Write], @@ -16,12 +20,13 @@ export const handleTrackTokens = createRoute({ input: body, }); - return c.json( - await runTrackV2({ - ctx, - body: trackBody, - featureDeductions, - }), - ); + const response = await runTrackWithRollout({ + ctx, + body: trackBody, + featureDeductions, + }); + const status = ctx.extraLogs.trackQueuedForReplay ? 202 : 200; + + return c.json(response, status); }, }); diff --git a/server/tests/unit/balances/track/handle-track-tokens.test.ts b/server/tests/unit/balances/track/handle-track-tokens.test.ts new file mode 100644 index 000000000..decca2c78 --- /dev/null +++ b/server/tests/unit/balances/track/handle-track-tokens.test.ts @@ -0,0 +1,126 @@ +import { beforeEach, describe, expect, mock, test } from "bun:test"; +import { Hono } from "hono"; +import type { AutumnContext, HonoEnv } from "@/honoUtils/HonoEnv.js"; + +const mockState = { + getTokenTrackParamsCalls: [] as Record[], + runTrackWithRolloutCalls: [] as Record[], + queuedForReplay: false, +}; + +const trackBody = { + customer_id: "cus_123", + entity_id: "ent_123", + feature_id: "ai_credits", + value: 3.5, +}; + +const featureDeductions = [ + { + feature: { id: "ai_credits" }, + deduction: 1, + precomputedCreditCost: 3.5, + }, +]; + +mock.module("@/internal/balances/track/utils/getTokenTrackParams.js", () => ({ + getTokenTrackParams: async (args: Record) => { + mockState.getTokenTrackParamsCalls.push(args); + return { body: trackBody, featureDeductions }; + }, +})); + +mock.module("@/internal/balances/track/runTrackWithRollout.js", () => ({ + runTrackWithRollout: async (args: { + ctx: AutumnContext; + body: typeof trackBody; + featureDeductions: typeof featureDeductions; + }) => { + mockState.runTrackWithRolloutCalls.push(args); + if (mockState.queuedForReplay) { + args.ctx.extraLogs.trackQueuedForReplay = true; + } + return { + customer_id: args.body.customer_id, + entity_id: args.body.entity_id, + value: args.body.value, + balance: null, + }; + }, +})); + +import { handleTrackTokens } from "@/internal/balances/handlers/handleTrackTokens.js"; + +const requestBody = { + customer_id: "cus_123", + entity_id: "ent_123", + model_id: "openai/gpt-4.1", + input_tokens: 100, + output_tokens: 50, +}; + +const createApp = ({ ctx }: { ctx: AutumnContext }) => { + const app = new Hono(); + app.use("*", async (c, next) => { + c.set("ctx", ctx); + await next(); + }); + app.post("/track_tokens", ...handleTrackTokens); + return app; +}; + +const createCtx = (): AutumnContext => + ({ + features: [], + extraLogs: {}, + scopes: [], + skipCache: false, + }) as unknown as AutumnContext; + +describe("handleTrackTokens", () => { + beforeEach(() => { + mockState.getTokenTrackParamsCalls = []; + mockState.runTrackWithRolloutCalls = []; + mockState.queuedForReplay = false; + }); + + test("tracks converted token usage through the rollout path", async () => { + const ctx = createCtx(); + const response = await createApp({ ctx }).request("/track_tokens", { + method: "POST", + headers: { "content-type": "application/json" }, + body: JSON.stringify(requestBody), + }); + + expect(response.status).toBe(200); + expect(await response.json()).toEqual({ + customer_id: "cus_123", + entity_id: "ent_123", + value: 3.5, + balance: null, + }); + expect(mockState.getTokenTrackParamsCalls).toHaveLength(1); + expect(mockState.getTokenTrackParamsCalls[0]).toMatchObject({ + input: requestBody, + }); + expect(mockState.runTrackWithRolloutCalls).toHaveLength(1); + expect(mockState.runTrackWithRolloutCalls[0]).toMatchObject({ + body: trackBody, + featureDeductions, + }); + }); + + test("returns 202 when rollout fallback queues token tracking for replay", async () => { + mockState.queuedForReplay = true; + const ctx = createCtx(); + const response = await createApp({ ctx }).request("/track_tokens", { + method: "POST", + headers: { "content-type": "application/json" }, + body: JSON.stringify(requestBody), + }); + + expect(response.status).toBe(202); + expect(ctx.extraLogs.trackQueuedForReplay).toBe(true); + expect(mockState.runTrackWithRolloutCalls).toHaveLength(1); + }); +});