fix: track tokens endpoint

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";

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.
 </DynamicParamField>
 
-<DynamicParamField body="modelId" type="string" required>
+<DynamicParamField body="model_id" type="string" required>
   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`).
 </DynamicParamField>
 

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 '{

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 '{

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",

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,

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.",
+});

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<Record<string, unknown>> {
 	registerInternalSchemas(UpdateBalanceParamsV0Schema);
 	registerInternalSchemas(CheckParamsSchema);
 	registerInternalSchemas(TrackParamsSchema);
+	registerInternalSchemas(TrackTokensParamsSchema);
 	registerInternalSchemas(BillingResponseSchema);
 	registerInternalSchemas(AttachPreviewResponseSchema);
 	registerInternalSchemas(PreviewUpdateSubscriptionResponseSchema);

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

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<string, unknown>[],
+	runTrackWithRolloutCalls: [] as Record<string, unknown>[],
+	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<string, unknown>) => {
+		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<HonoEnv>();
+	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);
+	});
+});