chore: minor ai credit system cleanup

.gitignore

@@ -22,6 +22,7 @@ supabase.sh
 tests/
 !server/tests
 !packages/mcp/tests
+!packages/ai-sdk/tests
 !apps/leaf/tests
 !vite/tests
 .secrets

apps/docs/mintlify/documentation/external-providers/ai-sdk.mdx

@@ -33,16 +33,16 @@ bun add @useautumn/ai-sdk
 
 #### 2. Wrap your model
 
-Use `withTokenTracking` to wrap any AI SDK language model. It intercepts generate and stream calls, reads the token usage from the response, and reports it to Autumn automatically.
+Use `withAutumn` to wrap any AI SDK language model. It intercepts generate and stream calls, reads the token usage from the response, and reports it to Autumn automatically.
 
 ```typescript
 import { Autumn } from "autumn-js";
 import { anthropic } from "@ai-sdk/anthropic";
-import { withTokenTracking } from "@useautumn/ai-sdk";
+import { withAutumn } from "@useautumn/ai-sdk";
 
 const autumn = new Autumn({ secretKey: "am_sk_test_1234" });
 
-const model = withTokenTracking({
+const model = withAutumn({
   autumn,
   model: anthropic("claude-sonnet-4-5-20250514"),
   customerId: "user_123",
@@ -89,7 +89,7 @@ import { createOpenRouter } from "@openrouter/ai-sdk-provider";
 
 const openrouter = createOpenRouter();
 
-const model = withTokenTracking({
+const model = withAutumn({
   autumn,
   model: openrouter("anthropic/claude-opus-4-6"),
   customerId: "user_123",
@@ -116,12 +116,12 @@ const model = withTokenTracking({
 import { Autumn } from "autumn-js";
 import { openai } from "@ai-sdk/openai";
 import { generateText } from "ai";
-import { withTokenTracking } from "@useautumn/ai-sdk";
+import { withAutumn } from "@useautumn/ai-sdk";
 
 const autumn = new Autumn({ secretKey: process.env.AUTUMN_SECRET_KEY! });
 
 async function chat(customerId: string, message: string) {
-  const model = withTokenTracking({
+  const model = withAutumn({
     autumn,
     model: openai("gpt-4o"),
     customerId,

bun.lock

@@ -1,5 +1,6 @@
 {
   "lockfileVersion": 1,
+  "configVersion": 0,
   "workspaces": {
     "": {
       "name": "autumn",
@@ -7099,6 +7100,8 @@
 
     "@typespec/ts-http-runtime/https-proxy-agent": ["https-proxy-agent@7.0.6", "", { "dependencies": { "agent-base": "^7.1.2", "debug": "4" } }, "sha512-vK9P5/iUfdl95AI+JVyUuIcVtd4ofvtrOr3HNtM2yxC9bnMbEdp3x01OhQNnjb8IJYi38VlTE3mBXwcfvywuSw=="],
 
+    "@useautumn/ai-sdk/typescript": ["typescript@5.9.3", "", { "bin": { "tsc": "bin/tsc", "tsserver": "bin/tsserver" } }, "sha512-jl1vZzPDinLr9eUt3J/t7V6FgNEw9QjvBPdysz9KfQDD41fQrC2Y4vKQdiaUpFT4bXlb1RHhLpp8wtm6M5TgSw=="],
+
     "@useautumn/sdk/zod": ["zod@4.4.3", "", {}, "sha512-ytENFjIJFl2UwYglde2jchW2Hwm4GJFLDiSXWdTrJQBIN9Fcyp7n4DhxJEiWNAJMV1/BqWfW/kkg71UDcHJyTQ=="],
 
     "@vercel/sdk/zod": ["zod@4.4.3", "", {}, "sha512-ytENFjIJFl2UwYglde2jchW2Hwm4GJFLDiSXWdTrJQBIN9Fcyp7n4DhxJEiWNAJMV1/BqWfW/kkg71UDcHJyTQ=="],

packages/ai-sdk/package.json

@@ -20,6 +20,7 @@
 	},
 	"scripts": {
 		"ts": "tsgo --noEmit --skipLibCheck",
+		"test": "bun test tests/unit",
 		"build": "rm -rf dist && tsup",
 		"prepublishOnly": "bun run build"
 	},

packages/ai-sdk/src/index.ts

@@ -4,6 +4,7 @@ import {
 	type LanguageModelUsage,
 	wrapLanguageModel,
 } from "ai";
+// @ts-expect-error autumn-js types resolve in consuming projects; this package only needs the peer type.
 import type { Autumn } from "autumn-js";
 
 // Standalone published package: must not import from the internal @autumn/shared workspace.
@@ -143,7 +144,7 @@ export const withAutumn = ({
 	const trackUsage = async (usage: AnyUsage) => {
 		try {
 			const pools = normalizeUsage(usage);
-			// @ts-expect-error trackTokens is generated from OpenAPI; local autumn-js types may not include it yet.
+			// @ts-ignore trackTokens is generated from OpenAPI; local autumn-js types may not include it yet.
 			await autumn.balances.trackTokens({
 				customerId,
 				modelId: modelName,

packages/ai-sdk/tests/unit/index.test.ts

@@ -0,0 +1,144 @@
+import { describe, expect, test } from "bun:test";
+import type { LanguageModelV3, LanguageModelV3Usage } from "@ai-sdk/provider";
+import { generateText, streamText } from "ai";
+import { withAutumn } from "../../src/index.js";
+
+type TrackTokensParams = {
+	customerId: string;
+	modelId: string;
+	inputTokens: number;
+	outputTokens: number;
+	cacheReadTokens?: number;
+	cacheWriteTokens?: number;
+	reasoningTokens?: number;
+	featureId?: string;
+	entityId?: string;
+	properties?: Record<string, unknown>;
+};
+
+const usage: LanguageModelV3Usage = {
+	inputTokens: {
+		total: 13,
+		noCache: 10,
+		cacheRead: 2,
+		cacheWrite: 1,
+	},
+	outputTokens: {
+		total: 7,
+		text: 5,
+		reasoning: 2,
+	},
+};
+
+const finishReason = { unified: "stop" as const, raw: "stop" };
+
+const createAutumn = () => {
+	const calls: TrackTokensParams[] = [];
+
+	return {
+		calls,
+		autumn: {
+			balances: {
+				trackTokens: async (params: TrackTokensParams) => {
+					calls.push(params);
+				},
+			},
+		},
+	};
+};
+
+const createModel = (): LanguageModelV3 => ({
+	specificationVersion: "v3",
+	provider: "openai",
+	modelId: "gpt-test",
+	supportedUrls: {},
+	async doGenerate() {
+		return {
+			content: [{ type: "text", text: "hello" }],
+			finishReason,
+			usage,
+			warnings: [],
+		};
+	},
+	async doStream() {
+		return {
+			stream: new ReadableStream({
+				start(controller) {
+					controller.enqueue({ type: "text-start", id: "text-1" });
+					controller.enqueue({
+						type: "text-delta",
+						id: "text-1",
+						delta: "hello",
+					});
+					controller.enqueue({ type: "text-end", id: "text-1" });
+					controller.enqueue({ type: "finish", finishReason, usage });
+					controller.close();
+				},
+			}),
+		};
+	},
+});
+
+describe("withAutumn", () => {
+	test("tracks token usage from generateText", async () => {
+		const { autumn, calls } = createAutumn();
+
+		const model = withAutumn({
+			autumn,
+			model: createModel(),
+			customerId: "cus_test",
+			featureId: "ai_credits",
+			entityId: "entity_test",
+			properties: { source: "test" },
+		});
+
+		const result = await generateText({ model, prompt: "Say hello" });
+
+		expect(result.text).toBe("hello");
+		expect(calls).toEqual([
+			{
+				customerId: "cus_test",
+				modelId: "openai/gpt-test",
+				inputTokens: 10,
+				outputTokens: 5,
+				cacheReadTokens: 2,
+				cacheWriteTokens: 1,
+				reasoningTokens: 2,
+				featureId: "ai_credits",
+				entityId: "entity_test",
+				properties: { source: "test" },
+			},
+		]);
+	});
+
+	test("tracks token usage from streamText when the stream finishes", async () => {
+		const { autumn, calls } = createAutumn();
+
+		const model = withAutumn({
+			autumn,
+			model: createModel(),
+			customerId: "cus_stream",
+			providerId: "custom-openai",
+		});
+
+		const result = streamText({ model, prompt: "Say hello" });
+		const chunks: string[] = [];
+
+		for await (const chunk of result.textStream) {
+			chunks.push(chunk);
+		}
+
+		expect(chunks.join("")).toBe("hello");
+		expect(calls).toEqual([
+			{
+				customerId: "cus_stream",
+				modelId: "custom-openai/gpt-test",
+				inputTokens: 10,
+				outputTokens: 5,
+				cacheReadTokens: 2,
+				cacheWriteTokens: 1,
+				reasoningTokens: 2,
+			},
+		]);
+	});
+});

packages/atmn/src/lib/transforms/sdkToApi/feature.ts

@@ -12,7 +12,7 @@ export interface ApiFeatureParams {
 		credit_cost: number;
 	}>;
 	model_markups?: Record<string, {
-		markup: number;
+		markup?: number;
 		input_cost?: number;
 		output_cost?: number;
 	}>;

packages/openapi/tsconfig.json

@@ -7,6 +7,7 @@
 		"moduleResolution": "bundler",
 		"target": "ES2020",
 		"noEmit": true,
+		"types": ["node", "bun"],
 		"paths": {
 			"@autumn/shared": ["../../shared/index.ts"],
 			"@api/*": ["../../shared/api/*"],
@@ -15,6 +16,5 @@
 		}
 	},
 	"include": ["./**/*"],
-	"types": ["node"],
 	"exclude": ["node_modules", "dist"]
 }

server/src/init.ts

@@ -45,7 +45,6 @@ import {
 import { preWarmOrgRedisConnections } from "./external/redis/orgRedisPool.js";
 import { createHonoApp } from "./initHono.js";
 import { otelSdk } from "./instrumentation.js";
-import { initializeDatabaseFunctions } from "./db/initializeDatabaseFunctions.js";
 import { checkEnvVars } from "./utils/initUtils.js";
 import { startMemoryMonitor } from "./utils/memoryMonitor.js";
 
@@ -67,7 +66,6 @@ const init = async ({ startupStartedAt }: { startupStartedAt: number }) => {
 	void preWarmOrgRedisConnections({ db }).catch((error) => {
 		logger.warn("[OrgRedis] Warmup failed", { error });
 	});
-	await initializeDatabaseFunctions();
 
 	await startAllEdgeConfigPolling({ logger });
 	await Promise.all([primeRedisMonitor(), primeRedisV2Monitor()]);

shared/utils/productUtils/priceUtils/convertPrice/priceToStripePrepaidV2Tiers.ts

@@ -1,11 +1,10 @@
 import type { Organization } from "@models/orgModels/orgTable";
 import type { Entitlement } from "@models/productModels/entModels/entModels";
-import type { UsagePriceConfig } from "@models/productModels/priceModels/priceConfig/usagePriceConfig";
 import type { Price } from "@models/productModels/priceModels/priceModels";
 import { orgToCurrency } from "@utils/orgUtils/convertOrgUtils";
 import {
-	isFinalTier,
 	isNotFinalTier,
+	isPrepaidPrice,
 } from "@utils/productUtils/priceUtils/classifyPriceUtils";
 import { atmnToStripeAmountDecimal } from "@utils/productUtils/priceUtils/convertAmountUtils";
 import { Decimal } from "decimal.js";
@@ -33,8 +32,13 @@ export const priceToStripePrepaidV2Tiers = ({
 	price: Price;
 	entitlement: Entitlement;
 	org: Organization;
-}) => {
-	const config = price.config as UsagePriceConfig;
+}): Stripe.PriceCreateParams.Tier[] => {
+	if (!isPrepaidPrice(price)) {
+		throw new Error(
+			`priceToStripePrepaidV2Tiers requires a prepaid price, got price ${price.id}`,
+		);
+	}
+	const config = price.config;
 
 	const tiers: Stripe.PriceCreateParams.Tier[] = [];
 
@@ -47,9 +51,8 @@ export const priceToStripePrepaidV2Tiers = ({
 		});
 	}
 
-	for (let i = 0; i < config.usage_tiers.length; i++) {
-		const tier = config.usage_tiers[i];
-		const atmnUnitAmount = new Decimal(tier.amount).div(
+	for (const tier of config.usage_tiers) {
+		const atmnUnitAmount = new Decimal(tier.amount ?? 0).div(
 			config.billing_units ?? 1,
 		);
 
@@ -58,14 +61,14 @@ export const priceToStripePrepaidV2Tiers = ({
 			currency: orgToCurrency({ org }),
 		});
 
-		let upTo = tier.to;
-		if (isNotFinalTier(tier) && entitlement.allowance) {
-			upTo = tier.to + entitlement.allowance;
+		let upTo: Stripe.PriceCreateParams.Tier["up_to"] = "inf";
+		if (isNotFinalTier(tier)) {
+			upTo = entitlement.allowance ? tier.to + entitlement.allowance : tier.to;
 		}
 
 		const stripeTier: Stripe.PriceCreateParams.Tier = {
 			unit_amount_decimal: stripeUnitAmountDecimal,
-			up_to: isFinalTier(tier) ? "inf" : upTo,
+			up_to: upTo,
 		};
 
 		if (tier.flat_amount) {
@@ -79,13 +82,13 @@ export const priceToStripePrepaidV2Tiers = ({
 	}
 
 	// Divide all tiers by billing units
-	const dividedTiers = tiers.map((tier, index: number) => ({
+	return tiers.map((tier, index) => ({
 		...tier,
 
 		up_to:
-			index === tiers.length - 1
+			index === tiers.length - 1 || tier.up_to === "inf"
 				? "inf"
-				: new Decimal(tier.up_to ?? 0)
+				: new Decimal(tier.up_to)
 						.div(config.billing_units ?? 1)
 						.ceil()
 						.toNumber(),
@@ -94,6 +97,4 @@ export const priceToStripePrepaidV2Tiers = ({
 			.mul(config.billing_units ?? 1)
 			.toString(),
 	}));
-
-	return dividedTiers;
 };

vite/src/views/products/features/credit-systems/components/AiCreditSchema.tsx

@@ -1,7 +1,12 @@
-import { PlusIcon } from "lucide-react";
+import { InfoIcon } from "lucide-react";
 import { FormLabel } from "@/components/v2/form/FormLabel";
 import { Input } from "@/components/v2/inputs/Input";
 import { SearchableSelect } from "@/components/v2/selects/SearchableSelect";
+import {
+	Tooltip,
+	TooltipContent,
+	TooltipTrigger,
+} from "@/components/v2/tooltips/Tooltip";
 import { useAiProviders } from "../hooks/useAiProviders";
 import type { CreditSystemFormInstance } from "../hooks/useCreditSystemForm";
 import { AiCreditSchemaTable } from "./AiCreditSchemaTable";
@@ -43,39 +48,51 @@ export function AiCreditSchema({ form }: AiCreditSchemaProps) {
 				/>
 			</div>
 
-			<div className="flex flex-col gap-3">
-				{activeProviderKeys.map((providerKey) => {
-					const provider = providers[providerKey];
-					const modelFullIds = providerGroups[providerKey] ?? [];
-					const providerName =
-						provider?.name ??
-						providerKey.charAt(0).toUpperCase() + providerKey.slice(1);
+			{activeProviderKeys.length > 0 && (
+				<div className="flex flex-col gap-3">
+					{activeProviderKeys.map((providerKey) => {
+						const provider = providers[providerKey];
+						const modelFullIds = providerGroups[providerKey] ?? [];
+						const providerName =
+							provider?.name ??
+							providerKey.charAt(0).toUpperCase() + providerKey.slice(1);
 
-					return (
-						<AiCreditSchemaTable
-							key={providerKey}
-							form={form}
-							providerKey={providerKey}
-							providerName={providerName}
-							modelFullIds={modelFullIds}
-							provider={
-								provider ?? { id: providerKey, name: providerKey, models: {} }
-							}
-							isLoading={isLoading}
-							removeKeys={removeKeys}
-							removeProvider={removeProvider}
-							setProviderMarkup={setProviderMarkup}
-							renameKey={renameKey}
-						/>
-					);
-				})}
-			</div>
+						return (
+							<AiCreditSchemaTable
+								key={providerKey}
+								form={form}
+								providerKey={providerKey}
+								providerName={providerName}
+								modelFullIds={modelFullIds}
+								provider={
+									provider ?? { id: providerKey, name: providerKey, models: {} }
+								}
+								isLoading={isLoading}
+								removeKeys={removeKeys}
+								removeProvider={removeProvider}
+								setProviderMarkup={setProviderMarkup}
+								renameKey={renameKey}
+							/>
+						);
+					})}
+				</div>
+			)}
 
 			<div
 				className="flex flex-col gap-1.5"
 				onWheel={(e) => e.stopPropagation()}
 			>
-				<FormLabel>Add Provider</FormLabel>
+				<FormLabel className="flex items-center gap-1.5">
+					Add Provider Override
+					<Tooltip>
+						<TooltipTrigger asChild>
+							<InfoIcon className="size-3.5 cursor-help text-tertiary-foreground" />
+						</TooltipTrigger>
+						<TooltipContent>
+							Add specific markup overrides for certain providers/models.
+						</TooltipContent>
+					</Tooltip>
+				</FormLabel>
 				<SearchableSelect
 					value={null}
 					onValueChange={addProvider}

vite/src/views/products/features/credit-systems/components/CreditSystemSchema.tsx

@@ -1,59 +1,14 @@
-import {
-	FeatureType,
-	isAiCreditSystem,
-	joinModelId,
-	type ModelsDevProvider,
-} from "@autumn/shared";
+import { FeatureType, isAiCreditSystem } from "@autumn/shared";
 import { useStore } from "@tanstack/react-form";
 import { useMemo } from "react";
 import { GroupedTabButton } from "@/components/v2/buttons/GroupedTabButton";
 import { SheetSection } from "@/components/v2/sheets/SharedSheetComponents";
-import { useModelsDevPricing } from "@/hooks/queries/useAiModelsQuery";
 import type { CreditSystemFormInstance } from "../hooks/useCreditSystemForm";
 import { AiCreditSchema } from "./AiCreditSchema";
 import { ClassicCreditSchema } from "./ClassicCreditSchema";
 
 type CreditSchemaMode = "classic" | "ai";
 
-const DEFAULT_AI_MODEL_COMPANIES = ["anthropic", "google", "openai"] as const;
-
-const getReleaseDateMs = (releaseDate?: string) => {
-	if (!releaseDate) return -1;
-	const timestamp = Date.parse(releaseDate);
-	return Number.isNaN(timestamp) ? -1 : timestamp;
-};
-
-function getDefaultModelMarkups(
-	providers: Record<string, ModelsDevProvider>,
-): Record<string, { markup?: number }> {
-	const result: Record<string, { markup?: number }> = {};
-	const preferredProvider =
-		providers["openrouter"] ?? Object.values(providers)[0];
-	if (!preferredProvider) return result;
-
-	const providerKey = preferredProvider.id;
-	for (const company of DEFAULT_AI_MODEL_COMPANIES) {
-		const companyModels = Object.entries(preferredProvider.models).filter(
-			([key]) => key.startsWith(company),
-		);
-
-		const latestModel = companyModels.reduce<
-			[string, ModelsDevProvider["models"][string]] | null
-		>((currentLatest, candidate) => {
-			if (!currentLatest) return candidate;
-			const currentRelease = getReleaseDateMs(currentLatest[1].release_date);
-			const candidateRelease = getReleaseDateMs(candidate[1].release_date);
-			return candidateRelease > currentRelease ? candidate : currentLatest;
-		}, null);
-
-		if (!latestModel) continue;
-
-		const [modelKey] = latestModel;
-		result[joinModelId(providerKey, modelKey)] = {};
-	}
-	return result;
-}
-
 interface CreditSystemSchemaProps {
 	form: CreditSystemFormInstance;
 	disableModeSwitch?: boolean;
@@ -63,20 +18,16 @@ export function CreditSystemSchema({
 	form,
 	disableModeSwitch = false,
 }: CreditSystemSchemaProps) {
-	const { providers } = useModelsDevPricing();
 	const type = useStore(form.store, (s) => s.values.type);
 
 	const mode: CreditSchemaMode = isAiCreditSystem(type) ? "ai" : "classic";
 
 	const handleModeChange = (newMode: string) => {
 		if (newMode === "ai") {
-			const modelMarkups = getDefaultModelMarkups(providers);
 			form.setFieldValue("type", FeatureType.AiCreditSystem);
 			form.setFieldValue("config", { ...form.state.values.config, schema: [] });
-			form.setFieldValue(
-				"model_markups",
-				Object.keys(modelMarkups).length > 0 ? modelMarkups : {},
-			);
+			form.setFieldValue("model_markups", {});
+			form.setFieldValue("provider_markups", {});
 		} else {
 			form.setFieldValue("type", FeatureType.CreditSystem);
 			form.setFieldValue("config", {
@@ -86,6 +37,7 @@ export function CreditSystemSchema({
 				],
 			});
 			form.setFieldValue("model_markups", {});
+			form.setFieldValue("provider_markups", {});
 		}
 	};
 

vite/src/views/products/plan/components/new-feature/NewFeatureType.tsx

@@ -2,6 +2,7 @@ import {
 	FeatureType as APIFeatureType,
 	type CreateFeature,
 	FeatureUsageType,
+	isAnyCreditSystem,
 } from "@autumn/shared";
 import { BarcodeIcon, CoinsIcon } from "@phosphor-icons/react";
 import { PanelButton } from "@/components/v2/buttons/PanelButton";
@@ -57,7 +58,7 @@ export function NewFeatureType({
 
 					<div className="flex w-full items-center gap-4">
 						<PanelButton
-							isSelected={feature.type === APIFeatureType.CreditSystem}
+							isSelected={isAnyCreditSystem(feature.type)}
 							onClick={() => {
 								setFeature({
 									...feature,