From 8822664df95534844eef0456633c0ed43de2437d Mon Sep 17 00:00:00 2001 From: Charlie Lamb Date: Tue, 19 May 2026 19:01:04 +0100 Subject: [PATCH 01/46] feat: AI credit system with model-based token pricing Adds a new FeatureType.AiCreditSystem that enables model-based token pricing with per-model markup configuration. Includes the @useautumn/ai-sdk package for Vercel AI SDK integration with automatic token tracking. Co-authored-by: TheUntraceable <73362400+TheUntraceable@users.noreply.github.com> --- ai | 2 +- .../balances/trackTokens.mdx | 61 +++ .../api-reference/balances/trackTokens.mdx | 171 +++++++++ apps/docs/mintlify/docs.json | 2 + .../customers/tracking-usage.mdx | 68 ++++ .../external-providers/ai-sdk.mdx | 141 +++++++ .../modelling-pricing/credit-systems.mdx | 111 ++++++ knip.json | 1 + package.json | 5 +- packages/ai-sdk/bun.lock | 253 +++++++++++++ packages/ai-sdk/package.json | 35 ++ packages/ai-sdk/src/index.ts | 104 ++++++ packages/ai-sdk/tsconfig.json | 19 + packages/ai-sdk/tsup.config.ts | 12 + .../atmn/src/compose/models/featureModels.ts | 20 +- .../transforms/apiToSdk/Transformer.test.ts | 65 +++- .../src/lib/transforms/apiToSdk/feature.ts | 50 ++- .../src/lib/transforms/sdkToApi/feature.ts | 18 + .../src/lib/transforms/sdkToCode/feature.ts | 5 + .../src/internal/balances/balancesRouter.ts | 3 + .../balances/handlers/handleTrackTokens.ts | 129 +++++++ .../deduction/executePostgresDeduction.ts | 2 +- .../utils/deduction/executeRedisDeduction.ts | 2 +- .../deduction/prepareFeatureDeduction.ts | 36 +- .../deductionV2/executePostgresDeductionV2.ts | 2 +- .../deductionV2/executeRedisDeductionV2.ts | 2 +- .../deductionV2/prepareFeatureDeductionV2.ts | 67 ++-- .../balances/utils/types/featureDeduction.ts | 11 +- .../internal/features/creditSystemUtils.ts | 126 ++++++- .../features/featureActions/createFeature.ts | 17 +- .../features/featureActions/updateFeature.ts | 69 +++- server/src/internal/features/featureRouter.ts | 2 + server/src/internal/features/featureUtils.ts | 13 +- .../handleUpdateFeatureV1.ts | 2 + .../internalHandlers/handleGetModelPricing.ts | 11 + .../features/utils/constructFeatureUtils.ts | 43 ++- .../features/utils/getModelPricing.ts | 40 ++ server/tests/advanced/usage/usage2.test.ts | 2 +- server/tests/advanced/usage/usage3.test.ts | 2 +- server/tests/advanced/usage/usage4.test.ts | 2 +- .../credit-systems/credit-systems1.test.ts | 6 +- .../check/send-event/send-event3.test.ts | 6 +- .../check/send-event/send-event4.test.ts | 2 +- .../auto-topup-credit-systems.test.ts | 6 +- .../check-entity-product-spend-limit.test.ts | 2 +- .../check-per-entity-spend-limit.test.ts | 2 +- .../check-with-lock-credit-system.test.ts | 12 +- .../track/basic/track-credit-system.test.ts | 20 +- .../track/basic/track-deductions.test.ts | 2 +- .../balances/track/basic/track-tokens.test.ts | 349 ++++++++++++++++++ .../track-overage-allowed-consumable.test.ts | 2 +- .../track-customer-spend-limit.test.ts | 2 +- .../track-entity-product-spend-limit.test.ts | 2 +- .../track-per-entity-spend-limit.test.ts | 2 +- .../track-postgres-entity-spend-limit.test.ts | 4 +- server/tests/setup/v2Features.ts | 37 ++ .../tests/utils/fixtures/db/entitlements.ts | 6 + server/tests/utils/fixtures/db/features.ts | 6 + .../api/balances/track/trackTokensParams.ts | 45 +++ shared/api/features/apiFeatureV1.ts | 7 +- .../features/changes/V1.2_FeatureChange.ts | 5 +- .../crud/common/baseFeatureParamsV1.ts | 8 +- shared/api/models.ts | 1 + shared/index.ts | 2 + shared/models/aiModels/modelsDevTypes.ts | 17 + .../featureConfig/creditConfig.ts | 13 +- shared/models/featureModels/featureEnums.ts | 1 + shared/models/featureModels/featureModels.ts | 3 + shared/models/featureModels/featureTable.ts | 3 +- shared/utils/agentTypes.ts | 17 +- shared/utils/featureUtils.ts | 2 +- .../featureUtils/apiFeatureToDbFeature.ts | 7 + .../utils/featureUtils/creditSystemUtils.ts | 3 +- shared/utils/productDisplayUtils.ts | 39 +- vite/src/hooks/queries/useAiModelsQuery.ts | 24 ++ .../product/product-item/formatProductItem.ts | 18 +- .../components/CreateFeatureSheet.tsx | 22 +- .../components/UpdateFeatureSheet.tsx | 16 +- .../components/AiCreditSchema.tsx | 191 ++++++++++ .../components/AiCreditSchemaRow.tsx | 172 +++++++++ .../components/AiModelSelectDropdown.tsx | 53 +++ .../components/ClassicCreditSchema.tsx | 161 ++++++++ .../components/CreditSystemSchema.tsx | 227 ++++++------ .../components/UpdateCreditSystemSheet.tsx | 20 +- .../credit-systems/hooks/useAiCreditSchema.ts | 259 +++++++++++++ .../utils/validateCreditSystem.ts | 32 +- .../feature-list/CreditListColumns.tsx | 34 +- .../feature-list/FeatureListTable.tsx | 7 +- .../edit-plan-feature/BillingType.tsx | 10 +- .../EditPlanFeatureSheet.tsx | 3 +- .../edit-plan-feature/IncludedUsage.tsx | 104 ++++-- .../new-feature/NewFeatureBehaviour.tsx | 5 +- .../plan-card/DummyPlanFeatureRow.tsx | 6 + .../components/plan-card/PlanFeatureRow.tsx | 8 +- 94 files changed, 3375 insertions(+), 364 deletions(-) create mode 100644 apps/docs/api-reference-generator/balances/trackTokens.mdx create mode 100644 apps/docs/mintlify/api-reference/balances/trackTokens.mdx create mode 100644 apps/docs/mintlify/documentation/external-providers/ai-sdk.mdx create mode 100644 packages/ai-sdk/bun.lock create mode 100644 packages/ai-sdk/package.json create mode 100644 packages/ai-sdk/src/index.ts create mode 100644 packages/ai-sdk/tsconfig.json create mode 100644 packages/ai-sdk/tsup.config.ts create mode 100644 server/src/internal/balances/handlers/handleTrackTokens.ts create mode 100644 server/src/internal/features/internalHandlers/handleGetModelPricing.ts create mode 100644 server/src/internal/features/utils/getModelPricing.ts create mode 100644 server/tests/integration/balances/track/basic/track-tokens.test.ts create mode 100644 shared/api/balances/track/trackTokensParams.ts create mode 100644 shared/models/aiModels/modelsDevTypes.ts create mode 100644 vite/src/hooks/queries/useAiModelsQuery.ts create mode 100644 vite/src/views/products/features/credit-systems/components/AiCreditSchema.tsx create mode 100644 vite/src/views/products/features/credit-systems/components/AiCreditSchemaRow.tsx create mode 100644 vite/src/views/products/features/credit-systems/components/AiModelSelectDropdown.tsx create mode 100644 vite/src/views/products/features/credit-systems/components/ClassicCreditSchema.tsx create mode 100644 vite/src/views/products/features/credit-systems/hooks/useAiCreditSchema.ts diff --git a/ai b/ai index b1efb8d30..761b84255 160000 --- a/ai +++ b/ai @@ -1 +1 @@ -Subproject commit b1efb8d303caf7fdc76d820b00ad87d305f6867f +Subproject commit 761b842553d890355b77e618c3aa703dddb661c1 diff --git a/apps/docs/api-reference-generator/balances/trackTokens.mdx b/apps/docs/api-reference-generator/balances/trackTokens.mdx new file mode 100644 index 000000000..9fed45b3e --- /dev/null +++ b/apps/docs/api-reference-generator/balances/trackTokens.mdx @@ -0,0 +1,61 @@ +--- +title: "Track Token Usage" +openapi: "openapi POST /v1/balances.trackTokens" +--- + +import { DynamicParamField } from "/components/dynamic-param-field.jsx"; +import { DynamicResponseExample } from "/components/dynamic-response-example.jsx"; +import { DynamicResponseField } from "/components/dynamic-response-field.jsx"; + + + Track AI token usage against a customer's AI credit system balance. Converts token counts to a dollar cost using [Models.dev](https://models.dev) pricing and your configured markup, then deducts from the customer's credit balance. + + +The `model_id` must use `provider/model` format, matching the provider and model keys from [Models.dev](https://models.dev). For providers with nested model paths (like OpenRouter), include the full path: `openrouter/anthropic/claude-opus-4.6`. + +### Common Use Cases + + + +```typescript Anthropic +await autumn.balances.trackTokens({ + customerId: "cus_123", + modelId: "anthropic/claude-opus-4-6", + inputTokens: 1000, + outputTokens: 500 +}); +``` + +```typescript OpenAI +await autumn.balances.trackTokens({ + customerId: "cus_123", + modelId: "openai/gpt-4o", + inputTokens: 500, + outputTokens: 200, + properties: { + conversation_id: "conv_abc123", + prompt_type: "summarization" + } +}); +``` + +```typescript OpenRouter (nested path) +await autumn.balances.trackTokens({ + customerId: "cus_123", + modelId: "openrouter/anthropic/claude-opus-4-6", + inputTokens: 2000, + outputTokens: 1000 +}); +``` + +```typescript With explicit feature +await autumn.balances.trackTokens({ + customerId: "cus_123", + featureId: "ai_credits", + modelId: "anthropic/claude-haiku-4-5", + inputTokens: 2000, + outputTokens: 1000 +}); +``` + + diff --git a/apps/docs/mintlify/api-reference/balances/trackTokens.mdx b/apps/docs/mintlify/api-reference/balances/trackTokens.mdx new file mode 100644 index 000000000..28f811123 --- /dev/null +++ b/apps/docs/mintlify/api-reference/balances/trackTokens.mdx @@ -0,0 +1,171 @@ +--- +title: "Track Token Usage" +openapi: "openapi POST /v1/balances.trackTokens" +--- + +import { DynamicParamField } from "/components/dynamic-param-field.jsx"; +import { DynamicResponseExample } from "/components/dynamic-response-example.jsx"; +import { DynamicResponseField } from "/components/dynamic-response-field.jsx"; + + + Track AI token usage against a customer's AI credit system balance. Converts token counts to a dollar cost using [Models.dev](https://models.dev) pricing and your configured markup, then deducts from the customer's credit balance. + + +The `model_id` must use `provider/model` format, matching the provider and model keys from [Models.dev](https://models.dev). For providers with nested model paths (like OpenRouter), include the full path: `openrouter/anthropic/claude-opus-4.6`. + +### Common Use Cases + + + +```typescript Anthropic +await autumn.balances.trackTokens({ + customerId: "cus_123", + modelId: "anthropic/claude-opus-4-6", + inputTokens: 1000, + outputTokens: 500 +}); +``` + +```typescript OpenAI +await autumn.balances.trackTokens({ + customerId: "cus_123", + modelId: "openai/gpt-4o", + inputTokens: 500, + outputTokens: 200, + properties: { + conversation_id: "conv_abc123", + prompt_type: "summarization" + } +}); +``` + +```typescript OpenRouter (nested path) +await autumn.balances.trackTokens({ + customerId: "cus_123", + modelId: "openrouter/anthropic/claude-opus-4.6", + inputTokens: 2000, + outputTokens: 1000 +}); +``` + +```typescript With explicit feature +await autumn.balances.trackTokens({ + customerId: "cus_123", + featureId: "ai_credits", + modelId: "anthropic/claude-haiku-4-5", + inputTokens: 2000, + outputTokens: 1000 +}); +``` + + + +### Body Parameters + + + 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`). + + + + Number of input tokens consumed. + + + + Number of output tokens consumed. + + + + The ID of the AI credit system feature. If omitted, automatically detects the organization's AI credit system feature. + + + + The ID of the entity for entity-scoped balances. + + + + Additional properties to attach to this usage event. The `model`, `input_tokens`, and `output_tokens` values are automatically included. + + +### Response + + + The ID of the customer whose token usage was tracked. + + + + The dollar cost that was deducted from the customer's AI credit balance. + + + + The updated balance for the AI credit system feature. + + + The feature ID this balance is for. + + + + Total balance granted (included + prepaid). + + + + Remaining balance available for use. + + + + Total usage consumed in the current period. + + + + Whether this feature has unlimited usage. + + + + Whether usage beyond the granted balance is allowed. + + + + Timestamp when the balance will reset, or null for no reset. + + + + + + + +```json 200 +{ + "customer_id": "cus_123", + "value": 0.06, + "balance": { + "feature_id": "ai_credits", + "granted": 10.00, + "remaining": 9.94, + "usage": 0.06, + "unlimited": false, + "overage_allowed": false, + "next_reset_at": 1773851121437, + "breakdown": [ + { + "id": "cus_ent_abc123", + "plan_id": "pro_plan", + "included_grant": 10.00, + "prepaid_grant": 0, + "remaining": 9.94, + "usage": 0.06, + "unlimited": false, + "reset": { + "interval": "month", + "resets_at": 1773851121437 + }, + "price": null, + "expires_at": null + } + ] + } +} +``` + diff --git a/apps/docs/mintlify/docs.json b/apps/docs/mintlify/docs.json index c4d4dfc6a..9985a5eea 100644 --- a/apps/docs/mintlify/docs.json +++ b/apps/docs/mintlify/docs.json @@ -133,6 +133,7 @@ "documentation/slack-discord-notifications", "documentation/fail-open", "documentation/rate-limits", + "documentation/external-providers/ai-sdk", "documentation/external-providers/convex", "documentation/external-providers/revenuecat", "documentation/external-providers/vercel-marketplace" @@ -205,6 +206,7 @@ "pages": [ "api-reference/core/check", "api-reference/core/track", + "api-reference/balances/trackTokens", "api-reference/balances/createBalance", "api-reference/balances/updateBalance", "api-reference/balances/deleteBalance", diff --git a/apps/docs/mintlify/documentation/customers/tracking-usage.mdx b/apps/docs/mintlify/documentation/customers/tracking-usage.mdx index e0809faac..90a65b734 100644 --- a/apps/docs/mintlify/documentation/customers/tracking-usage.mdx +++ b/apps/docs/mintlify/documentation/customers/tracking-usage.mdx @@ -131,6 +131,74 @@ curl -X POST "https://api.useautumn.com/v1/balances/update" \ can reset or override incremental usage recorded through events. +## Tracking AI Token Usage + +If you're using an [AI credit system](/examples/monetary-credits), you can track token usage directly with `trackTokens`. This automatically converts token counts to a dollar cost using [Models.dev](https://models.dev) pricing and your configured markup, then deducts from the customer's credit balance. + +The `modelId` must be in `provider/model` format, matching the provider and model keys from [Models.dev](https://models.dev). For example: +- `anthropic/claude-sonnet-4-5-20250514` +- `openai/gpt-4o` +- `google/gemini-2.5-pro` + +For providers with nested model paths (like OpenRouter), include the full path after the provider: `openrouter/anthropic/claude-opus-4.6`. + + + +```typescript TypeScript +import { Autumn } from "autumn-js"; + +const autumn = new Autumn({ secretKey: "am_sk_test_1234" }); + +await autumn.balances.trackTokens({ + customerId: "user_123", + modelId: "anthropic/claude-opus-4-6", + inputTokens: 1000, + outputTokens: 500, +}); +``` + +```python Python +from autumn_sdk import Autumn + +autumn = Autumn("am_sk_test_1234") + +await autumn.balances.track_tokens( + customer_id="user_123", + model_id="anthropic/claude-opus-4-6", + input_tokens=1000, + output_tokens=500, +) +``` + +```bash cURL +curl -X POST "https://api.useautumn.com/v1/balances.trackTokens" \ + -H "Authorization: Bearer am_sk_test_1234" \ + -H "Content-Type: application/json" \ + -d '{ + "customer_id": "user_123", + "model_id": "anthropic/claude-opus-4-6", + "input_tokens": 1000, + "output_tokens": 500 + }' +``` + + + + + If your organization has only one AI credit system feature, you can omit the `featureId` parameter — it will be auto-detected. + + +### Vercel AI SDK integration + +If you're using the [Vercel AI SDK](https://sdk.vercel.ai), the `@useautumn/ai-sdk` package can automatically track token usage for every `generateText` or `streamText` call — no manual `trackTokens` calls needed. + + + ## Using Event Names In the above examples, we used the `featureId` to identify the feature. You can instead use the `eventName` parameter to link balances to different events in your application. This can be useful when: diff --git a/apps/docs/mintlify/documentation/external-providers/ai-sdk.mdx b/apps/docs/mintlify/documentation/external-providers/ai-sdk.mdx new file mode 100644 index 000000000..bb33e85f0 --- /dev/null +++ b/apps/docs/mintlify/documentation/external-providers/ai-sdk.mdx @@ -0,0 +1,141 @@ +--- +title: "Vercel AI SDK" +description: "Automatically track AI token usage with the Vercel AI SDK" +--- + +The `@useautumn/ai-sdk` package integrates Autumn with the [Vercel AI SDK](https://sdk.vercel.ai), automatically tracking token usage for every `generateText` or `streamText` call. No manual `trackTokens` calls needed. + +## Setup + +#### 1. Install the package + + +```bash npm +npm install @useautumn/ai-sdk +``` + +```bash pnpm +pnpm add @useautumn/ai-sdk +``` + +```bash yarn +yarn add @useautumn/ai-sdk +``` + +```bash bun +bun add @useautumn/ai-sdk +``` + + + + Requires `autumn-js` and `ai` (v6+) as peer dependencies. + + +#### 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. + +```typescript +import { Autumn } from "autumn-js"; +import { anthropic } from "@ai-sdk/anthropic"; +import { withTokenTracking } from "@useautumn/ai-sdk"; + +const autumn = new Autumn({ secretKey: "am_sk_test_1234" }); + +const model = withTokenTracking({ + autumn, + model: anthropic("claude-sonnet-4-5-20250514"), + customerId: "user_123", +}); +``` + +#### 3. Use as normal + +The wrapped model works exactly like a regular AI SDK model. Token usage is tracked in the background after each call. + +```typescript +import { generateText, streamText } from "ai"; + +// Generate — usage tracked automatically +const { text } = await generateText({ + model, + prompt: "Explain quantum computing in one paragraph", +}); + +// Stream — usage tracked when the stream finishes +const result = streamText({ + model, + prompt: "Write a short poem about recursion", +}); + +for await (const chunk of result.textStream) { + process.stdout.write(chunk); +} +``` + +## Model ID format + +The wrapped model constructs the `modelId` sent to Autumn using `provider/model` format, derived from the AI SDK model's `provider` and `modelId` fields. This must match a valid provider and model key from [Models.dev](https://models.dev). + +For example: +- `@ai-sdk/anthropic` → `anthropic/claude-sonnet-4-5-20250514` +- `@ai-sdk/openai` → `openai/gpt-4o` +- `@ai-sdk/google` → `google/gemini-2.5-pro` + +If the AI SDK provider name doesn't match the Models.dev provider key, use the `providerId` option to override it: + +```typescript +import { createOpenRouter } from "@openrouter/ai-sdk-provider"; + +const openrouter = createOpenRouter(); + +const model = withTokenTracking({ + autumn, + model: openrouter("anthropic/claude-opus-4-6"), + customerId: "user_123", + providerId: "openrouter", // Override provider prefix +}); +// Sends modelId as "openrouter/anthropic/claude-opus-4-6" +``` + +## Options + +| Parameter | Type | Required | Description | +|-----------|------|----------|-------------| +| `autumn` | `Autumn` | Yes | Your Autumn SDK client instance | +| `model` | `LanguageModelV3` | Yes | The AI SDK language model to wrap | +| `customerId` | `string` | Yes | The Autumn customer ID to attribute usage to | +| `providerId` | `string` | No | Override the provider prefix in the model name sent to Autumn. Falls back to the model's `provider` field | +| `featureId` | `string` | No | Target a specific AI credit system feature. Auto-detected if you only have one | +| `entityId` | `string` | No | Entity ID for entity-scoped balance tracking | +| `properties` | `Record` | No | Additional properties to attach to each usage event | + +## Full example + +```typescript +import { Autumn } from "autumn-js"; +import { openai } from "@ai-sdk/openai"; +import { generateText } from "ai"; +import { withTokenTracking } from "@useautumn/ai-sdk"; + +const autumn = new Autumn({ secretKey: process.env.AUTUMN_SECRET_KEY! }); + +async function chat(customerId: string, message: string) { + const model = withTokenTracking({ + autumn, + model: openai("gpt-4o"), + customerId, + }); + + const { text } = await generateText({ + model, + prompt: message, + }); + + return text; +} +``` + + + Tracking failures are caught and logged to the console — they won't break your AI features. Check your server logs if usage isn't appearing in Autumn. + diff --git a/apps/docs/mintlify/documentation/modelling-pricing/credit-systems.mdx b/apps/docs/mintlify/documentation/modelling-pricing/credit-systems.mdx index e165aede4..b4382ec75 100644 --- a/apps/docs/mintlify/documentation/modelling-pricing/credit-systems.mdx +++ b/apps/docs/mintlify/documentation/modelling-pricing/credit-systems.mdx @@ -271,3 +271,114 @@ See the credits pricing guide for a more detailed example of setting up a moneta href="/examples/monetary-credits" icon="money-bills" /> + +## AI Credit Systems + +For AI applications that need to track token usage with per-model pricing, you can create an AI credit system. This lets you define markup percentages for each model and automatically calculate costs based on input/output tokens. + + + + +Define an AI credit system with `modelMarkups` that maps model IDs to pricing configuration: + +```ts autumn.config.ts +import { feature, item, plan } from 'atmn'; + +export const aiCredits = feature({ + id: 'ai_credits', + name: 'AI Credits', + type: 'ai_credit_system', + modelMarkups: { + 'anthropic/claude-opus-4-5': { markup: 20 }, + 'anthropic/claude-sonnet-4-5': { markup: 15 }, + 'openai/gpt-4o': { markup: 15 }, + // For custom/self-hosted models, specify input/output costs in $/M tokens + 'custom/my-model': { markup: 25, inputCost: 0.01, outputCost: 0.03 }, + }, +}); + +export const pro = plan({ + id: 'pro', + name: 'Pro', + price: { amount: 50, interval: 'month' }, + items: [ + item({ + featureId: aiCredits.id, + included: 10, // $10 worth of AI credits + reset: { interval: 'month' }, + }), + ], +}); +``` + +Push changes with `atmn push`. + + + + +1. Navigate to the features page, under Plans. +2. Click "Create Credit System" +3. Toggle "AI Credit System" to enable model-based pricing +4. Add the models you want to support with their markup percentages +5. For custom models, also specify input/output costs per million tokens +6. Click "Create" + + + + +### Model ID Format + +Model IDs follow the `provider/model` format: +- Standard models: `anthropic/claude-opus-4-5`, `openai/gpt-4o` +- OpenRouter models: `openrouter/anthropic/claude-opus-4.6` +- Custom models: `custom/my-model-name` + +For standard models, pricing is automatically fetched from models.dev. For custom models, you must specify `inputCost` and `outputCost` in dollars per million tokens. + +### Tracking Token Usage + +Use the `trackTokens` endpoint to deduct credits based on token usage: + + + +```typescript TypeScript +import { Autumn } from "autumn-js"; + +const autumn = new Autumn({ secretKey: "am_sk_test_1234" }); + +await autumn.balances.trackTokens({ + customerId: "user_123", + modelId: "anthropic/claude-opus-4-5", + inputTokens: 1500, + outputTokens: 500, +}); +``` + +```python Python +from autumn_sdk import Autumn + +autumn = Autumn("am_sk_test_1234") + +await autumn.balances.track_tokens( + customer_id="user_123", + model_id="anthropic/claude-opus-4-5", + input_tokens=1500, + output_tokens=500, +) +``` + +```bash cURL +curl -X POST "https://api.useautumn.com/v1/balances.trackTokens" \ + -H "Authorization: Bearer am_sk_test_1234" \ + -H "Content-Type: application/json" \ + -d '{ + "customer_id": "user_123", + "model_id": "anthropic/claude-opus-4-5", + "input_tokens": 1500, + "output_tokens": 500 + }' +``` + + + +The cost is calculated automatically based on the model's pricing plus your configured markup percentage. diff --git a/knip.json b/knip.json index 74649a7a2..1775aaa17 100644 --- a/knip.json +++ b/knip.json @@ -13,6 +13,7 @@ ], "ignore": ["ai/**", "others/**", ".trigger/**"], "ignoreWorkspaces": [ + "packages/ai-sdk", "packages/atmn", "packages/autumn-js", "packages/openapi", diff --git a/package.json b/package.json index d2abaf7cd..13c56e25f 100644 --- a/package.json +++ b/package.json @@ -18,7 +18,8 @@ "packages/autumn-js", "packages/openapi", "packages/ksuid", - "packages/stripe-sync" + "packages/stripe-sync", + "packages/ai-sdk" ], "catalog": { "stripe": "19.3.0-beta.1", @@ -131,7 +132,7 @@ "site": "cd apps/website && bun dev && cd ../..", "docs": "bun -F @autumn/docs dev", "docs:build": "bun -F @autumn/docs build", - "ts": "turbo run ts --filter=@autumn/server --filter=autumn-js --filter=@autumn/openapi --filter=atmn --filter=checkout", + "ts": "turbo run ts --filter=@autumn/server --filter=autumn-js --filter=@autumn/openapi --filter=atmn --filter=checkout --filter=@useautumn/ai-sdk", "kill:ts": "while pgrep -f tsgo > /dev/null; do pkill -9 -f tsgo; sleep 0.1; done", "atmn:build": "bun -F atmn build", "openapi:ts": "bun -F @autumn/openapi ts", diff --git a/packages/ai-sdk/bun.lock b/packages/ai-sdk/bun.lock new file mode 100644 index 000000000..34f668488 --- /dev/null +++ b/packages/ai-sdk/bun.lock @@ -0,0 +1,253 @@ +{ + "lockfileVersion": 1, + "configVersion": 1, + "workspaces": { + "": { + "name": "@useautumn/ai-sdk", + "devDependencies": { + "@types/node": "^24.9.1", + "tsup": "^8.4.0", + "typescript": "^5.8.3", + }, + "peerDependencies": { + "ai": "^6.0.116", + "autumn-js": "*", + }, + }, + }, + "packages": { + "@ai-sdk/gateway": ["@ai-sdk/gateway@3.0.66", "", { "dependencies": { "@ai-sdk/provider": "3.0.8", "@ai-sdk/provider-utils": "4.0.19", "@vercel/oidc": "3.1.0" }, "peerDependencies": { "zod": "^3.25.76 || ^4.1.8" } }, "sha512-SIQ0YY0iMuv+07HLsZ+bB990zUJ6S4ujORAh+Jv1V2KGNn73qQKnGO0JBk+w+Res8YqOFSycwDoWcFlQrVxS4A=="], + + "@ai-sdk/provider": ["@ai-sdk/provider@3.0.8", "", { "dependencies": { "json-schema": "^0.4.0" } }, "sha512-oGMAgGoQdBXbZqNG0Ze56CHjDZ1IDYOwGYxYjO5KLSlz5HiNQ9udIXsPZ61VWaHGZ5XW/jyjmr6t2xz2jGVwbQ=="], + + "@ai-sdk/provider-utils": ["@ai-sdk/provider-utils@4.0.19", "", { "dependencies": { "@ai-sdk/provider": "3.0.8", "@standard-schema/spec": "^1.1.0", "eventsource-parser": "^3.0.6" }, "peerDependencies": { "zod": "^3.25.76 || ^4.1.8" } }, "sha512-3eG55CrSWCu2SXlqq2QCsFjo3+E7+Gmg7i/oRVoSZzIodTuDSfLb3MRje67xE9RFea73Zao7Lm4mADIfUETKGg=="], + + "@esbuild/aix-ppc64": ["@esbuild/aix-ppc64@0.27.4", "", { "os": "aix", "cpu": "ppc64" }, "sha512-cQPwL2mp2nSmHHJlCyoXgHGhbEPMrEEU5xhkcy3Hs/O7nGZqEpZ2sUtLaL9MORLtDfRvVl2/3PAuEkYZH0Ty8Q=="], + + "@esbuild/android-arm": ["@esbuild/android-arm@0.27.4", "", { "os": "android", "cpu": "arm" }, "sha512-X9bUgvxiC8CHAGKYufLIHGXPJWnr0OCdR0anD2e21vdvgCI8lIfqFbnoeOz7lBjdrAGUhqLZLcQo6MLhTO2DKQ=="], + + "@esbuild/android-arm64": ["@esbuild/android-arm64@0.27.4", "", { "os": "android", "cpu": "arm64" }, "sha512-gdLscB7v75wRfu7QSm/zg6Rx29VLdy9eTr2t44sfTW7CxwAtQghZ4ZnqHk3/ogz7xao0QAgrkradbBzcqFPasw=="], + + "@esbuild/android-x64": ["@esbuild/android-x64@0.27.4", "", { "os": "android", "cpu": "x64" }, "sha512-PzPFnBNVF292sfpfhiyiXCGSn9HZg5BcAz+ivBuSsl6Rk4ga1oEXAamhOXRFyMcjwr2DVtm40G65N3GLeH1Lvw=="], + + "@esbuild/darwin-arm64": ["@esbuild/darwin-arm64@0.27.4", "", { "os": "darwin", "cpu": "arm64" }, "sha512-b7xaGIwdJlht8ZFCvMkpDN6uiSmnxxK56N2GDTMYPr2/gzvfdQN8rTfBsvVKmIVY/X7EM+/hJKEIbbHs9oA4tQ=="], + + "@esbuild/darwin-x64": ["@esbuild/darwin-x64@0.27.4", "", { "os": "darwin", "cpu": "x64" }, "sha512-sR+OiKLwd15nmCdqpXMnuJ9W2kpy0KigzqScqHI3Hqwr7IXxBp3Yva+yJwoqh7rE8V77tdoheRYataNKL4QrPw=="], + + "@esbuild/freebsd-arm64": ["@esbuild/freebsd-arm64@0.27.4", "", { "os": "freebsd", "cpu": "arm64" }, "sha512-jnfpKe+p79tCnm4GVav68A7tUFeKQwQyLgESwEAUzyxk/TJr4QdGog9sqWNcUbr/bZt/O/HXouspuQDd9JxFSw=="], + + "@esbuild/freebsd-x64": ["@esbuild/freebsd-x64@0.27.4", "", { "os": "freebsd", "cpu": "x64" }, "sha512-2kb4ceA/CpfUrIcTUl1wrP/9ad9Atrp5J94Lq69w7UwOMolPIGrfLSvAKJp0RTvkPPyn6CIWrNy13kyLikZRZQ=="], + + "@esbuild/linux-arm": ["@esbuild/linux-arm@0.27.4", "", { "os": "linux", "cpu": "arm" }, "sha512-aBYgcIxX/wd5n2ys0yESGeYMGF+pv6g0DhZr3G1ZG4jMfruU9Tl1i2Z+Wnj9/KjGz1lTLCcorqE2viePZqj4Eg=="], + + "@esbuild/linux-arm64": ["@esbuild/linux-arm64@0.27.4", "", { "os": "linux", "cpu": "arm64" }, "sha512-7nQOttdzVGth1iz57kxg9uCz57dxQLHWxopL6mYuYthohPKEK0vU0C3O21CcBK6KDlkYVcnDXY099HcCDXd9dA=="], + + "@esbuild/linux-ia32": ["@esbuild/linux-ia32@0.27.4", "", { "os": "linux", "cpu": "ia32" }, "sha512-oPtixtAIzgvzYcKBQM/qZ3R+9TEUd1aNJQu0HhGyqtx6oS7qTpvjheIWBbes4+qu1bNlo2V4cbkISr8q6gRBFA=="], + + "@esbuild/linux-loong64": ["@esbuild/linux-loong64@0.27.4", "", { "os": "linux", "cpu": "none" }, "sha512-8mL/vh8qeCoRcFH2nM8wm5uJP+ZcVYGGayMavi8GmRJjuI3g1v6Z7Ni0JJKAJW+m0EtUuARb6Lmp4hMjzCBWzA=="], + + "@esbuild/linux-mips64el": ["@esbuild/linux-mips64el@0.27.4", "", { "os": "linux", "cpu": "none" }, "sha512-1RdrWFFiiLIW7LQq9Q2NES+HiD4NyT8Itj9AUeCl0IVCA459WnPhREKgwrpaIfTOe+/2rdntisegiPWn/r/aAw=="], + + "@esbuild/linux-ppc64": ["@esbuild/linux-ppc64@0.27.4", "", { "os": "linux", "cpu": "ppc64" }, "sha512-tLCwNG47l3sd9lpfyx9LAGEGItCUeRCWeAx6x2Jmbav65nAwoPXfewtAdtbtit/pJFLUWOhpv0FpS6GQAmPrHA=="], + + "@esbuild/linux-riscv64": ["@esbuild/linux-riscv64@0.27.4", "", { "os": "linux", "cpu": "none" }, "sha512-BnASypppbUWyqjd1KIpU4AUBiIhVr6YlHx/cnPgqEkNoVOhHg+YiSVxM1RLfiy4t9cAulbRGTNCKOcqHrEQLIw=="], + + "@esbuild/linux-s390x": ["@esbuild/linux-s390x@0.27.4", "", { "os": "linux", "cpu": "s390x" }, "sha512-+eUqgb/Z7vxVLezG8bVB9SfBie89gMueS+I0xYh2tJdw3vqA/0ImZJ2ROeWwVJN59ihBeZ7Tu92dF/5dy5FttA=="], + + "@esbuild/linux-x64": ["@esbuild/linux-x64@0.27.4", "", { "os": "linux", "cpu": "x64" }, "sha512-S5qOXrKV8BQEzJPVxAwnryi2+Iq5pB40gTEIT69BQONqR7JH1EPIcQ/Uiv9mCnn05jff9umq/5nqzxlqTOg9NA=="], + + "@esbuild/netbsd-arm64": ["@esbuild/netbsd-arm64@0.27.4", "", { "os": "none", "cpu": "arm64" }, "sha512-xHT8X4sb0GS8qTqiwzHqpY00C95DPAq7nAwX35Ie/s+LO9830hrMd3oX0ZMKLvy7vsonee73x0lmcdOVXFzd6Q=="], + + "@esbuild/netbsd-x64": ["@esbuild/netbsd-x64@0.27.4", "", { "os": "none", "cpu": "x64" }, "sha512-RugOvOdXfdyi5Tyv40kgQnI0byv66BFgAqjdgtAKqHoZTbTF2QqfQrFwa7cHEORJf6X2ht+l9ABLMP0dnKYsgg=="], + + "@esbuild/openbsd-arm64": ["@esbuild/openbsd-arm64@0.27.4", "", { "os": "openbsd", "cpu": "arm64" }, "sha512-2MyL3IAaTX+1/qP0O1SwskwcwCoOI4kV2IBX1xYnDDqthmq5ArrW94qSIKCAuRraMgPOmG0RDTA74mzYNQA9ow=="], + + "@esbuild/openbsd-x64": ["@esbuild/openbsd-x64@0.27.4", "", { "os": "openbsd", "cpu": "x64" }, "sha512-u8fg/jQ5aQDfsnIV6+KwLOf1CmJnfu1ShpwqdwC0uA7ZPwFws55Ngc12vBdeUdnuWoQYx/SOQLGDcdlfXhYmXQ=="], + + "@esbuild/openharmony-arm64": ["@esbuild/openharmony-arm64@0.27.4", "", { "os": "none", "cpu": "arm64" }, "sha512-JkTZrl6VbyO8lDQO3yv26nNr2RM2yZzNrNHEsj9bm6dOwwu9OYN28CjzZkH57bh4w0I2F7IodpQvUAEd1mbWXg=="], + + "@esbuild/sunos-x64": ["@esbuild/sunos-x64@0.27.4", "", { "os": "sunos", "cpu": "x64" }, "sha512-/gOzgaewZJfeJTlsWhvUEmUG4tWEY2Spp5M20INYRg2ZKl9QPO3QEEgPeRtLjEWSW8FilRNacPOg8R1uaYkA6g=="], + + "@esbuild/win32-arm64": ["@esbuild/win32-arm64@0.27.4", "", { "os": "win32", "cpu": "arm64" }, "sha512-Z9SExBg2y32smoDQdf1HRwHRt6vAHLXcxD2uGgO/v2jK7Y718Ix4ndsbNMU/+1Qiem9OiOdaqitioZwxivhXYg=="], + + "@esbuild/win32-ia32": ["@esbuild/win32-ia32@0.27.4", "", { "os": "win32", "cpu": "ia32" }, "sha512-DAyGLS0Jz5G5iixEbMHi5KdiApqHBWMGzTtMiJ72ZOLhbu/bzxgAe8Ue8CTS3n3HbIUHQz/L51yMdGMeoxXNJw=="], + + "@esbuild/win32-x64": ["@esbuild/win32-x64@0.27.4", "", { "os": "win32", "cpu": "x64" }, "sha512-+knoa0BDoeXgkNvvV1vvbZX4+hizelrkwmGJBdT17t8FNPwG2lKemmuMZlmaNQ3ws3DKKCxpb4zRZEIp3UxFCg=="], + + "@jridgewell/gen-mapping": ["@jridgewell/gen-mapping@0.3.13", "", { "dependencies": { "@jridgewell/sourcemap-codec": "^1.5.0", "@jridgewell/trace-mapping": "^0.3.24" } }, "sha512-2kkt/7niJ6MgEPxF0bYdQ6etZaA+fQvDcLKckhy1yIQOzaoKjBBjSj63/aLVjYE3qhRt5dvM+uUyfCg6UKCBbA=="], + + "@jridgewell/resolve-uri": ["@jridgewell/resolve-uri@3.1.2", "", {}, "sha512-bRISgCIjP20/tbWSPWMEi54QVPRZExkuD9lJL+UIxUKtwVJA8wW1Trb1jMs1RFXo1CBTNZ/5hpC9QvmKWdopKw=="], + + "@jridgewell/sourcemap-codec": ["@jridgewell/sourcemap-codec@1.5.5", "", {}, "sha512-cYQ9310grqxueWbl+WuIUIaiUaDcj7WOq5fVhEljNVgRfOUhY9fy2zTvfoqWsnebh8Sl70VScFbICvJnLKB0Og=="], + + "@jridgewell/trace-mapping": ["@jridgewell/trace-mapping@0.3.31", "", { "dependencies": { "@jridgewell/resolve-uri": "^3.1.0", "@jridgewell/sourcemap-codec": "^1.4.14" } }, "sha512-zzNR+SdQSDJzc8joaeP8QQoCQr8NuYx2dIIytl1QeBEZHJ9uW6hebsrYgbz8hJwUQao3TWCMtmfV8Nu1twOLAw=="], + + "@opentelemetry/api": ["@opentelemetry/api@1.9.0", "", {}, "sha512-3giAOQvZiH5F9bMlMiv8+GSPMeqg0dbaeo58/0SlA9sxSqZhnUtxzX9/2FzyhS9sWQf5S0GJE0AKBrFqjpeYcg=="], + + "@rollup/rollup-android-arm-eabi": ["@rollup/rollup-android-arm-eabi@4.59.0", "", { "os": "android", "cpu": "arm" }, "sha512-upnNBkA6ZH2VKGcBj9Fyl9IGNPULcjXRlg0LLeaioQWueH30p6IXtJEbKAgvyv+mJaMxSm1l6xwDXYjpEMiLMg=="], + + "@rollup/rollup-android-arm64": ["@rollup/rollup-android-arm64@4.59.0", "", { "os": "android", "cpu": "arm64" }, "sha512-hZ+Zxj3SySm4A/DylsDKZAeVg0mvi++0PYVceVyX7hemkw7OreKdCvW2oQ3T1FMZvCaQXqOTHb8qmBShoqk69Q=="], + + "@rollup/rollup-darwin-arm64": ["@rollup/rollup-darwin-arm64@4.59.0", "", { "os": "darwin", "cpu": "arm64" }, "sha512-W2Psnbh1J8ZJw0xKAd8zdNgF9HRLkdWwwdWqubSVk0pUuQkoHnv7rx4GiF9rT4t5DIZGAsConRE3AxCdJ4m8rg=="], + + "@rollup/rollup-darwin-x64": ["@rollup/rollup-darwin-x64@4.59.0", "", { "os": "darwin", "cpu": "x64" }, "sha512-ZW2KkwlS4lwTv7ZVsYDiARfFCnSGhzYPdiOU4IM2fDbL+QGlyAbjgSFuqNRbSthybLbIJ915UtZBtmuLrQAT/w=="], + + "@rollup/rollup-freebsd-arm64": ["@rollup/rollup-freebsd-arm64@4.59.0", "", { "os": "freebsd", "cpu": "arm64" }, "sha512-EsKaJ5ytAu9jI3lonzn3BgG8iRBjV4LxZexygcQbpiU0wU0ATxhNVEpXKfUa0pS05gTcSDMKpn3Sx+QB9RlTTA=="], + + "@rollup/rollup-freebsd-x64": ["@rollup/rollup-freebsd-x64@4.59.0", "", { "os": "freebsd", "cpu": "x64" }, "sha512-d3DuZi2KzTMjImrxoHIAODUZYoUUMsuUiY4SRRcJy6NJoZ6iIqWnJu9IScV9jXysyGMVuW+KNzZvBLOcpdl3Vg=="], + + "@rollup/rollup-linux-arm-gnueabihf": ["@rollup/rollup-linux-arm-gnueabihf@4.59.0", "", { "os": "linux", "cpu": "arm" }, "sha512-t4ONHboXi/3E0rT6OZl1pKbl2Vgxf9vJfWgmUoCEVQVxhW6Cw/c8I6hbbu7DAvgp82RKiH7TpLwxnJeKv2pbsw=="], + + "@rollup/rollup-linux-arm-musleabihf": ["@rollup/rollup-linux-arm-musleabihf@4.59.0", "", { "os": "linux", "cpu": "arm" }, "sha512-CikFT7aYPA2ufMD086cVORBYGHffBo4K8MQ4uPS/ZnY54GKj36i196u8U+aDVT2LX4eSMbyHtyOh7D7Zvk2VvA=="], + + "@rollup/rollup-linux-arm64-gnu": ["@rollup/rollup-linux-arm64-gnu@4.59.0", "", { "os": "linux", "cpu": "arm64" }, "sha512-jYgUGk5aLd1nUb1CtQ8E+t5JhLc9x5WdBKew9ZgAXg7DBk0ZHErLHdXM24rfX+bKrFe+Xp5YuJo54I5HFjGDAA=="], + + "@rollup/rollup-linux-arm64-musl": ["@rollup/rollup-linux-arm64-musl@4.59.0", "", { "os": "linux", "cpu": "arm64" }, "sha512-peZRVEdnFWZ5Bh2KeumKG9ty7aCXzzEsHShOZEFiCQlDEepP1dpUl/SrUNXNg13UmZl+gzVDPsiCwnV1uI0RUA=="], + + "@rollup/rollup-linux-loong64-gnu": ["@rollup/rollup-linux-loong64-gnu@4.59.0", "", { "os": "linux", "cpu": "none" }, "sha512-gbUSW/97f7+r4gHy3Jlup8zDG190AuodsWnNiXErp9mT90iCy9NKKU0Xwx5k8VlRAIV2uU9CsMnEFg/xXaOfXg=="], + + "@rollup/rollup-linux-loong64-musl": ["@rollup/rollup-linux-loong64-musl@4.59.0", "", { "os": "linux", "cpu": "none" }, "sha512-yTRONe79E+o0FWFijasoTjtzG9EBedFXJMl888NBEDCDV9I2wGbFFfJQQe63OijbFCUZqxpHz1GzpbtSFikJ4Q=="], + + "@rollup/rollup-linux-ppc64-gnu": ["@rollup/rollup-linux-ppc64-gnu@4.59.0", "", { "os": "linux", "cpu": "ppc64" }, "sha512-sw1o3tfyk12k3OEpRddF68a1unZ5VCN7zoTNtSn2KndUE+ea3m3ROOKRCZxEpmT9nsGnogpFP9x6mnLTCaoLkA=="], + + "@rollup/rollup-linux-ppc64-musl": ["@rollup/rollup-linux-ppc64-musl@4.59.0", "", { "os": "linux", "cpu": "ppc64" }, "sha512-+2kLtQ4xT3AiIxkzFVFXfsmlZiG5FXYW7ZyIIvGA7Bdeuh9Z0aN4hVyXS/G1E9bTP/vqszNIN/pUKCk/BTHsKA=="], + + "@rollup/rollup-linux-riscv64-gnu": ["@rollup/rollup-linux-riscv64-gnu@4.59.0", "", { "os": "linux", "cpu": "none" }, "sha512-NDYMpsXYJJaj+I7UdwIuHHNxXZ/b/N2hR15NyH3m2qAtb/hHPA4g4SuuvrdxetTdndfj9b1WOmy73kcPRoERUg=="], + + "@rollup/rollup-linux-riscv64-musl": ["@rollup/rollup-linux-riscv64-musl@4.59.0", "", { "os": "linux", "cpu": "none" }, "sha512-nLckB8WOqHIf1bhymk+oHxvM9D3tyPndZH8i8+35p/1YiVoVswPid2yLzgX7ZJP0KQvnkhM4H6QZ5m0LzbyIAg=="], + + "@rollup/rollup-linux-s390x-gnu": ["@rollup/rollup-linux-s390x-gnu@4.59.0", "", { "os": "linux", "cpu": "s390x" }, "sha512-oF87Ie3uAIvORFBpwnCvUzdeYUqi2wY6jRFWJAy1qus/udHFYIkplYRW+wo+GRUP4sKzYdmE1Y3+rY5Gc4ZO+w=="], + + "@rollup/rollup-linux-x64-gnu": ["@rollup/rollup-linux-x64-gnu@4.59.0", "", { "os": "linux", "cpu": "x64" }, "sha512-3AHmtQq/ppNuUspKAlvA8HtLybkDflkMuLK4DPo77DfthRb71V84/c4MlWJXixZz4uruIH4uaa07IqoAkG64fg=="], + + "@rollup/rollup-linux-x64-musl": ["@rollup/rollup-linux-x64-musl@4.59.0", "", { "os": "linux", "cpu": "x64" }, "sha512-2UdiwS/9cTAx7qIUZB/fWtToJwvt0Vbo0zmnYt7ED35KPg13Q0ym1g442THLC7VyI6JfYTP4PiSOWyoMdV2/xg=="], + + "@rollup/rollup-openbsd-x64": ["@rollup/rollup-openbsd-x64@4.59.0", "", { "os": "openbsd", "cpu": "x64" }, "sha512-M3bLRAVk6GOwFlPTIxVBSYKUaqfLrn8l0psKinkCFxl4lQvOSz8ZrKDz2gxcBwHFpci0B6rttydI4IpS4IS/jQ=="], + + "@rollup/rollup-openharmony-arm64": ["@rollup/rollup-openharmony-arm64@4.59.0", "", { "os": "none", "cpu": "arm64" }, "sha512-tt9KBJqaqp5i5HUZzoafHZX8b5Q2Fe7UjYERADll83O4fGqJ49O1FsL6LpdzVFQcpwvnyd0i+K/VSwu/o/nWlA=="], + + "@rollup/rollup-win32-arm64-msvc": ["@rollup/rollup-win32-arm64-msvc@4.59.0", "", { "os": "win32", "cpu": "arm64" }, "sha512-V5B6mG7OrGTwnxaNUzZTDTjDS7F75PO1ae6MJYdiMu60sq0CqN5CVeVsbhPxalupvTX8gXVSU9gq+Rx1/hvu6A=="], + + "@rollup/rollup-win32-ia32-msvc": ["@rollup/rollup-win32-ia32-msvc@4.59.0", "", { "os": "win32", "cpu": "ia32" }, "sha512-UKFMHPuM9R0iBegwzKF4y0C4J9u8C6MEJgFuXTBerMk7EJ92GFVFYBfOZaSGLu6COf7FxpQNqhNS4c4icUPqxA=="], + + "@rollup/rollup-win32-x64-gnu": ["@rollup/rollup-win32-x64-gnu@4.59.0", "", { "os": "win32", "cpu": "x64" }, "sha512-laBkYlSS1n2L8fSo1thDNGrCTQMmxjYY5G0WFWjFFYZkKPjsMBsgJfGf4TLxXrF6RyhI60L8TMOjBMvXiTcxeA=="], + + "@rollup/rollup-win32-x64-msvc": ["@rollup/rollup-win32-x64-msvc@4.59.0", "", { "os": "win32", "cpu": "x64" }, "sha512-2HRCml6OztYXyJXAvdDXPKcawukWY2GpR5/nxKp4iBgiO3wcoEGkAaqctIbZcNB6KlUQBIqt8VYkNSj2397EfA=="], + + "@standard-schema/spec": ["@standard-schema/spec@1.1.0", "", {}, "sha512-l2aFy5jALhniG5HgqrD6jXLi/rUWrKvqN/qJx6yoJsgKhblVd+iqqU4RCXavm/jPityDo5TCvKMnpjKnOriy0w=="], + + "@types/estree": ["@types/estree@1.0.8", "", {}, "sha512-dWHzHa2WqEXI/O1E9OjrocMTKJl2mSrEolh1Iomrv6U+JuNwaHXsXx9bLu5gG7BUWFIN0skIQJQ/L1rIex4X6w=="], + + "@types/node": ["@types/node@24.12.0", "", { "dependencies": { "undici-types": "~7.16.0" } }, "sha512-GYDxsZi3ChgmckRT9HPU0WEhKLP08ev/Yfcq2AstjrDASOYCSXeyjDsHg4v5t4jOj7cyDX3vmprafKlWIG9MXQ=="], + + "@vercel/oidc": ["@vercel/oidc@3.1.0", "", {}, "sha512-Fw28YZpRnA3cAHHDlkt7xQHiJ0fcL+NRcIqsocZQUSmbzeIKRpwttJjik5ZGanXP+vlA4SbTg+AbA3bP363l+w=="], + + "acorn": ["acorn@8.16.0", "", { "bin": { "acorn": "bin/acorn" } }, "sha512-UVJyE9MttOsBQIDKw1skb9nAwQuR5wuGD3+82K6JgJlm/Y+KI92oNsMNGZCYdDsVtRHSak0pcV5Dno5+4jh9sw=="], + + "ai": ["ai@6.0.116", "", { "dependencies": { "@ai-sdk/gateway": "3.0.66", "@ai-sdk/provider": "3.0.8", "@ai-sdk/provider-utils": "4.0.19", "@opentelemetry/api": "1.9.0" }, "peerDependencies": { "zod": "^3.25.76 || ^4.1.8" } }, "sha512-7yM+cTmyRLeNIXwt4Vj+mrrJgVQ9RMIW5WO0ydoLoYkewIvsMcvUmqS4j2RJTUXaF1HphwmSKUMQ/HypNRGOmA=="], + + "any-promise": ["any-promise@1.3.0", "", {}, "sha512-7UvmKalWRt1wgjL1RrGxoSJW/0QZFIegpeGvZG9kjp8vrRu55XTHbwnqq2GpXm9uLbcuhxm3IqX9OB4MZR1b2A=="], + + "autumn-js": ["autumn-js@1.0.5", "", { "dependencies": { "query-string": "^9.2.2", "rou3": "^0.6.1", "zod": "^4.0.0" }, "peerDependencies": { "better-auth": "^1.3.17", "better-call": "^1.0.12", "express": "^5.2.1", "hono": "^4.0.0", "next": "^14.0.0 || ^15.0.0", "react": "^18.0.0 || ^19.0.0" }, "optionalPeers": ["better-auth", "better-call", "express", "hono", "next", "react"] }, "sha512-Xh4hqx7EO+hilVjCuilLjUt5iw6RP8vxKmkKmwwHM8WhrHH4StAol17Qis9FtZ/ixeHfe4TOFojzgLhP74kETw=="], + + "bundle-require": ["bundle-require@5.1.0", "", { "dependencies": { "load-tsconfig": "^0.2.3" }, "peerDependencies": { "esbuild": ">=0.18" } }, "sha512-3WrrOuZiyaaZPWiEt4G3+IffISVC9HYlWueJEBWED4ZH4aIAC2PnkdnuRrR94M+w6yGWn4AglWtJtBI8YqvgoA=="], + + "cac": ["cac@6.7.14", "", {}, "sha512-b6Ilus+c3RrdDk+JhLKUAQfzzgLEPy6wcXqS7f/xe1EETvsDP6GORG7SFuOs6cID5YkqchW/LXZbX5bc8j7ZcQ=="], + + "chokidar": ["chokidar@4.0.3", "", { "dependencies": { "readdirp": "^4.0.1" } }, "sha512-Qgzu8kfBvo+cA4962jnP1KkS6Dop5NS6g7R5LFYJr4b8Ub94PPQXUksCw9PvXoeXPRRddRNC5C1JQUR2SMGtnA=="], + + "commander": ["commander@4.1.1", "", {}, "sha512-NOKm8xhkzAjzFx8B2v5OAHT+u5pRQc2UCa2Vq9jYL/31o2wi9mxBA7LIFs3sV5VSC49z6pEhfbMULvShKj26WA=="], + + "confbox": ["confbox@0.1.8", "", {}, "sha512-RMtmw0iFkeR4YV+fUOSucriAQNb9g8zFR52MWCtl+cCZOFRNL6zeB395vPzFhEjjn4fMxXudmELnl/KF/WrK6w=="], + + "consola": ["consola@3.4.2", "", {}, "sha512-5IKcdX0nnYavi6G7TtOhwkYzyjfJlatbjMjuLSfE2kYT5pMDOilZ4OvMhi637CcDICTmz3wARPoyhqyX1Y+XvA=="], + + "debug": ["debug@4.4.3", "", { "dependencies": { "ms": "^2.1.3" } }, "sha512-RGwwWnwQvkVfavKVt22FGLw+xYSdzARwm0ru6DhTVA3umU5hZc28V3kO4stgYryrTlLpuvgI9GiijltAjNbcqA=="], + + "decode-uri-component": ["decode-uri-component@0.4.1", "", {}, "sha512-+8VxcR21HhTy8nOt6jf20w0c9CADrw1O8d+VZ/YzzCt4bJ3uBjw+D1q2osAB8RnpwwaeYBxy0HyKQxD5JBMuuQ=="], + + "esbuild": ["esbuild@0.27.4", "", { "optionalDependencies": { "@esbuild/aix-ppc64": "0.27.4", "@esbuild/android-arm": "0.27.4", "@esbuild/android-arm64": "0.27.4", "@esbuild/android-x64": "0.27.4", "@esbuild/darwin-arm64": "0.27.4", "@esbuild/darwin-x64": "0.27.4", "@esbuild/freebsd-arm64": "0.27.4", "@esbuild/freebsd-x64": "0.27.4", "@esbuild/linux-arm": "0.27.4", "@esbuild/linux-arm64": "0.27.4", "@esbuild/linux-ia32": "0.27.4", "@esbuild/linux-loong64": "0.27.4", "@esbuild/linux-mips64el": "0.27.4", "@esbuild/linux-ppc64": "0.27.4", "@esbuild/linux-riscv64": "0.27.4", "@esbuild/linux-s390x": "0.27.4", "@esbuild/linux-x64": "0.27.4", "@esbuild/netbsd-arm64": "0.27.4", "@esbuild/netbsd-x64": "0.27.4", "@esbuild/openbsd-arm64": "0.27.4", "@esbuild/openbsd-x64": "0.27.4", "@esbuild/openharmony-arm64": "0.27.4", "@esbuild/sunos-x64": "0.27.4", "@esbuild/win32-arm64": "0.27.4", "@esbuild/win32-ia32": "0.27.4", "@esbuild/win32-x64": "0.27.4" }, "bin": { "esbuild": "bin/esbuild" } }, "sha512-Rq4vbHnYkK5fws5NF7MYTU68FPRE1ajX7heQ/8QXXWqNgqqJ/GkmmyxIzUnf2Sr/bakf8l54716CcMGHYhMrrQ=="], + + "eventsource-parser": ["eventsource-parser@3.0.6", "", {}, "sha512-Vo1ab+QXPzZ4tCa8SwIHJFaSzy4R6SHf7BY79rFBDf0idraZWAkYrDjDj8uWaSm3S2TK+hJ7/t1CEmZ7jXw+pg=="], + + "fdir": ["fdir@6.5.0", "", { "peerDependencies": { "picomatch": "^3 || ^4" }, "optionalPeers": ["picomatch"] }, "sha512-tIbYtZbucOs0BRGqPJkshJUYdL+SDH7dVM8gjy+ERp3WAUjLEFJE+02kanyHtwjWOnwrKYBiwAmM0p4kLJAnXg=="], + + "filter-obj": ["filter-obj@5.1.0", "", {}, "sha512-qWeTREPoT7I0bifpPUXtxkZJ1XJzxWtfoWWkdVGqa+eCr3SHW/Ocp89o8vLvbUuQnadybJpjOKu4V+RwO6sGng=="], + + "fix-dts-default-cjs-exports": ["fix-dts-default-cjs-exports@1.0.1", "", { "dependencies": { "magic-string": "^0.30.17", "mlly": "^1.7.4", "rollup": "^4.34.8" } }, "sha512-pVIECanWFC61Hzl2+oOCtoJ3F17kglZC/6N94eRWycFgBH35hHx0Li604ZIzhseh97mf2p0cv7vVrOZGoqhlEg=="], + + "fsevents": ["fsevents@2.3.3", "", { "os": "darwin" }, "sha512-5xoDfX+fL7faATnagmWPpbFtwh/R77WmMMqqHGS65C3vvB0YHrgF+B1YmZ3441tMj5n63k0212XNoJwzlhffQw=="], + + "joycon": ["joycon@3.1.1", "", {}, "sha512-34wB/Y7MW7bzjKRjUKTa46I2Z7eV62Rkhva+KkopW7Qvv/OSWBqvkSY7vusOPrNuZcUG3tApvdVgNB8POj3SPw=="], + + "json-schema": ["json-schema@0.4.0", "", {}, "sha512-es94M3nTIfsEPisRafak+HDLfHXnKBhV3vU5eqPcS3flIWqcxJWgXHXiey3YrpaNsanY5ei1VoYEbOzijuq9BA=="], + + "lilconfig": ["lilconfig@3.1.3", "", {}, "sha512-/vlFKAoH5Cgt3Ie+JLhRbwOsCQePABiU3tJ1egGvyQ+33R/vcwM2Zl2QR/LzjsBeItPt3oSVXapn+m4nQDvpzw=="], + + "lines-and-columns": ["lines-and-columns@1.2.4", "", {}, "sha512-7ylylesZQ/PV29jhEDl3Ufjo6ZX7gCqJr5F7PKrqc93v7fzSymt1BpwEU8nAUXs8qzzvqhbjhK5QZg6Mt/HkBg=="], + + "load-tsconfig": ["load-tsconfig@0.2.5", "", {}, "sha512-IXO6OCs9yg8tMKzfPZ1YmheJbZCiEsnBdcB03l0OcfK9prKnJb96siuHCr5Fl37/yo9DnKU+TLpxzTUspw9shg=="], + + "magic-string": ["magic-string@0.30.21", "", { "dependencies": { "@jridgewell/sourcemap-codec": "^1.5.5" } }, "sha512-vd2F4YUyEXKGcLHoq+TEyCjxueSeHnFxyyjNp80yg0XV4vUhnDer/lvvlqM/arB5bXQN5K2/3oinyCRyx8T2CQ=="], + + "mlly": ["mlly@1.8.1", "", { "dependencies": { "acorn": "^8.16.0", "pathe": "^2.0.3", "pkg-types": "^1.3.1", "ufo": "^1.6.3" } }, "sha512-SnL6sNutTwRWWR/vcmCYHSADjiEesp5TGQQ0pXyLhW5IoeibRlF/CbSLailbB3CNqJUk9cVJ9dUDnbD7GrcHBQ=="], + + "ms": ["ms@2.1.3", "", {}, "sha512-6FlzubTLZG3J2a/NVCAleEhjzq5oxgHyaCU9yYXvcLsvoVaHJq/s5xXI6/XXP6tz7R9xAOtHnSO/tXtF3WRTlA=="], + + "mz": ["mz@2.7.0", "", { "dependencies": { "any-promise": "^1.0.0", "object-assign": "^4.0.1", "thenify-all": "^1.0.0" } }, "sha512-z81GNO7nnYMEhrGh9LeymoE4+Yr0Wn5McHIZMK5cfQCl+NDX08sCZgUc9/6MHni9IWuFLm1Z3HTCXu2z9fN62Q=="], + + "object-assign": ["object-assign@4.1.1", "", {}, "sha512-rJgTQnkUnH1sFw8yT6VSU3zD3sWmu6sZhIseY8VX+GRu3P6F7Fu+JNDoXfklElbLJSnc3FUQHVe4cU5hj+BcUg=="], + + "pathe": ["pathe@2.0.3", "", {}, "sha512-WUjGcAqP1gQacoQe+OBJsFA7Ld4DyXuUIjZ5cc75cLHvJ7dtNsTugphxIADwspS+AraAUePCKrSVtPLFj/F88w=="], + + "picocolors": ["picocolors@1.1.1", "", {}, "sha512-xceH2snhtb5M9liqDsmEw56le376mTZkEX/jEb/RxNFyegNul7eNslCXP9FDj/Lcu0X8KEyMceP2ntpaHrDEVA=="], + + "picomatch": ["picomatch@4.0.3", "", {}, "sha512-5gTmgEY/sqK6gFXLIsQNH19lWb4ebPDLA4SdLP7dsWkIXHWlG66oPuVvXSGFPppYZz8ZDZq0dYYrbHfBCVUb1Q=="], + + "pirates": ["pirates@4.0.7", "", {}, "sha512-TfySrs/5nm8fQJDcBDuUng3VOUKsd7S+zqvbOTiGXHfxX4wK31ard+hoNuvkicM/2YFzlpDgABOevKSsB4G/FA=="], + + "pkg-types": ["pkg-types@1.3.1", "", { "dependencies": { "confbox": "^0.1.8", "mlly": "^1.7.4", "pathe": "^2.0.1" } }, "sha512-/Jm5M4RvtBFVkKWRu2BLUTNP8/M2a+UwuAX+ae4770q1qVGtfjG+WTCupoZixokjmHiry8uI+dlY8KXYV5HVVQ=="], + + "postcss-load-config": ["postcss-load-config@6.0.1", "", { "dependencies": { "lilconfig": "^3.1.1" }, "peerDependencies": { "jiti": ">=1.21.0", "postcss": ">=8.0.9", "tsx": "^4.8.1", "yaml": "^2.4.2" }, "optionalPeers": ["jiti", "postcss", "tsx", "yaml"] }, "sha512-oPtTM4oerL+UXmx+93ytZVN82RrlY/wPUV8IeDxFrzIjXOLF1pN+EmKPLbubvKHT2HC20xXsCAH2Z+CKV6Oz/g=="], + + "query-string": ["query-string@9.3.1", "", { "dependencies": { "decode-uri-component": "^0.4.1", "filter-obj": "^5.1.0", "split-on-first": "^3.0.0" } }, "sha512-5fBfMOcDi5SA9qj5jZhWAcTtDfKF5WFdd2uD9nVNlbxVv1baq65aALy6qofpNEGELHvisjjasxQp7BlM9gvMzw=="], + + "react": ["react@19.2.4", "", {}, "sha512-9nfp2hYpCwOjAN+8TZFGhtWEwgvWHXqESH8qT89AT/lWklpLON22Lc8pEtnpsZz7VmawabSU0gCjnj8aC0euHQ=="], + + "readdirp": ["readdirp@4.1.2", "", {}, "sha512-GDhwkLfywWL2s6vEjyhri+eXmfH6j1L7JE27WhqLeYzoh/A3DBaYGEj2H/HFZCn/kMfim73FXxEJTw06WtxQwg=="], + + "resolve-from": ["resolve-from@5.0.0", "", {}, "sha512-qYg9KP24dD5qka9J47d0aVky0N+b4fTU89LN9iDnjB5waksiC49rvMB0PrUJQGoTmH50XPiqOvAjDfaijGxYZw=="], + + "rollup": ["rollup@4.59.0", "", { "dependencies": { "@types/estree": "1.0.8" }, "optionalDependencies": { "@rollup/rollup-android-arm-eabi": "4.59.0", "@rollup/rollup-android-arm64": "4.59.0", "@rollup/rollup-darwin-arm64": "4.59.0", "@rollup/rollup-darwin-x64": "4.59.0", "@rollup/rollup-freebsd-arm64": "4.59.0", "@rollup/rollup-freebsd-x64": "4.59.0", "@rollup/rollup-linux-arm-gnueabihf": "4.59.0", "@rollup/rollup-linux-arm-musleabihf": "4.59.0", "@rollup/rollup-linux-arm64-gnu": "4.59.0", "@rollup/rollup-linux-arm64-musl": "4.59.0", "@rollup/rollup-linux-loong64-gnu": "4.59.0", "@rollup/rollup-linux-loong64-musl": "4.59.0", "@rollup/rollup-linux-ppc64-gnu": "4.59.0", "@rollup/rollup-linux-ppc64-musl": "4.59.0", "@rollup/rollup-linux-riscv64-gnu": "4.59.0", "@rollup/rollup-linux-riscv64-musl": "4.59.0", "@rollup/rollup-linux-s390x-gnu": "4.59.0", "@rollup/rollup-linux-x64-gnu": "4.59.0", "@rollup/rollup-linux-x64-musl": "4.59.0", "@rollup/rollup-openbsd-x64": "4.59.0", "@rollup/rollup-openharmony-arm64": "4.59.0", "@rollup/rollup-win32-arm64-msvc": "4.59.0", "@rollup/rollup-win32-ia32-msvc": "4.59.0", "@rollup/rollup-win32-x64-gnu": "4.59.0", "@rollup/rollup-win32-x64-msvc": "4.59.0", "fsevents": "~2.3.2" }, "bin": { "rollup": "dist/bin/rollup" } }, "sha512-2oMpl67a3zCH9H79LeMcbDhXW/UmWG/y2zuqnF2jQq5uq9TbM9TVyXvA4+t+ne2IIkBdrLpAaRQAvo7YI/Yyeg=="], + + "rou3": ["rou3@0.6.3", "", {}, "sha512-1HSG1ENTj7Kkm5muMnXuzzfdDOf7CFnbSYFA+H3Fp/rB9lOCxCPgy1jlZxTKyFoC5jJay8Mmc+VbPLYRjzYLrA=="], + + "source-map": ["source-map@0.7.6", "", {}, "sha512-i5uvt8C3ikiWeNZSVZNWcfZPItFQOsYTUAOkcUPGd8DqDy1uOUikjt5dG+uRlwyvR108Fb9DOd4GvXfT0N2/uQ=="], + + "split-on-first": ["split-on-first@3.0.0", "", {}, "sha512-qxQJTx2ryR0Dw0ITYyekNQWpz6f8dGd7vffGNflQQ3Iqj9NJ6qiZ7ELpZsJ/QBhIVAiDfXdag3+Gp8RvWa62AA=="], + + "sucrase": ["sucrase@3.35.1", "", { "dependencies": { "@jridgewell/gen-mapping": "^0.3.2", "commander": "^4.0.0", "lines-and-columns": "^1.1.6", "mz": "^2.7.0", "pirates": "^4.0.1", "tinyglobby": "^0.2.11", "ts-interface-checker": "^0.1.9" }, "bin": { "sucrase": "bin/sucrase", "sucrase-node": "bin/sucrase-node" } }, "sha512-DhuTmvZWux4H1UOnWMB3sk0sbaCVOoQZjv8u1rDoTV0HTdGem9hkAZtl4JZy8P2z4Bg0nT+YMeOFyVr4zcG5Tw=="], + + "thenify": ["thenify@3.3.1", "", { "dependencies": { "any-promise": "^1.0.0" } }, "sha512-RVZSIV5IG10Hk3enotrhvz0T9em6cyHBLkH/YAZuKqd8hRkKhSfCGIcP2KUY0EPxndzANBmNllzWPwak+bheSw=="], + + "thenify-all": ["thenify-all@1.6.0", "", { "dependencies": { "thenify": ">= 3.1.0 < 4" } }, "sha512-RNxQH/qI8/t3thXJDwcstUO4zeqo64+Uy/+sNVRBx4Xn2OX+OZ9oP+iJnNFqplFra2ZUVeKCSa2oVWi3T4uVmA=="], + + "tinyexec": ["tinyexec@0.3.2", "", {}, "sha512-KQQR9yN7R5+OSwaK0XQoj22pwHoTlgYqmUscPYoknOoWCWfj/5/ABTMRi69FrKU5ffPVh5QcFikpWJI/P1ocHA=="], + + "tinyglobby": ["tinyglobby@0.2.15", "", { "dependencies": { "fdir": "^6.5.0", "picomatch": "^4.0.3" } }, "sha512-j2Zq4NyQYG5XMST4cbs02Ak8iJUdxRM0XI5QyxXuZOzKOINmWurp3smXu3y5wDcJrptwpSjgXHzIQxR0omXljQ=="], + + "tree-kill": ["tree-kill@1.2.2", "", { "bin": { "tree-kill": "cli.js" } }, "sha512-L0Orpi8qGpRG//Nd+H90vFB+3iHnue1zSSGmNOOCh1GLJ7rUKVwV2HvijphGQS2UmhUZewS9VgvxYIdgr+fG1A=="], + + "ts-interface-checker": ["ts-interface-checker@0.1.13", "", {}, "sha512-Y/arvbn+rrz3JCKl9C4kVNfTfSm2/mEp5FSz5EsZSANGPSlQrpRI5M4PKF+mJnE52jOO90PnPSc3Ur3bTQw0gA=="], + + "tsup": ["tsup@8.5.1", "", { "dependencies": { "bundle-require": "^5.1.0", "cac": "^6.7.14", "chokidar": "^4.0.3", "consola": "^3.4.0", "debug": "^4.4.0", "esbuild": "^0.27.0", "fix-dts-default-cjs-exports": "^1.0.0", "joycon": "^3.1.1", "picocolors": "^1.1.1", "postcss-load-config": "^6.0.1", "resolve-from": "^5.0.0", "rollup": "^4.34.8", "source-map": "^0.7.6", "sucrase": "^3.35.0", "tinyexec": "^0.3.2", "tinyglobby": "^0.2.11", "tree-kill": "^1.2.2" }, "peerDependencies": { "@microsoft/api-extractor": "^7.36.0", "@swc/core": "^1", "postcss": "^8.4.12", "typescript": ">=4.5.0" }, "optionalPeers": ["@microsoft/api-extractor", "@swc/core", "postcss", "typescript"], "bin": { "tsup": "dist/cli-default.js", "tsup-node": "dist/cli-node.js" } }, "sha512-xtgkqwdhpKWr3tKPmCkvYmS9xnQK3m3XgxZHwSUjvfTjp7YfXe5tT3GgWi0F2N+ZSMsOeWeZFh7ZZFg5iPhing=="], + + "typescript": ["typescript@5.9.3", "", { "bin": { "tsc": "bin/tsc", "tsserver": "bin/tsserver" } }, "sha512-jl1vZzPDinLr9eUt3J/t7V6FgNEw9QjvBPdysz9KfQDD41fQrC2Y4vKQdiaUpFT4bXlb1RHhLpp8wtm6M5TgSw=="], + + "ufo": ["ufo@1.6.3", "", {}, "sha512-yDJTmhydvl5lJzBmy/hyOAA0d+aqCBuwl818haVdYCRrWV84o7YyeVm4QlVHStqNrrJSTb6jKuFAVqAFsr+K3Q=="], + + "undici-types": ["undici-types@7.16.0", "", {}, "sha512-Zz+aZWSj8LE6zoxD+xrjh4VfkIG8Ya6LvYkZqtUQGJPZjYl53ypCaUwWqo7eI0x66KBGeRo+mlBEkMSeSZ38Nw=="], + + "zod": ["zod@4.3.6", "", {}, "sha512-rftlrkhHZOcjDwkGlnUtZZkvaPHCsDATp4pGpuOOMDaTdDDXF91wuVDJoWoPsKX/3YPQ5fHuF3STjcYyKr+Qhg=="], + } +} diff --git a/packages/ai-sdk/package.json b/packages/ai-sdk/package.json new file mode 100644 index 000000000..f8adf149d --- /dev/null +++ b/packages/ai-sdk/package.json @@ -0,0 +1,35 @@ +{ + "name": "@useautumn/ai-sdk", + "version": "0.0.1", + "type": "module", + "description": "AI SDK for Autumn", + "main": "./dist/sdk/index.cjs", + "module": "./dist/sdk/index.js", + "types": "./dist/sdk/index.d.ts", + "files": [ + "dist", + "README.md", + "LICENSE.md" + ], + "exports": { + ".": { + "types": "./dist/sdk/index.d.ts", + "require": "./dist/sdk/index.cjs", + "import": "./dist/sdk/index.js" + } + }, + "scripts": { + "ts": "tsgo --noEmit --skipLibCheck", + "build": "rm -rf dist && tsup", + "prepublishOnly": "bun run build" + }, + "peerDependencies": { + "ai": "^6.0.116", + "autumn-js": "*" + }, + "devDependencies": { + "@types/node": "^24.9.1", + "tsup": "^8.4.0", + "typescript": "^5.8.3" + } +} diff --git a/packages/ai-sdk/src/index.ts b/packages/ai-sdk/src/index.ts new file mode 100644 index 000000000..544423e5a --- /dev/null +++ b/packages/ai-sdk/src/index.ts @@ -0,0 +1,104 @@ +import type { LanguageModelV3, LanguageModelV3Usage } from "@ai-sdk/provider"; +import { + type LanguageModelMiddleware, + type LanguageModelUsage, + wrapLanguageModel, +} from "ai"; +import type { Autumn } from "autumn-js"; + +type TokenCount = + | LanguageModelV3Usage["inputTokens"] + | LanguageModelUsage["inputTokens"]; + +export const withTokenTracking = ({ + autumn, + model, + customerId, + providerId, + featureId, + entityId, + properties, +}: { + /** Autumn SDK client instance. */ + autumn: Autumn; + /** The AI SDK language model to wrap. */ + model: LanguageModelV3; + /** The Autumn customer ID to attribute usage to. */ + customerId: string; + /** Override the provider prefix used in the model name. Falls back to `model.provider`. */ + providerId?: "custom" | string; + /** Target a specific AI credit system feature. Auto-detected if omitted. */ + featureId?: string; + /** Entity ID for entity-scoped balance tracking. */ + entityId?: string; + /** Additional properties to attach to the usage event. */ + properties?: Record; +}) => { + const provider = providerId ?? model.provider; + const modelName = `${provider}/${model.modelId}`; + + const resolveTokens = (tokens: TokenCount, label: string): number => { + const value = typeof tokens === "number" ? tokens : tokens?.total; + if (value == null) + throw new Error( + `[Autumn] ${label} token usage was not returned by the model provider (${modelName}). This provider may not support usage tracking.`, + ); + return value; + }; + + const trackUsage = async ( + usage: LanguageModelV3Usage | LanguageModelUsage, + ) => { + try { + await autumn.balances.trackTokens({ + customerId, + modelId: modelName, + inputTokens: resolveTokens(usage.inputTokens, "Input"), + outputTokens: resolveTokens(usage.outputTokens, "Output"), + featureId, + entityId, + properties, + }); + } catch (error) { + console.error("[Autumn Tracking] Failed to track usage:", error); + } + }; + + const middleware: LanguageModelMiddleware = { + specificationVersion: "v3", + wrapGenerate: async ({ doGenerate }) => { + const result = await doGenerate(); + await trackUsage(result.usage); + return result; + }, + wrapStream: async ({ doStream }) => { + const { stream, ...rest } = await doStream(); + + let trackingPromise: Promise | undefined; + + type StreamChunk = + typeof stream extends ReadableStream ? T : never; + + const transformStream = new TransformStream({ + transform(chunk, controller) { + if (chunk.type === "finish" && chunk.usage) { + trackingPromise = trackUsage(chunk.usage); + } + controller.enqueue(chunk); + }, + async flush() { + if (trackingPromise) { + await trackingPromise; + } + }, + }); + + return { + stream: stream.pipeThrough(transformStream), + ...rest, + }; + }, + }; + + return wrapLanguageModel({ model: model, middleware }); +}; diff --git a/packages/ai-sdk/tsconfig.json b/packages/ai-sdk/tsconfig.json new file mode 100644 index 000000000..c3ea081e4 --- /dev/null +++ b/packages/ai-sdk/tsconfig.json @@ -0,0 +1,19 @@ +{ + "compilerOptions": { + "strict": true, + "noImplicitAny": true, + "esModuleInterop": true, + "strictNullChecks": true, + "target": "ES2022", + "moduleResolution": "NodeNext", + "module": "NodeNext", + "declaration": true, + "isolatedModules": true, + "noEmit": true, + "outDir": "dist", + "lib": ["ES2022"], + "skipLibCheck": true + }, + "include": ["src"], + "exclude": ["dist", "node_modules"] +} diff --git a/packages/ai-sdk/tsup.config.ts b/packages/ai-sdk/tsup.config.ts new file mode 100644 index 000000000..5ea218dce --- /dev/null +++ b/packages/ai-sdk/tsup.config.ts @@ -0,0 +1,12 @@ +import { defineConfig } from "tsup"; + +export default defineConfig({ + entry: { + "sdk/index": "src/index.ts", + }, + format: ["cjs", "esm"], + dts: true, + splitting: false, + sourcemap: false, + clean: true, +}); diff --git a/packages/atmn/src/compose/models/featureModels.ts b/packages/atmn/src/compose/models/featureModels.ts index 5a2e63bdb..fac022ce7 100644 --- a/packages/atmn/src/compose/models/featureModels.ts +++ b/packages/atmn/src/compose/models/featureModels.ts @@ -85,11 +85,17 @@ export type CreditSystemFeature = FeatureBase & { }>; }; -/** - * Feature definition with type-safe constraints: - * - Boolean features cannot have consumable - * - Metered features require consumable (true = single_use style, false = continuous_use style) - * - Credit system features are always consumable and require creditSchema - */ -export type Feature = BooleanFeature | MeteredFeature | CreditSystemFeature; +export type ModelMarkupEntry = { + markup: number; + inputCost?: number; + outputCost?: number; +}; + +/** AI credit system feature - uses model-based pricing */ +export type AiCreditSystemFeature = FeatureBase & { + type: "ai_credit_system"; + modelMarkups?: Record; +}; + +export type Feature = BooleanFeature | MeteredFeature | CreditSystemFeature | AiCreditSystemFeature; diff --git a/packages/atmn/src/lib/transforms/apiToSdk/Transformer.test.ts b/packages/atmn/src/lib/transforms/apiToSdk/Transformer.test.ts index 0d77ee067..d8488addc 100644 --- a/packages/atmn/src/lib/transforms/apiToSdk/Transformer.test.ts +++ b/packages/atmn/src/lib/transforms/apiToSdk/Transformer.test.ts @@ -6,14 +6,14 @@ import { createTransformer } from "./Transformer.js"; describe("Transformer", () => { describe("Feature transforms", () => { test("boolean feature", () => { - const apiFeature = { + const result = transformApiFeature({ id: "enabled", name: "Feature Enabled", type: "boolean", + consumable: false, + archived: false, event_names: [], - }; - - const result = transformApiFeature(apiFeature); + }); expect(result.type).toBe("boolean"); expect(result.id).toBe("enabled"); @@ -21,47 +21,72 @@ describe("Transformer", () => { }); test("single_use → metered with consumable=true", () => { - const apiFeature = { + const result = transformApiFeature({ id: "api_calls", name: "API Calls", type: "single_use", + consumable: true, + archived: false, event_names: ["api.call"], - }; - - const result = transformApiFeature(apiFeature); + }); expect(result.type).toBe("metered"); - expect(result.consumable).toBe(true); + if (result.type === "metered") { + expect(result.consumable).toBe(true); + } expect(result.id).toBe("api_calls"); }); test("continuous_use → metered with consumable=false", () => { - const apiFeature = { + const result = transformApiFeature({ id: "seats", name: "Seats", type: "continuous_use", + consumable: false, + archived: false, event_names: [], - }; - - const result = transformApiFeature(apiFeature); + }); expect(result.type).toBe("metered"); - expect(result.consumable).toBe(false); + if (result.type === "metered") { + expect(result.consumable).toBe(false); + } }); test("credit_system", () => { - const apiFeature = { + const result = transformApiFeature({ id: "credits", name: "Credits", type: "credit_system", + consumable: true, + archived: false, credit_schema: [{ metered_feature_id: "api_calls", credit_cost: 10 }], - }; - - const result = transformApiFeature(apiFeature); + }); expect(result.type).toBe("credit_system"); - expect(result.consumable).toBe(true); - expect(result.creditSchema).toHaveLength(1); + if (result.type === "credit_system") { + expect(result.consumable).toBe(true); + expect(result.creditSchema).toHaveLength(1); + } + }); + + test("ai_credit_system", () => { + const result = transformApiFeature({ + id: "ai_credits", + name: "AI Credits", + type: "ai_credit_system", + consumable: true, + archived: false, + model_markups: { + "anthropic/claude-opus-4-5": { markup: 20 }, + }, + }); + + expect(result.type).toBe("ai_credit_system"); + if (result.type === "ai_credit_system") { + expect(result.modelMarkups).toBeDefined(); + expect(result.modelMarkups!["anthropic/claude-opus-4-5"].markup).toBe(20); + } }); }); diff --git a/packages/atmn/src/lib/transforms/apiToSdk/feature.ts b/packages/atmn/src/lib/transforms/apiToSdk/feature.ts index 9d5519e30..e4ba94ffb 100644 --- a/packages/atmn/src/lib/transforms/apiToSdk/feature.ts +++ b/packages/atmn/src/lib/transforms/apiToSdk/feature.ts @@ -1,26 +1,38 @@ -import type { Feature } from "../../../compose/models/featureModels.js"; +import type { Feature, ModelMarkupEntry } from "../../../compose/models/featureModels.js"; +import type { ApiFeature } from "../../api/types/feature.js"; import { createTransformer } from "./Transformer.js"; +type RawApiFeature = Omit & { type: string }; + function mapCreditSchema( - api: any, + api: RawApiFeature, ): Array<{ meteredFeatureId: string; creditCost: number }> { - return (api.credit_schema ?? []).map( - (cs: { metered_feature_id: string; credit_cost: number }) => ({ - meteredFeatureId: cs.metered_feature_id, - creditCost: cs.credit_cost, - }), + return (api.credit_schema ?? []).map((cs) => ({ + meteredFeatureId: cs.metered_feature_id, + creditCost: cs.credit_cost, + })); +} + +function mapModelMarkups(api: RawApiFeature): Record | undefined { + if (!api.model_markups) return undefined; + return Object.fromEntries( + Object.entries(api.model_markups).map(([modelId, entry]) => [ + modelId, + { + markup: entry.markup, + inputCost: entry.input_cost, + outputCost: entry.output_cost, + }, + ]) ); } const BASE_COMPUTE = { - eventNames: (api: any) => + eventNames: (api: RawApiFeature) => api.event_names && api.event_names.length > 0 ? api.event_names : undefined, }; -/** - * Declarative feature transformer - replaces 79 lines with 40 lines of config - */ -export const featureTransformer = createTransformer({ +export const featureTransformer = createTransformer({ discriminator: "type", cases: { // Boolean features: just copy base fields, no consumable @@ -32,14 +44,22 @@ export const featureTransformer = createTransformer({ }, }, - // Credit system features: always consumable credit_system: { copy: ["id", "name", "archived"], compute: { ...BASE_COMPUTE, type: () => "credit_system" as const, consumable: () => true, - creditSchema: mapCreditSchema, + creditSchema: (api) => mapCreditSchema(api), + }, + }, + + ai_credit_system: { + copy: ["id", "name", "archived"], + compute: { + ...BASE_COMPUTE, + type: () => "ai_credit_system" as const, + modelMarkups: (api) => mapModelMarkups(api), }, }, @@ -85,6 +105,6 @@ export const featureTransformer = createTransformer({ }, }); -export function transformApiFeature(apiFeature: any): Feature { +export function transformApiFeature(apiFeature: RawApiFeature): Feature { return featureTransformer.transform(apiFeature); } diff --git a/packages/atmn/src/lib/transforms/sdkToApi/feature.ts b/packages/atmn/src/lib/transforms/sdkToApi/feature.ts index aa3b95f86..a67702701 100644 --- a/packages/atmn/src/lib/transforms/sdkToApi/feature.ts +++ b/packages/atmn/src/lib/transforms/sdkToApi/feature.ts @@ -11,6 +11,11 @@ export interface ApiFeatureParams { metered_feature_id: string; credit_cost: number; }>; + model_markups?: Record; } export function transformFeatureToApi(feature: Feature): ApiFeatureParams { @@ -39,5 +44,18 @@ export function transformFeatureToApi(feature: Feature): ApiFeatureParams { })); } + if (feature.type === "ai_credit_system" && feature.modelMarkups) { + base.model_markups = Object.fromEntries( + Object.entries(feature.modelMarkups).map(([modelId, entry]) => [ + modelId, + { + markup: entry.markup, + input_cost: entry.inputCost, + output_cost: entry.outputCost, + }, + ]) + ); + } + return base; } diff --git a/packages/atmn/src/lib/transforms/sdkToCode/feature.ts b/packages/atmn/src/lib/transforms/sdkToCode/feature.ts index 3e61d4e1a..54a728f17 100644 --- a/packages/atmn/src/lib/transforms/sdkToCode/feature.ts +++ b/packages/atmn/src/lib/transforms/sdkToCode/feature.ts @@ -42,6 +42,11 @@ export function buildFeatureCode(feature: Feature, varNameOverride?: string): st lines.push(`\tcreditSchema: ${formatValue(feature.creditSchema)},`); } + // Add modelMarkups for ai_credit_system features + if (feature.type === "ai_credit_system" && feature.modelMarkups) { + lines.push(`\tmodelMarkups: ${formatValue(feature.modelMarkups)},`); + } + lines.push(`});`); return lines.join("\n"); diff --git a/server/src/internal/balances/balancesRouter.ts b/server/src/internal/balances/balancesRouter.ts index 56f096be8..95d6c54ce 100644 --- a/server/src/internal/balances/balancesRouter.ts +++ b/server/src/internal/balances/balancesRouter.ts @@ -7,6 +7,7 @@ import { handleFinalizeLock } from "./handlers/handleFinalizeLock.js"; import { handleListBalances } from "./handlers/handleListBalances.js"; import { handleSetUsage } from "./handlers/handleSetUsage.js"; import { handleTrack } from "./handlers/handleTrack.js"; +import { handleTrackTokens } from "./handlers/handleTrackTokens.js"; import { handleUpdateBalance } from "./handlers/handleUpdateBalance.js"; // Create a Hono app for products @@ -19,6 +20,7 @@ balancesRouter.post("/balances/update", ...handleUpdateBalance); // Track balancesRouter.post("/events", ...handleTrack); balancesRouter.post("/track", ...handleTrack); +balancesRouter.post("/trackTokens", ...handleTrackTokens); // Check balancesRouter.post("/entitled", ...handleCheck); @@ -33,5 +35,6 @@ balancesRpcRouter.post("/balances.update", ...handleUpdateBalance); balancesRpcRouter.post("/balances.delete", ...handleDeleteBalance); balancesRpcRouter.post("/balances.track", ...handleTrack); +balancesRpcRouter.post("/balances.trackTokens", ...handleTrackTokens); balancesRpcRouter.post("/balances.check", ...handleCheck); balancesRpcRouter.post("/balances.finalize", ...handleFinalizeLock); diff --git a/server/src/internal/balances/handlers/handleTrackTokens.ts b/server/src/internal/balances/handlers/handleTrackTokens.ts new file mode 100644 index 000000000..b0f6df62f --- /dev/null +++ b/server/src/internal/balances/handlers/handleTrackTokens.ts @@ -0,0 +1,129 @@ +import { createRoute } from "@/honoMiddlewares/routeHandler.js"; +import { runTrackV2 } from "@/internal/balances/track/runTrackV2.js"; +import { getCreditCost } from "@/internal/features/creditSystemUtils.js"; +import { + AffectedResource, + ErrCode, + type Feature, + FeatureType, + RecaseError, + Scopes, + type TrackParams, + TrackTokensParamsSchema, +} from "@autumn/shared"; +import type { FeatureDeduction } from "../utils/types/featureDeduction.js"; + +const resolveAiCreditFeature = ({ + features, + featureId, +}: { + features: Feature[]; + featureId?: string; +}): Feature => { + if (featureId) { + const candidate = features.find((f) => f.id === featureId); + if (!candidate) { + throw new RecaseError({ + message: `Feature ${featureId} not found`, + code: ErrCode.FeatureNotFound, + statusCode: 404, + }); + } + if (candidate.type !== FeatureType.AiCreditSystem) { + throw new RecaseError({ + message: `Feature ${featureId} is not an AI credit system feature`, + code: ErrCode.InvalidRequest, + statusCode: 400, + }); + } + return candidate; + } + + const matches = features.filter((f) => f.type === FeatureType.AiCreditSystem); + if (matches.length === 0) { + throw new RecaseError({ + message: "No AI credit system feature found for this organization", + code: ErrCode.FeatureNotFound, + statusCode: 404, + }); + } + if (matches.length > 1) { + throw new RecaseError({ + message: + "Multiple AI credit system features found. Please specify a feature_id to disambiguate.", + code: ErrCode.InvalidRequest, + statusCode: 400, + }); + } + return matches[0]; +}; + +export const handleTrackTokens = createRoute({ + scopes: [Scopes.Balances.Write], + body: TrackTokensParamsSchema, + resource: AffectedResource.Track, + handler: async (c) => { + const body = c.req.valid("json"); + const ctx = c.get("ctx"); + + const aiCreditFeature = resolveAiCreditFeature({ + features: ctx.features, + featureId: body.feature_id, + }); + + const rawModelName = body.model_id; + + // Compute the dollar cost once and reuse it for both the response/event row + // and the deduction layer (avoids a second getCreditCost call per entitlement). + const cost = await getCreditCost({ + featureId: aiCreditFeature.id, + creditSystem: aiCreditFeature, + modelName: rawModelName, + tokens: { + input: body.input_tokens, + output: body.output_tokens, + }, + }); + + const featureDeductions: FeatureDeduction[] = [ + { + feature: aiCreditFeature, + deduction: 1, // multiplied by per-entitlement credit_cost in the deduction layer (Postgres) + tokenUsage: { + modelName: rawModelName, + inputTokens: body.input_tokens, + outputTokens: body.output_tokens, + }, + precomputedCreditCost: cost, + }, + ]; + + // Build TrackParams body — store model/tokens in properties for audit + const trackBody: TrackParams = { + customer_id: body.customer_id, + entity_id: body.entity_id, + feature_id: aiCreditFeature.id, + value: cost, + properties: { + ...body.properties, + model: rawModelName, + input_tokens: body.input_tokens, + output_tokens: body.output_tokens, + cost, + }, + idempotency_key: body.idempotency_key, + overage_behavior: body.overage_behavior, + customer_data: body.customer_data, + entity_data: body.entity_data, + skip_event: body.skip_event, + }; + + return c.json( + await runTrackV2({ + ctx, + body: trackBody, + featureDeductions, + }), + ); + }, +}); diff --git a/server/src/internal/balances/utils/deduction/executePostgresDeduction.ts b/server/src/internal/balances/utils/deduction/executePostgresDeduction.ts index d0f3a84ac..05c015f1e 100644 --- a/server/src/internal/balances/utils/deduction/executePostgresDeduction.ts +++ b/server/src/internal/balances/utils/deduction/executePostgresDeduction.ts @@ -111,7 +111,7 @@ export const executePostgresDeduction = async ({ customerEntitlements, unlimitedFeatureIds, lock: preparedLock, - } = prepareFeatureDeduction({ + } = await prepareFeatureDeduction({ ctx, fullCustomer, deduction, diff --git a/server/src/internal/balances/utils/deduction/executeRedisDeduction.ts b/server/src/internal/balances/utils/deduction/executeRedisDeduction.ts index 9ac4e797a..43a5040b7 100644 --- a/server/src/internal/balances/utils/deduction/executeRedisDeduction.ts +++ b/server/src/internal/balances/utils/deduction/executeRedisDeduction.ts @@ -109,7 +109,7 @@ export const executeRedisDeduction = async ({ customerEntitlements, unlimitedFeatureIds, lock: preparedLock, - } = prepareFeatureDeduction({ + } = await prepareFeatureDeduction({ ctx, fullCustomer, deduction, diff --git a/server/src/internal/balances/utils/deduction/prepareFeatureDeduction.ts b/server/src/internal/balances/utils/deduction/prepareFeatureDeduction.ts index f89bb8c3a..1d4cba8f3 100644 --- a/server/src/internal/balances/utils/deduction/prepareFeatureDeduction.ts +++ b/server/src/internal/balances/utils/deduction/prepareFeatureDeduction.ts @@ -27,7 +27,7 @@ import type { FeatureDeduction } from "../types/featureDeduction.js"; * Prepares all the inputs needed to execute a deduction for a single feature. * Shared by both Redis (Lua) and Postgres (SQL) deduction paths. */ -export const prepareFeatureDeduction = ({ +export const prepareFeatureDeduction = async ({ ctx, fullCustomer, deduction, @@ -37,7 +37,7 @@ export const prepareFeatureDeduction = ({ fullCustomer: FullCustomer; deduction: FeatureDeduction; options?: DeductionOptions; -}): PreparedFeatureDeduction => { +}): Promise => { const { org } = ctx; const { env } = ctx; const { feature, lock, targetBalance } = deduction; @@ -101,14 +101,31 @@ export const prepareFeatureDeduction = ({ .map((ce) => ce.entitlement.feature.id), ); + // Compute credit cost once per customer entitlement + const creditCostByEntitlementId = new Map(); + await Promise.all( + cusEnts.map(async (ce) => { + const creditCost = + deduction.precomputedCreditCost ?? + (await getCreditCost({ + featureId: feature.id, + creditSystem: ce.entitlement.feature, + modelName: deduction.tokenUsage?.modelName, + tokens: deduction.tokenUsage + ? { + input: deduction.tokenUsage.inputTokens, + output: deduction.tokenUsage.outputTokens, + } + : undefined, + })); + creditCostByEntitlementId.set(ce.id, creditCost); + }), + ); + // Build input for each customer entitlement const customerEntitlementDeductions: CustomerEntitlementDeduction[] = cusEnts.map((ce) => { - const creditCost = getCreditCost({ - featureId: feature.id, - creditSystem: ce.entitlement.feature, - }); - + const creditCost = creditCostByEntitlementId.get(ce.id)!; const maxOverage = getMaxOverage({ cusEnt: ce }); const isFreeAllocated = @@ -148,10 +165,7 @@ export const prepareFeatureDeduction = ({ // Collect and sort rollovers by expires_at (oldest first), including credit_cost from parent entitlement const sortedRollovers = cusEnts .flatMap((ce) => { - const creditCost = getCreditCost({ - featureId: feature.id, - creditSystem: ce.entitlement.feature, - }); + const creditCost = creditCostByEntitlementId.get(ce.id)!; return (ce.rollovers || []).map((r) => ({ ...r, credit_cost: creditCost, diff --git a/server/src/internal/balances/utils/deductionV2/executePostgresDeductionV2.ts b/server/src/internal/balances/utils/deductionV2/executePostgresDeductionV2.ts index 9f025e7fb..1b6d5ca84 100644 --- a/server/src/internal/balances/utils/deductionV2/executePostgresDeductionV2.ts +++ b/server/src/internal/balances/utils/deductionV2/executePostgresDeductionV2.ts @@ -111,7 +111,7 @@ export const executePostgresDeductionV2 = async ({ unlimitedFeatureIds, unlimitedCusEnt, lock: preparedLock, - } = prepareFeatureDeductionV2({ + } = await prepareFeatureDeductionV2({ ctx, fullSubject, deduction, diff --git a/server/src/internal/balances/utils/deductionV2/executeRedisDeductionV2.ts b/server/src/internal/balances/utils/deductionV2/executeRedisDeductionV2.ts index a9501a6e4..bfd9cdcb2 100644 --- a/server/src/internal/balances/utils/deductionV2/executeRedisDeductionV2.ts +++ b/server/src/internal/balances/utils/deductionV2/executeRedisDeductionV2.ts @@ -122,7 +122,7 @@ export const executeRedisDeductionV2 = async ({ unlimitedFeatureIds, unlimitedCusEnt, lock: preparedLock, - } = prepareFeatureDeductionV2({ + } = await prepareFeatureDeductionV2({ ctx, fullSubject, deduction, diff --git a/server/src/internal/balances/utils/deductionV2/prepareFeatureDeductionV2.ts b/server/src/internal/balances/utils/deductionV2/prepareFeatureDeductionV2.ts index 4b550de10..44147d7ba 100644 --- a/server/src/internal/balances/utils/deductionV2/prepareFeatureDeductionV2.ts +++ b/server/src/internal/balances/utils/deductionV2/prepareFeatureDeductionV2.ts @@ -29,7 +29,7 @@ import type { FeatureDeduction } from "../types/featureDeduction.js"; * Prepares all the inputs needed to execute a deduction for a single feature. * Mirrors the legacy helper, but reads from FullSubject. */ -export const prepareFeatureDeductionV2 = ({ +export const prepareFeatureDeductionV2 = async ({ ctx, fullSubject, deduction, @@ -39,7 +39,7 @@ export const prepareFeatureDeductionV2 = ({ fullSubject: FullSubject; deduction: FeatureDeduction; options?: DeductionOptions; -}): PreparedFeatureDeduction => { +}): Promise => { const { org, env } = ctx; const { feature, lock, targetBalance } = deduction; const { overageBehaviour = "cap", customerEntitlementFilters } = options; @@ -115,12 +115,31 @@ export const prepareFeatureDeductionV2 = ({ .map((customerEntitlement) => customerEntitlement.entitlement.feature.id), ); + const creditCostByCustomerEntitlementId = new Map(); + await Promise.all( + customerEntitlements.map(async (customerEntitlement) => { + const creditCost = + deduction.precomputedCreditCost ?? + (await getCreditCost({ + featureId: feature.id, + creditSystem: customerEntitlement.entitlement.feature, + modelName: deduction.tokenUsage?.modelName, + tokens: deduction.tokenUsage + ? { + input: deduction.tokenUsage.inputTokens, + output: deduction.tokenUsage.outputTokens, + } + : undefined, + })); + creditCostByCustomerEntitlementId.set(customerEntitlement.id, creditCost); + }), + ); + const customerEntitlementDeductions: CustomerEntitlementDeduction[] = customerEntitlements.map((customerEntitlement) => { - const creditCost = getCreditCost({ - featureId: feature.id, - creditSystem: customerEntitlement.entitlement.feature, - }); + const creditCost = creditCostByCustomerEntitlementId.get( + customerEntitlement.id, + )!; const maxOverage = getMaxOverage({ cusEnt: customerEntitlement, @@ -162,26 +181,24 @@ export const prepareFeatureDeductionV2 = ({ }; }); - const sortedRollovers = customerEntitlements - .flatMap((customerEntitlement) => { - const creditCost = getCreditCost({ - featureId: feature.id, - creditSystem: customerEntitlement.entitlement.feature, - }); + const rolloverArrays = customerEntitlements.map((customerEntitlement) => { + const creditCost = creditCostByCustomerEntitlementId.get( + customerEntitlement.id, + )!; + return (customerEntitlement.rollovers || []).map((rollover) => ({ + ...rollover, + credit_cost: creditCost, + })); + }); - return (customerEntitlement.rollovers || []).map((rollover) => ({ - ...rollover, - credit_cost: creditCost, - })); - }) - .sort((left, right) => { - if (left.expires_at && right.expires_at) { - return left.expires_at - right.expires_at; - } - if (left.expires_at && !right.expires_at) return -1; - if (!left.expires_at && right.expires_at) return 1; - return 0; - }); + const sortedRollovers = rolloverArrays.flat().sort((left, right) => { + if (left.expires_at && right.expires_at) { + return left.expires_at - right.expires_at; + } + if (left.expires_at && !right.expires_at) return -1; + if (!left.expires_at && right.expires_at) return 1; + return 0; + }); const oneDaySeconds = 24 * 60 * 60; const oneHourSeconds = 60 * 60; diff --git a/server/src/internal/balances/utils/types/featureDeduction.ts b/server/src/internal/balances/utils/types/featureDeduction.ts index b440158d1..eb6d3b0eb 100644 --- a/server/src/internal/balances/utils/types/featureDeduction.ts +++ b/server/src/internal/balances/utils/types/featureDeduction.ts @@ -4,8 +4,17 @@ export type FeatureDeduction = { feature: Feature; deduction: number; targetBalance?: number; - lock?: LockParams; + tokenUsage?: { + modelName: string; + inputTokens: number; + outputTokens: number; + }; + + /** Pre-computed dollar cost; if set, the deduction layer skips its own getCreditCost call. */ + precomputedCreditCost?: number; + + lock?: LockParams; lockReceipt?: LockReceipt; lockReceiptKey?: string; unwindValue?: number; diff --git a/server/src/internal/features/creditSystemUtils.ts b/server/src/internal/features/creditSystemUtils.ts index 12a0363ad..367af65e6 100644 --- a/server/src/internal/features/creditSystemUtils.ts +++ b/server/src/internal/features/creditSystemUtils.ts @@ -1,9 +1,13 @@ import { type CreditSchemaItem, + ErrCode, type Feature, FeatureType, + InternalError, + RecaseError, } from "@autumn/shared"; import { Decimal } from "decimal.js"; +import { getModelsDevPricing } from "@/internal/features/utils/getModelPricing"; const creditSystemContainsFeature = ({ creditSystem, @@ -15,7 +19,8 @@ const creditSystemContainsFeature = ({ if (creditSystem.type !== FeatureType.CreditSystem) { return false; } - const schema: CreditSchemaItem[] = creditSystem.config.schema; + const schema: CreditSchemaItem[] | undefined = creditSystem.config?.schema; + if (!schema) return false; for (const schemaItem of schema) { if (schemaItem.metered_feature_id === meteredFeatureId) { @@ -69,21 +74,127 @@ export const featureToCreditSystem = ({ return amount; }; +// Costs are in $/M tokens; markup is a percentage (e.g. 20 = +20%). +const computeMarkedUpCost = ({ + inputCostPerMillion, + outputCostPerMillion, + input, + output, + markup, +}: { + inputCostPerMillion: Decimal.Value; + outputCostPerMillion: Decimal.Value; + input: number; + output: number; + markup: number; +}) => + new Decimal(inputCostPerMillion) + .mul(input) + .add(new Decimal(outputCostPerMillion).mul(output)) + .div(1_000_000) + .mul(new Decimal(1).add(new Decimal(markup).div(100))) + .toNumber(); -export const getCreditCost = ({ +const getModelCreditCost = async ({ + modelName, + creditSystem, + input, + output, +}: { + modelName: string; + creditSystem: Feature; + input: number; + output: number; +}) => { + const markups = creditSystem.model_markups || {}; + const markupEntry = markups[modelName]; + const { markup } = markupEntry ?? { markup: 0 }; + + if (modelName.startsWith("custom/")) { + if (markupEntry?.input_cost == null || markupEntry?.output_cost == null) { + throw new RecaseError({ + message: `Custom model ${modelName} is missing input_cost or output_cost in model_markups`, + code: ErrCode.InvalidRequest, + data: { modelName }, + }); + } + return computeMarkedUpCost({ + inputCostPerMillion: markupEntry.input_cost, + outputCostPerMillion: markupEntry.output_cost, + input, + output, + markup, + }); + } + + const pricingData = await getModelsDevPricing(); + if (!pricingData) { + throw new InternalError({ + message: "Failed to fetch models.dev pricing data", + code: ErrCode.InternalError, + }); + } + + const [providerKey, ...modelParts] = modelName.split("/"); + const modelKey = modelParts.join("/"); + const model = pricingData[providerKey]?.models[modelKey]; + + if (!model) { + throw new RecaseError({ + message: `Model ${modelName} not found in models.dev pricing data ${providerKey} provider config.`, + code: ErrCode.InvalidRequest, + statusCode: 400, + data: { modelName }, + }); + } + + return computeMarkedUpCost({ + inputCostPerMillion: model.cost.input, + outputCostPerMillion: model.cost.output, + input, + output, + markup, + }); +}; + +export const getCreditCost = async ({ featureId, creditSystem, amount = 1, + tokens, + modelName, }: { featureId: string; creditSystem: Feature; amount?: number; + modelName?: string; + tokens?: { + input: number; + output: number; + }; }) => { - if (creditSystem.type !== FeatureType.CreditSystem) { + if (creditSystem.type !== FeatureType.CreditSystem && creditSystem.type !== FeatureType.AiCreditSystem) { + return amount; + } + if (creditSystem.type === FeatureType.AiCreditSystem) { + if (!tokens || !modelName) { + throw new RecaseError({ + message: "modelName and tokens must be provided for AI credit systems", + code: ErrCode.InvalidRequest, + statusCode: 400, + }); + } + return await getModelCreditCost({ + modelName, + creditSystem, + ...tokens, + }); + } + // If tracking the credit system feature itself, 1:1 mapping + if (featureId === creditSystem.id) { return amount; } const schema: CreditSchemaItem[] = creditSystem.config.schema; - for (const schemaItem of schema) { if (schemaItem.metered_feature_id === featureId) { return new Decimal(schemaItem.credit_amount) @@ -93,5 +204,10 @@ export const getCreditCost = ({ } } - return 1; + throw new RecaseError({ + message: "Feature is not included in credit system schema", + code: ErrCode.InvalidRequest, + statusCode: 400, + data: { featureId, creditSystemId: creditSystem.id }, + }); }; diff --git a/server/src/internal/features/featureActions/createFeature.ts b/server/src/internal/features/featureActions/createFeature.ts index d2622976f..1a1bbc1b2 100644 --- a/server/src/internal/features/featureActions/createFeature.ts +++ b/server/src/internal/features/featureActions/createFeature.ts @@ -9,15 +9,15 @@ import { } from "../featureUtils.js"; const validateFeature = (data: any) => { - const featureType = data.type; - - // validateFeatureId(data.id); + const featureType = data.type as FeatureType; let config = data.config; if (featureType === FeatureType.Metered) { config = validateMeteredConfig(config); } else if (featureType === FeatureType.CreditSystem) { - config = validateCreditSystem(config); + config = validateCreditSystem(config, { isAiCreditSystem: false }); + } else if (featureType === FeatureType.AiCreditSystem) { + config = validateCreditSystem(config, { isAiCreditSystem: true }); } const parsedFeature = CreateFeatureSchema.parse({ ...data, config }); @@ -32,6 +32,14 @@ interface CreateFeatureParams { type: string; config?: any; event_names?: string[]; + model_markups?: Record< + string, + { + markup: number; + input_cost?: number; + output_cost?: number; + } + > | null; }; skipGenerateDisplay?: boolean; } @@ -54,6 +62,7 @@ export const createFeature = async ({ created_at: Date.now(), env: ctx.env, ...parsedFeature, + model_markups: data.model_markups ?? null, }; const insertedData = await FeatureService.insert({ diff --git a/server/src/internal/features/featureActions/updateFeature.ts b/server/src/internal/features/featureActions/updateFeature.ts index 1321b5140..514d9b164 100644 --- a/server/src/internal/features/featureActions/updateFeature.ts +++ b/server/src/internal/features/featureActions/updateFeature.ts @@ -3,6 +3,7 @@ import { ErrCode, type Feature, FeatureType, + type ModelMarkups, notNullish, } from "@autumn/shared"; import type { AutumnContext } from "@/honoUtils/HonoEnv.js"; @@ -31,6 +32,34 @@ interface UpdateFeatureParams { * Checks if the credit schema has changed between old and new config. * Returns true if schema changed (different items or different credit amounts). */ +const areModelMarkupsEqual = ({ + a, + b, +}: { + a: ModelMarkups; + b: ModelMarkups; +}): boolean => { + const aIsAbsent = a == null; + const bIsAbsent = b == null; + if (aIsAbsent && bIsAbsent) return true; + if (aIsAbsent || bIsAbsent) return false; + + const aKeys = Object.keys(a); + const bKeys = Object.keys(b); + if (aKeys.length !== bKeys.length) return false; + + for (const key of aKeys) { + const aEntry = a[key]; + const bEntry = b[key]; + if (!bEntry) return false; + if (aEntry.markup !== bEntry.markup) return false; + if (aEntry.input_cost !== bEntry.input_cost) return false; + if (aEntry.output_cost !== bEntry.output_cost) return false; + } + + return true; +}; + const hasCreditSchemaChanged = ({ oldSchema, newSchema, @@ -146,12 +175,14 @@ export const updateFeature = async ({ } } - // Validate config based on feature type + const effectiveType = updates.type ?? feature.type; + const isAiCreditSystem = effectiveType === FeatureType.AiCreditSystem; + const newConfig = updates.config !== undefined - ? feature.type === FeatureType.CreditSystem - ? validateCreditSystem(updates.config) - : feature.type === FeatureType.Metered + ? effectiveType === FeatureType.CreditSystem || isAiCreditSystem + ? validateCreditSystem(updates.config, { isAiCreditSystem }) + : effectiveType === FeatureType.Metered ? validateMeteredConfig(updates.config) : updates.config : feature.config; @@ -165,10 +196,11 @@ export const updateFeature = async ({ updates: { id: updates.id, name: updates.name, - type: updates.type, + type: effectiveType, archived: updates.archived, event_names: updates.event_names, config: newConfig, + model_markups: updates.model_markups, }, }); @@ -181,18 +213,23 @@ export const updateFeature = async ({ }); } - // Queue cache clear for credit system if schema changed - if ( - feature.type === FeatureType.CreditSystem && - updates.config?.schema && - updatedFeature - ) { - const schemaChanged = hasCreditSchemaChanged({ - oldSchema: feature.config?.schema, - newSchema: updates.config.schema, - }); + // Queue cache clear for credit system if schema or model markups changed + if ((feature.type === FeatureType.CreditSystem || feature.type === FeatureType.AiCreditSystem) && updatedFeature) { + const schemaChanged = + updates.config != null && + hasCreditSchemaChanged({ + oldSchema: feature.config?.schema, + newSchema: updates.config.schema, + }); - if (schemaChanged) { + const markupsChanged = + updates.model_markups !== undefined && + !areModelMarkupsEqual({ + a: updates.model_markups, + b: feature.model_markups, + }); + + if (schemaChanged || markupsChanged) { await addTaskToQueue({ jobName: JobName.ClearCreditSystemCustomerCache, payload: { diff --git a/server/src/internal/features/featureRouter.ts b/server/src/internal/features/featureRouter.ts index 75113786f..75d11b38a 100644 --- a/server/src/internal/features/featureRouter.ts +++ b/server/src/internal/features/featureRouter.ts @@ -10,10 +10,12 @@ import { handleListFeaturesV1 } from "./handlers/handleListFeatures/handleListFe import { handleUpdateFeatureV1 } from "./handlers/handleUpdateFeature/handleUpdateFeatureV1"; import { handleUpdateFeatureV2 } from "./handlers/handleUpdateFeature/handleUpdateFeatureV2"; import { handleGetFeatureDeletionInfo } from "./internalHandlers/handleGetFeatureDeletionInfo"; +import { handleGetModelPricing } from "./internalHandlers/handleGetModelPricing"; export const featureRouter = new Hono(); featureRouter.get("", ...handleListFeaturesV1); featureRouter.post("", ...handleCreateFeatureV1); +featureRouter.get("/ai/model_pricing", ...handleGetModelPricing); featureRouter.get("/:feature_id", ...handleGetFeatureV1); featureRouter.post("/:feature_id", ...handleUpdateFeatureV1); featureRouter.delete("/:feature_id", ...handleDeleteFeatureV1); diff --git a/server/src/internal/features/featureUtils.ts b/server/src/internal/features/featureUtils.ts index 02015399d..761b84e04 100644 --- a/server/src/internal/features/featureUtils.ts +++ b/server/src/internal/features/featureUtils.ts @@ -41,9 +41,13 @@ export const validateMeteredConfig = (config: MeteredConfig) => { return newConfig as MeteredConfig; }; -export const validateCreditSystem = (config: CreditSystemConfig) => { +export const validateCreditSystem = ( + config: CreditSystemConfig, + { isAiCreditSystem = false }: { isAiCreditSystem?: boolean } = {}, +) => { const schema = config.schema; - if (!schema || schema.length === 0) { + + if (!isAiCreditSystem && (!schema || schema.length === 0)) { throw new RecaseError({ message: `At least one metered feature is required for credit system`, code: ErrCode.InvalidFeature, @@ -51,11 +55,14 @@ export const validateCreditSystem = (config: CreditSystemConfig) => { }); } + if (isAiCreditSystem) { + return { ...config, usage_type: FeatureUsageType.Single }; + } + // Check if multiple of the same feature const meteredFeatureIds = schema.map( (schemaItem) => schemaItem.metered_feature_id, ); - // console.log("Metered feature ids:", meteredFeatureIds); const uniqueMeteredFeatureIds = Array.from(new Set(meteredFeatureIds)); if (meteredFeatureIds.length !== uniqueMeteredFeatureIds.length) { throw new RecaseError({ diff --git a/server/src/internal/features/handlers/handleUpdateFeature/handleUpdateFeatureV1.ts b/server/src/internal/features/handlers/handleUpdateFeature/handleUpdateFeatureV1.ts index 17fceaaea..2c4bd4263 100644 --- a/server/src/internal/features/handlers/handleUpdateFeature/handleUpdateFeatureV1.ts +++ b/server/src/internal/features/handlers/handleUpdateFeature/handleUpdateFeatureV1.ts @@ -58,6 +58,8 @@ export const handleUpdateFeatureV1 = createRoute({ archived: body.archived, event_names: body.event_names, display: body.display, + model_markups: + body.model_markups === undefined ? undefined : body.model_markups, }, }); diff --git a/server/src/internal/features/internalHandlers/handleGetModelPricing.ts b/server/src/internal/features/internalHandlers/handleGetModelPricing.ts new file mode 100644 index 000000000..91f9ba866 --- /dev/null +++ b/server/src/internal/features/internalHandlers/handleGetModelPricing.ts @@ -0,0 +1,11 @@ +import { createRoute } from "@/honoMiddlewares/routeHandler"; +import { getModelsDevPricing } from "@/internal/features/utils/getModelPricing"; +import { Scopes } from "@autumn/shared"; + +export const handleGetModelPricing = createRoute({ + scopes: [Scopes.Features.Read], + handler: async (c) => { + const data = await getModelsDevPricing(); + return c.json(data); + }, +}); diff --git a/server/src/internal/features/utils/constructFeatureUtils.ts b/server/src/internal/features/utils/constructFeatureUtils.ts index 7ffab0b81..29c712eed 100644 --- a/server/src/internal/features/utils/constructFeatureUtils.ts +++ b/server/src/internal/features/utils/constructFeatureUtils.ts @@ -4,6 +4,7 @@ import { type Feature, FeatureType, FeatureUsageType, + type ModelMarkups, } from "@autumn/shared"; import { generateId, keyToTitle } from "@server/utils/genUtils"; @@ -36,8 +37,9 @@ const constructFeature = ({ display, archived: false, event_names: [], + model_markups: null, }; - + // This function isn't used anywhere, maybe delete it? return newFeature; }; @@ -64,6 +66,7 @@ export const constructBooleanFeature = ({ config: null, archived: false, event_names: [], + model_markups: null, }; return newFeature; @@ -109,6 +112,7 @@ export const constructMeteredFeature = ({ }, archived: false, event_names: eventNames, + model_markups: null, }; return newFeature; @@ -151,6 +155,43 @@ export const constructCreditSystem = ({ config, archived: false, event_names: [], + model_markups: null, + }; + + return newFeature; +}; + +export const constructAiCreditSystem = ({ + featureId, + name, + orgId, + env, + modelMarkups, +}: { + featureId: string; + name?: string; + orgId: string; + env: AppEnv; + modelMarkups: ModelMarkups; +}) => { + const config = { + schema: [], + usage_type: FeatureUsageType.Single, + }; + + const newFeature: Feature = { + internal_id: generateId("fe"), + org_id: orgId, + env, + created_at: Date.now(), + + id: featureId, + name: name || keyToTitle(featureId), + type: FeatureType.AiCreditSystem, + config, + archived: false, + event_names: [], + model_markups: modelMarkups, }; return newFeature; diff --git a/server/src/internal/features/utils/getModelPricing.ts b/server/src/internal/features/utils/getModelPricing.ts new file mode 100644 index 000000000..ac61bb8a3 --- /dev/null +++ b/server/src/internal/features/utils/getModelPricing.ts @@ -0,0 +1,40 @@ +import { CacheManager } from "@/utils/cacheUtils/CacheManager"; +import { ErrCode, InternalError, type ModelsDevProvider } from "@autumn/shared"; + +const MODELS_DEV_CACHE_KEY = "models_dev_pricing"; + +export const getModelsDevPricing = async () => { + try { + const cached = + await CacheManager.getJson>( + MODELS_DEV_CACHE_KEY, + ); + if (cached) return cached; + const response = await fetch("https://models.dev/api.json"); + if (!response.ok) + throw new InternalError({ + message: `models.dev returned ${response.status}`, + code: ErrCode.InternalError, + }); + + const data: Record = await response.json(); + Promise.all([ + CacheManager.setJson(MODELS_DEV_CACHE_KEY, data, 60 * 60 * 3), + CacheManager.setJson( + `${MODELS_DEV_CACHE_KEY}_stale`, + data, + 60 * 60 * 24 * 3, + ), + ]).catch(() => {}); + return data; + } catch { + const stale = await CacheManager.getJson>( + `${MODELS_DEV_CACHE_KEY}_stale`, + ); + if (stale) return stale; + throw new InternalError({ + message: "Failed to fetch models.dev pricing and no cache available", + code: ErrCode.InternalError, + }); + } +}; diff --git a/server/tests/advanced/usage/usage2.test.ts b/server/tests/advanced/usage/usage2.test.ts index 886ce2e4d..75344a9dd 100644 --- a/server/tests/advanced/usage/usage2.test.ts +++ b/server/tests/advanced/usage/usage2.test.ts @@ -104,7 +104,7 @@ describe(`${chalk.yellowBright("usage2: Testing basic usage product")}`, () => { .toNumber(); const featureId = i % 2 === 0 ? TestFeature.Action1 : TestFeature.Action2; - const creditsUsed = getCreditCost({ + const creditsUsed = await getCreditCost({ creditSystem: creditsFeature, featureId: featureId, amount: randomVal, diff --git a/server/tests/advanced/usage/usage3.test.ts b/server/tests/advanced/usage/usage3.test.ts index c1474e470..6a48c6ebe 100644 --- a/server/tests/advanced/usage/usage3.test.ts +++ b/server/tests/advanced/usage/usage3.test.ts @@ -121,7 +121,7 @@ describe(`${chalk.yellowBright( .toNumber(); const featureId = i % 2 === 0 ? TestFeature.Action1 : TestFeature.Action2; - const creditsUsed = getCreditCost({ + const creditsUsed = await getCreditCost({ creditSystem: creditsFeature, featureId: featureId, amount: randomVal, diff --git a/server/tests/advanced/usage/usage4.test.ts b/server/tests/advanced/usage/usage4.test.ts index 3a3d6f1a8..7199ecc5b 100644 --- a/server/tests/advanced/usage/usage4.test.ts +++ b/server/tests/advanced/usage/usage4.test.ts @@ -105,7 +105,7 @@ describe(`${chalk.yellowBright("usage4: GPU starter annual")}`, () => { .toNumber(); const featureId = i % 2 === 0 ? TestFeature.Action1 : TestFeature.Action2; - const creditsUsed = getCreditCost({ + const creditsUsed = await getCreditCost({ creditSystem: creditsFeature, featureId: featureId, amount: randomVal, diff --git a/server/tests/balances/check/credit-systems/credit-systems1.test.ts b/server/tests/balances/check/credit-systems/credit-systems1.test.ts index 22f1a2fef..6fe182bee 100644 --- a/server/tests/balances/check/credit-systems/credit-systems1.test.ts +++ b/server/tests/balances/check/credit-systems/credit-systems1.test.ts @@ -71,7 +71,7 @@ describe(`${chalk.yellowBright("credit-systems1: test /check on action that uses required_balance: requiredActionUnits, })) as unknown as CheckResponseV2; - const creditCost = getCreditCost({ + const creditCost = await getCreditCost({ featureId: action, creditSystem: creditFeature!, amount: requiredActionUnits, @@ -172,7 +172,7 @@ describe(`${chalk.yellowBright("credit-systems1: test /check on action that uses required_balance: requiredAction1Units, })) as unknown as CheckResponseV0; - const meteredCost = getCreditCost({ + const meteredCost = await getCreditCost({ featureId: TestFeature.Action1, creditSystem: creditFeature!, amount: requiredAction1Units, @@ -196,7 +196,7 @@ describe(`${chalk.yellowBright("credit-systems1: test /check on action that uses required_balance: requiredAction2Units, })) as unknown as CheckResponseV0; - const meteredCost = getCreditCost({ + const meteredCost = await getCreditCost({ featureId: TestFeature.Action2, creditSystem: creditFeature!, amount: requiredAction2Units, diff --git a/server/tests/balances/check/send-event/send-event3.test.ts b/server/tests/balances/check/send-event/send-event3.test.ts index 866be0b0b..4d84de2ca 100644 --- a/server/tests/balances/check/send-event/send-event3.test.ts +++ b/server/tests/balances/check/send-event/send-event3.test.ts @@ -93,7 +93,7 @@ describe(`${chalk.yellowBright("send-event3: Testing check with track, credit sy send_event: true, })) as unknown as CheckResponseV2; - const creditCost = getCreditCost({ + const creditCost = await getCreditCost({ featureId: TestFeature.Action1, creditSystem: creditFeature!, amount: 10, @@ -156,7 +156,7 @@ describe(`${chalk.yellowBright("send-event3: Testing check with track, credit sy allowed: true, customer_id: customerId, feature_id: TestFeature.Credits, - required_balance: getCreditCost({ + required_balance: await getCreditCost({ featureId: TestFeature.Action1, creditSystem: creditFeature!, amount: 10, @@ -167,7 +167,7 @@ describe(`${chalk.yellowBright("send-event3: Testing check with track, credit sy test("should check with track and deduct from credits", async () => { const value = 2.5; - const creditCost = getCreditCost({ + const creditCost = await getCreditCost({ featureId: TestFeature.Action1, creditSystem: creditFeature!, amount: value, diff --git a/server/tests/balances/check/send-event/send-event4.test.ts b/server/tests/balances/check/send-event/send-event4.test.ts index aa1a91409..05ce02f20 100644 --- a/server/tests/balances/check/send-event/send-event4.test.ts +++ b/server/tests/balances/check/send-event/send-event4.test.ts @@ -82,7 +82,7 @@ describe(`${chalk.yellowBright("send-event4: Testing check with track, unlimited send_event: true, }); - const requiredBalance = getCreditCost({ + const requiredBalance = await getCreditCost({ featureId: TestFeature.Action1, creditSystem: ctx.features.find((f) => f.id === TestFeature.Credits)!, amount: 1000, diff --git a/server/tests/integration/balances/auto-topup/auto-topup-credit-systems.test.ts b/server/tests/integration/balances/auto-topup/auto-topup-credit-systems.test.ts index 85399f694..5336d752e 100644 --- a/server/tests/integration/balances/auto-topup/auto-topup-credit-systems.test.ts +++ b/server/tests/integration/balances/auto-topup/auto-topup-credit-systems.test.ts @@ -77,7 +77,7 @@ test.concurrent(`${chalk.yellowBright("auto-topup cs1: action track depletes cre // Track 845 units → 845 × 0.2 = 169 credits deducted // Balance: 200 - 169 = 31 → strictly above threshold (30) → does NOT trigger // (exact threshold uses <= in code, so landing on 30 would fire auto top-up) - const action1Cost = getCreditCost({ + const action1Cost = await getCreditCost({ featureId: TestFeature.Action1, creditSystem: creditFeature!, amount: 845, @@ -98,7 +98,7 @@ test.concurrent(`${chalk.yellowBright("auto-topup cs1: action track depletes cre // Track 10 units of action1 → 10 × 0.2 = 2 credits // Balance: 31 - 2 = 29 → 29 <= threshold → auto top-up fires → 29 + 100 = 129 - const action1CostSmall = getCreditCost({ + const action1CostSmall = await getCreditCost({ featureId: TestFeature.Action1, creditSystem: creditFeature!, amount: 10, @@ -169,7 +169,7 @@ test.concurrent(`${chalk.yellowBright("auto-topup cs2: action track depletes cre // Action1 costs 0.2 credits per unit // Track 900 units of action1 → 900 × 0.2 = 180 credits deducted // Balance: 200 - 180 = 20 → auto top-up fires - const action1Cost = getCreditCost({ + const action1Cost = await getCreditCost({ featureId: TestFeature.Action1, creditSystem: creditFeature!, amount: 900, diff --git a/server/tests/integration/balances/check/spend-limit/check-entity-product-spend-limit.test.ts b/server/tests/integration/balances/check/spend-limit/check-entity-product-spend-limit.test.ts index 94487bbe0..988e87ed8 100644 --- a/server/tests/integration/balances/check/spend-limit/check-entity-product-spend-limit.test.ts +++ b/server/tests/integration/balances/check/spend-limit/check-entity-product-spend-limit.test.ts @@ -362,7 +362,7 @@ test.concurrent(`${chalk.yellowBright("check-entity-product-spend-limit5: credit const creditsFeature = ctx.features.find( (f) => f.id === TestFeature.Credits, )!; - const action1CreditCost = getCreditCost({ + const action1CreditCost = await getCreditCost({ featureId: TestFeature.Action1, creditSystem: creditsFeature, amount: 1, diff --git a/server/tests/integration/balances/check/spend-limit/check-per-entity-spend-limit.test.ts b/server/tests/integration/balances/check/spend-limit/check-per-entity-spend-limit.test.ts index 83d0f53c8..dab02bfdf 100644 --- a/server/tests/integration/balances/check/spend-limit/check-per-entity-spend-limit.test.ts +++ b/server/tests/integration/balances/check/spend-limit/check-per-entity-spend-limit.test.ts @@ -210,7 +210,7 @@ test.concurrent(`${chalk.yellowBright("check-per-entity-spend-limit4: credit-sys const creditsFeature = ctx.features.find( (f) => f.id === TestFeature.Credits, )!; - const action1CreditCost = getCreditCost({ + const action1CreditCost = await getCreditCost({ featureId: TestFeature.Action1, creditSystem: creditsFeature, amount: 1, diff --git a/server/tests/integration/balances/lock/check-with-lock-credit-system.test.ts b/server/tests/integration/balances/lock/check-with-lock-credit-system.test.ts index 3c4598022..b60c09c95 100644 --- a/server/tests/integration/balances/lock/check-with-lock-credit-system.test.ts +++ b/server/tests/integration/balances/lock/check-with-lock-credit-system.test.ts @@ -131,12 +131,12 @@ test.concurrent(`${chalk.yellowBright("lock-credit CS-2: cross-boundary lock=8 c }); const creditFeature = features.find((f) => f.id === TestFeature.Credits)!; - const lockCreditCost = getCreditCost({ + const lockCreditCost = await getCreditCost({ featureId: TestFeature.Action1, creditSystem: creditFeature, amount: 3, // overflow during lock: 8 - 5 remaining = 3 }); - const extraCreditCost = getCreditCost({ + const extraCreditCost = await getCreditCost({ featureId: TestFeature.Action1, creditSystem: creditFeature, amount: 4, // confirm delta: 12 - 8 = 4 more units @@ -289,7 +289,7 @@ test.concurrent(`${chalk.yellowBright("lock-credit CS-4: lock within action1, co // Lock deducted 10 from action1 (→90). Confirm delta=+105: // exhaust remaining 90 from action1 (→0), then 15 overflow → 15×0.2=3 credits. const creditFeature = features.find((f) => f.id === TestFeature.Credits)!; - const overflowCreditCost = getCreditCost({ + const overflowCreditCost = await getCreditCost({ featureId: TestFeature.Action1, creditSystem: creditFeature, amount: 15, @@ -628,12 +628,12 @@ test.concurrent(`${chalk.yellowBright("lock-credit CS-9: cross-boundary lock=8 c // Lock deducted: 5 from action1 + 3 overflow (0.6 credits). // Confirm delta = 20 - 8 = 12 more units, action1 is already 0, all go to credits: 12×0.2=2.4. const creditFeature = features.find((f) => f.id === TestFeature.Credits)!; - const lockOverflowCost = getCreditCost({ + const lockOverflowCost = await getCreditCost({ featureId: TestFeature.Action1, creditSystem: creditFeature, amount: 3, }); - const confirmExtraCost = getCreditCost({ + const confirmExtraCost = await getCreditCost({ featureId: TestFeature.Action1, creditSystem: creditFeature, amount: 12, @@ -721,7 +721,7 @@ test.concurrent(`${chalk.yellowBright("lock-credit CS-10: confirm no override_va // Balances unchanged from what the lock left const creditFeature = features.find((f) => f.id === TestFeature.Credits)!; - const lockOverflowCost = getCreditCost({ + const lockOverflowCost = await getCreditCost({ featureId: TestFeature.Action1, creditSystem: creditFeature, amount: 3, // overflow during lock: 8 - 5 remaining = 3 diff --git a/server/tests/integration/balances/track/basic/track-credit-system.test.ts b/server/tests/integration/balances/track/basic/track-credit-system.test.ts index 7cd345c7d..a84933213 100644 --- a/server/tests/integration/balances/track/basic/track-credit-system.test.ts +++ b/server/tests/integration/balances/track/basic/track-credit-system.test.ts @@ -88,7 +88,7 @@ test.concurrent(`${chalk.yellowBright("track-credit-system2: track metered featu expect(customerBefore.features[TestFeature.Credits].balance).toBe(200); const action1Value = 50.25; - const expectedAction1CreditCost = getCreditCost({ + const expectedAction1CreditCost = await getCreditCost({ featureId: TestFeature.Action1, creditSystem: creditFeature!, amount: action1Value, @@ -108,7 +108,7 @@ test.concurrent(`${chalk.yellowBright("track-credit-system2: track metered featu }); const action2Value = 33.67; - const expectedAction2CreditCost = getCreditCost({ + const expectedAction2CreditCost = await getCreditCost({ featureId: TestFeature.Action2, creditSystem: creditFeature!, amount: action2Value, @@ -209,7 +209,7 @@ test.concurrent(`${chalk.yellowBright("track-credit-system3: test deduction orde const deduct2 = 80; const remainingAction1 = 100 - deduct1; const overflowAmount = deduct2 - remainingAction1; - const creditCostForOverflow = getCreditCost({ + const creditCostForOverflow = await getCreditCost({ featureId: TestFeature.Action1, creditSystem: creditFeature!, amount: overflowAmount, @@ -244,7 +244,7 @@ test.concurrent(`${chalk.yellowBright("track-credit-system3: test deduction orde const creditsBefore = customer2.features[TestFeature.Credits].balance; const deduct3 = 50.75; - const creditCost3 = getCreditCost({ + const creditCost3 = await getCreditCost({ featureId: TestFeature.Action1, creditSystem: creditFeature!, amount: deduct3, @@ -367,13 +367,13 @@ test.concurrent(`${chalk.yellowBright("track-credit-system4: test deduction with const overflowAction1 = deduct2 - remainingAction1; const overflowAction3 = deduct2 - remainingAction3; - const creditCostAction1 = getCreditCost({ + const creditCostAction1 = await getCreditCost({ featureId: TestFeature.Action1, creditSystem: creditFeature!, amount: overflowAction1, }); - const creditCostAction3 = getCreditCost({ + const creditCostAction3 = await getCreditCost({ featureId: TestFeature.Action3, creditSystem: credit2Feature!, amount: overflowAction3, @@ -419,13 +419,13 @@ test.concurrent(`${chalk.yellowBright("track-credit-system4: test deduction with const deduct3 = 40.25; - const creditCostAction1Deduct3 = getCreditCost({ + const creditCostAction1Deduct3 = await getCreditCost({ featureId: TestFeature.Action1, creditSystem: creditFeature!, amount: deduct3, }); - const creditCostAction3Deduct3 = getCreditCost({ + const creditCostAction3Deduct3 = await getCreditCost({ featureId: TestFeature.Action3, creditSystem: credit2Feature!, amount: deduct3, @@ -556,7 +556,7 @@ test.concurrent(`${chalk.yellowBright("track-credit-system5: test deduction orde const deduct2 = 80; const remainingAction1 = 100 - deduct1; const overflowAmount = deduct2 - remainingAction1; - const creditCostForOverflow = getCreditCost({ + const creditCostForOverflow = await getCreditCost({ featureId: TestFeature.Action1, creditSystem: creditFeature!, amount: overflowAmount, @@ -596,7 +596,7 @@ test.concurrent(`${chalk.yellowBright("track-credit-system5: test deduction orde const creditsBefore = customer2.features[TestFeature.Credits].balance; const deduct3 = 50.75; - const creditCost3 = getCreditCost({ + const creditCost3 = await getCreditCost({ featureId: TestFeature.Action1, creditSystem: creditFeature!, amount: deduct3, diff --git a/server/tests/integration/balances/track/basic/track-deductions.test.ts b/server/tests/integration/balances/track/basic/track-deductions.test.ts index cdf3bd0b0..47dc727f2 100644 --- a/server/tests/integration/balances/track/basic/track-deductions.test.ts +++ b/server/tests/integration/balances/track/basic/track-deductions.test.ts @@ -304,7 +304,7 @@ test.concurrent( }); const overflowAmount = 50; - const expectedCreditCost = getCreditCost({ + const expectedCreditCost = await getCreditCost({ featureId: TestFeature.Action1, creditSystem: creditFeature!, amount: overflowAmount, diff --git a/server/tests/integration/balances/track/basic/track-tokens.test.ts b/server/tests/integration/balances/track/basic/track-tokens.test.ts new file mode 100644 index 000000000..fe7b3bc8e --- /dev/null +++ b/server/tests/integration/balances/track/basic/track-tokens.test.ts @@ -0,0 +1,349 @@ +import { expect, test } from "bun:test"; + +import type { ApiCustomerV3, TrackResponseV2 } from "@autumn/shared"; +import { TestFeature } from "@tests/setup/v2Features.js"; +import { items } from "@tests/utils/fixtures/items.js"; +import { products } from "@tests/utils/fixtures/products.js"; +import { initScenario, s } from "@tests/utils/testInitUtils/initScenario.js"; +import chalk from "chalk"; +import { Decimal } from "decimal.js"; +import { getCreditCost } from "@/internal/features/creditSystemUtils.js"; + +// ═══════════════════════════════════════════════════════════════════ +// TRACK-TOKENS-1: Basic trackTokens with models.dev pricing +// ═══════════════════════════════════════════════════════════════════ + +test.concurrent(`${chalk.yellowBright("track-tokens-1: basic trackTokens with models.dev pricing")}`, async () => { + const aiCreditsItem = items.free({ + featureId: TestFeature.AiCredits, + includedUsage: 1000, + }); + const freeProd = products.base({ + id: "free", + items: [aiCreditsItem], + }); + + const { customerId, autumnV1, autumnV2, ctx } = await initScenario({ + customerId: "track-tokens-1", + setup: [s.customer({ testClock: false }), s.products({ list: [freeProd] })], + actions: [s.attach({ productId: freeProd.id })], + }); + + const aiCreditFeature = ctx.features.find( + (f) => f.id === TestFeature.AiCredits, + ); + + const customerBefore = + await autumnV1.customers.get(customerId); + expect(customerBefore.features[TestFeature.AiCredits].balance).toBe(1000); + + const inputTokens = 1000; + const outputTokens = 500; + const modelId = "anthropic/claude-sonnet-4-20250514"; + + const expectedCost = await getCreditCost({ + featureId: aiCreditFeature!.id, + creditSystem: aiCreditFeature!, + modelName: modelId, + tokens: { input: inputTokens, output: outputTokens }, + }); + + const trackRes: TrackResponseV2 = await autumnV2.post("/trackTokens", { + customer_id: customerId, + feature_id: TestFeature.AiCredits, + model_id: modelId, + input_tokens: inputTokens, + output_tokens: outputTokens, + }); + + expect(trackRes.customer_id).toBe(customerId); + expect(trackRes.value).toBeCloseTo(expectedCost, 10); + + const customerAfter = await autumnV1.customers.get(customerId); + expect(customerAfter.features[TestFeature.AiCredits]).toMatchObject({ + balance: new Decimal(1000).minus(expectedCost).toNumber(), + usage: expectedCost, + }); +}); + +// ═══════════════════════════════════════════════════════════════════ +// TRACK-TOKENS-2: Disambiguation error + explicit feature_id resolution +// ═══════════════════════════════════════════════════════════════════ + +test.concurrent(`${chalk.yellowBright("track-tokens-2: disambiguation error and explicit feature_id resolution")}`, async () => { + const aiCreditsItem = items.free({ + featureId: TestFeature.AiCredits, + includedUsage: 500, + }); + const aiCredits2Item = items.free({ + featureId: TestFeature.AiCredits2, + includedUsage: 500, + }); + const freeProd = products.base({ + id: "free", + items: [aiCreditsItem, aiCredits2Item], + }); + + const { customerId, autumnV1, autumnV2, ctx } = await initScenario({ + customerId: "track-tokens-2", + setup: [s.customer({ testClock: false }), s.products({ list: [freeProd] })], + actions: [s.attach({ productId: freeProd.id })], + }); + + // Without feature_id, should fail with disambiguation error + let error: any; + try { + await autumnV2.post("/trackTokens", { + customer_id: customerId, + model_id: "anthropic/claude-sonnet-4-20250514", + input_tokens: 100, + output_tokens: 50, + }); + } catch (e) { + error = e; + } + expect(error).toBeDefined(); + expect(error.message).toContain("Multiple AI credit system features"); + + // With explicit feature_id, should succeed and only deduct from AiCredits + const aiCreditFeature = ctx.features.find( + (f) => f.id === TestFeature.AiCredits, + ); + + const inputTokens = 2000; + const outputTokens = 1000; + const modelId = "anthropic/claude-sonnet-4-20250514"; + + const expectedCost = await getCreditCost({ + featureId: aiCreditFeature!.id, + creditSystem: aiCreditFeature!, + modelName: modelId, + tokens: { input: inputTokens, output: outputTokens }, + }); + + const trackRes: TrackResponseV2 = await autumnV2.post("/trackTokens", { + customer_id: customerId, + feature_id: TestFeature.AiCredits, + model_id: modelId, + input_tokens: inputTokens, + output_tokens: outputTokens, + }); + + expect(trackRes.customer_id).toBe(customerId); + expect(trackRes.value).toBeCloseTo(expectedCost, 10); + + const customer = await autumnV1.customers.get(customerId); + expect(customer.features[TestFeature.AiCredits]).toMatchObject({ + balance: new Decimal(500).minus(expectedCost).toNumber(), + usage: expectedCost, + }); + expect(customer.features[TestFeature.AiCredits2]).toMatchObject({ + balance: 500, + usage: 0, + }); +}); + +// ═══════════════════════════════════════════════════════════════════ +// TRACK-TOKENS-3: custom/* model pricing (with and without markup) +// ═══════════════════════════════════════════════════════════════════ + +test.concurrent(`${chalk.yellowBright("track-tokens-3: custom model pricing with and without markup")}`, async () => { + const aiCreditsItem = items.free({ + featureId: TestFeature.AiCredits, + includedUsage: 1000, + }); + const freeProd = products.base({ + id: "free", + items: [aiCreditsItem], + }); + + const { customerId, autumnV1, autumnV2 } = await initScenario({ + customerId: "track-tokens-3", + setup: [s.customer({ testClock: false }), s.products({ list: [freeProd] })], + actions: [s.attach({ productId: freeProd.id })], + }); + + // custom/internal-model: input_cost=5 $/M, output_cost=15 $/M, markup=0% + const expectedCostNoMarkup = new Decimal(5) + .mul(10000) + .add(new Decimal(15).mul(5000)) + .div(1_000_000) + .toNumber(); // 0.125 + + const trackRes1: TrackResponseV2 = await autumnV2.post("/trackTokens", { + customer_id: customerId, + feature_id: TestFeature.AiCredits, + model_id: "custom/internal-model", + input_tokens: 10000, + output_tokens: 5000, + }); + + expect(trackRes1.value).toBeCloseTo(expectedCostNoMarkup, 10); + + // custom/marked-up-model: input_cost=10 $/M, output_cost=30 $/M, markup=50% + const baseCost = new Decimal(10) + .mul(8000) + .add(new Decimal(30).mul(2000)) + .div(1_000_000); + const expectedCostWithMarkup = baseCost + .mul(new Decimal(1).add(new Decimal(50).div(100))) + .toNumber(); // 0.21 + + const trackRes2: TrackResponseV2 = await autumnV2.post("/trackTokens", { + customer_id: customerId, + feature_id: TestFeature.AiCredits, + model_id: "custom/marked-up-model", + input_tokens: 8000, + output_tokens: 2000, + }); + + expect(trackRes2.value).toBeCloseTo(expectedCostWithMarkup, 10); + + const totalCost = new Decimal(expectedCostNoMarkup) + .plus(expectedCostWithMarkup) + .toNumber(); + + const customer = await autumnV1.customers.get(customerId); + expect(customer.features[TestFeature.AiCredits]).toMatchObject({ + balance: new Decimal(1000).minus(totalCost).toNumber(), + usage: totalCost, + }); +}); + +// ═══════════════════════════════════════════════════════════════════ +// TRACK-TOKENS-4: models.dev pricing with markup + error for non-AI feature +// ═══════════════════════════════════════════════════════════════════ + +test.concurrent(`${chalk.yellowBright("track-tokens-4: models.dev markup and non-AI feature_id error")}`, async () => { + const aiCreditsItem = items.free({ + featureId: TestFeature.AiCredits, + includedUsage: 1000, + }); + const creditsItem = items.free({ + featureId: TestFeature.Credits, + includedUsage: 100, + }); + const freeProd = products.base({ + id: "free", + items: [aiCreditsItem, creditsItem], + }); + + const { customerId, autumnV1, autumnV2, ctx } = await initScenario({ + customerId: "track-tokens-4", + setup: [s.customer({ testClock: false }), s.products({ list: [freeProd] })], + actions: [s.attach({ productId: freeProd.id })], + }); + + const aiCreditFeature = ctx.features.find( + (f) => f.id === TestFeature.AiCredits, + ); + + // anthropic/claude-haiku-3.5 has 20% markup in test config + const inputTokens = 50000; + const outputTokens = 10000; + const modelId = "anthropic/claude-3-5-haiku-20241022"; + + const expectedCost = await getCreditCost({ + featureId: aiCreditFeature!.id, + creditSystem: aiCreditFeature!, + modelName: modelId, + tokens: { input: inputTokens, output: outputTokens }, + }); + + const trackRes: TrackResponseV2 = await autumnV2.post("/trackTokens", { + customer_id: customerId, + feature_id: TestFeature.AiCredits, + model_id: modelId, + input_tokens: inputTokens, + output_tokens: outputTokens, + }); + + expect(trackRes.value).toBeCloseTo(expectedCost, 10); + + const customer = await autumnV1.customers.get(customerId); + expect(customer.features[TestFeature.AiCredits]).toMatchObject({ + balance: new Decimal(1000).minus(expectedCost).toNumber(), + usage: expectedCost, + }); + + // Pointing at a regular credit system should fail + let error: any; + try { + await autumnV2.post("/trackTokens", { + customer_id: customerId, + feature_id: TestFeature.Credits, + model_id: "anthropic/claude-sonnet-4-20250514", + input_tokens: 100, + output_tokens: 50, + }); + } catch (e) { + error = e; + } + expect(error).toBeDefined(); + expect(error.message).toContain("not an AI credit system"); +}); + +// ═══════════════════════════════════════════════════════════════════ +// TRACK-TOKENS-5: Multiple tracks accumulate correctly +// ═══════════════════════════════════════════════════════════════════ + +test.concurrent(`${chalk.yellowBright("track-tokens-5: multiple tracks accumulate balance deductions")}`, async () => { + const aiCreditsItem = items.free({ + featureId: TestFeature.AiCredits, + includedUsage: 1000, + }); + const freeProd = products.base({ + id: "free", + items: [aiCreditsItem], + }); + + const { customerId, autumnV1, autumnV2, ctx } = await initScenario({ + customerId: "track-tokens-5", + setup: [s.customer({ testClock: false }), s.products({ list: [freeProd] })], + actions: [s.attach({ productId: freeProd.id })], + }); + + const aiCreditFeature = ctx.features.find( + (f) => f.id === TestFeature.AiCredits, + ); + + // First track: custom/internal-model (input_cost=5, output_cost=15, markup=0%) + const cost1 = await getCreditCost({ + featureId: aiCreditFeature!.id, + creditSystem: aiCreditFeature!, + modelName: "custom/internal-model", + tokens: { input: 5000, output: 2000 }, + }); + + await autumnV2.post("/trackTokens", { + customer_id: customerId, + feature_id: TestFeature.AiCredits, + model_id: "custom/internal-model", + input_tokens: 5000, + output_tokens: 2000, + }); + + // Second track: custom/marked-up-model (input_cost=10, output_cost=30, markup=50%) + const cost2 = await getCreditCost({ + featureId: aiCreditFeature!.id, + creditSystem: aiCreditFeature!, + modelName: "custom/marked-up-model", + tokens: { input: 3000, output: 1000 }, + }); + + await autumnV2.post("/trackTokens", { + customer_id: customerId, + feature_id: TestFeature.AiCredits, + model_id: "custom/marked-up-model", + input_tokens: 3000, + output_tokens: 1000, + }); + + const totalCost = new Decimal(cost1).plus(cost2).toNumber(); + + const customer = await autumnV1.customers.get(customerId); + expect(customer.features[TestFeature.AiCredits]).toMatchObject({ + balance: new Decimal(1000).minus(totalCost).toNumber(), + usage: totalCost, + }); +}); diff --git a/server/tests/integration/balances/track/overage-allowed/track-overage-allowed-consumable.test.ts b/server/tests/integration/balances/track/overage-allowed/track-overage-allowed-consumable.test.ts index 6e3242fad..636513eb9 100644 --- a/server/tests/integration/balances/track/overage-allowed/track-overage-allowed-consumable.test.ts +++ b/server/tests/integration/balances/track/overage-allowed/track-overage-allowed-consumable.test.ts @@ -368,7 +368,7 @@ test.concurrent(`${chalk.yellowBright("track-consumable-overage-8: credit system (feature) => feature.id === TestFeature.Credits, )!; - const action1CreditCost = getCreditCost({ + const action1CreditCost = await getCreditCost({ featureId: TestFeature.Action1, creditSystem: creditsFeature, amount: 1, diff --git a/server/tests/integration/balances/track/spend-limit/track-customer-spend-limit.test.ts b/server/tests/integration/balances/track/spend-limit/track-customer-spend-limit.test.ts index 64e0b3622..25c11d2ec 100644 --- a/server/tests/integration/balances/track/spend-limit/track-customer-spend-limit.test.ts +++ b/server/tests/integration/balances/track/spend-limit/track-customer-spend-limit.test.ts @@ -430,7 +430,7 @@ test.concurrent(`${chalk.yellowBright("track-customer-spend-limit6: credit-syste const creditsFeature = ctx.features.find( (feature) => feature.id === TestFeature.Credits, )!; - const action1CreditCost = getCreditCost({ + const action1CreditCost = await getCreditCost({ featureId: TestFeature.Action1, creditSystem: creditsFeature, amount: 1, diff --git a/server/tests/integration/balances/track/spend-limit/track-entity-product-spend-limit.test.ts b/server/tests/integration/balances/track/spend-limit/track-entity-product-spend-limit.test.ts index 85c5ab273..79fae21fa 100644 --- a/server/tests/integration/balances/track/spend-limit/track-entity-product-spend-limit.test.ts +++ b/server/tests/integration/balances/track/spend-limit/track-entity-product-spend-limit.test.ts @@ -463,7 +463,7 @@ test.concurrent(`${chalk.yellowBright("track-entity-product-spend-limit5: credit const creditsFeature = ctx.features.find( (feature) => feature.id === TestFeature.Credits, )!; - const action1CreditCost = getCreditCost({ + const action1CreditCost = await getCreditCost({ featureId: TestFeature.Action1, creditSystem: creditsFeature, amount: 1, diff --git a/server/tests/integration/balances/track/spend-limit/track-per-entity-spend-limit.test.ts b/server/tests/integration/balances/track/spend-limit/track-per-entity-spend-limit.test.ts index e61f687bd..89822ddad 100644 --- a/server/tests/integration/balances/track/spend-limit/track-per-entity-spend-limit.test.ts +++ b/server/tests/integration/balances/track/spend-limit/track-per-entity-spend-limit.test.ts @@ -489,7 +489,7 @@ test.concurrent(`${chalk.yellowBright("track-per-entity-spend-limit5: credit-sys const creditsFeature = ctx.features.find( (feature) => feature.id === TestFeature.Credits, )!; - const action1CreditCost = getCreditCost({ + const action1CreditCost = await getCreditCost({ featureId: TestFeature.Action1, creditSystem: creditsFeature, amount: 1, diff --git a/server/tests/integration/balances/track/spend-limit/track-postgres-entity-spend-limit.test.ts b/server/tests/integration/balances/track/spend-limit/track-postgres-entity-spend-limit.test.ts index 00422b328..72ebef0b4 100644 --- a/server/tests/integration/balances/track/spend-limit/track-postgres-entity-spend-limit.test.ts +++ b/server/tests/integration/balances/track/spend-limit/track-postgres-entity-spend-limit.test.ts @@ -216,7 +216,7 @@ test.concurrent(`${chalk.yellowBright("track-postgres-entity-spend-limit3: credi const creditsFeature = ctx.features.find( (feature) => feature.id === TestFeature.Credits, )!; - const action1CreditCost = getCreditCost({ + const action1CreditCost = await getCreditCost({ featureId: TestFeature.Action1, creditSystem: creditsFeature, amount: 1, @@ -334,7 +334,7 @@ test.concurrent(`${chalk.yellowBright("track-postgres-entity-spend-limit4: prepa const creditsFeature = ctx.features.find( (feature) => feature.id === TestFeature.Credits, )!; - const action1CreditCost = getCreditCost({ + const action1CreditCost = await getCreditCost({ featureId: TestFeature.Action1, creditSystem: creditsFeature, amount: 1, diff --git a/server/tests/setup/v2Features.ts b/server/tests/setup/v2Features.ts index 26b8aa106..2a4150b83 100644 --- a/server/tests/setup/v2Features.ts +++ b/server/tests/setup/v2Features.ts @@ -4,6 +4,7 @@ dotenv.config(); import { AppEnv, FeatureUsageType } from "@autumn/shared"; import { + constructAiCreditSystem, constructBooleanFeature, constructCreditSystem, constructMeteredFeature, @@ -25,6 +26,9 @@ export enum TestFeature { Action3 = "action3", // single use (pay per use) Credits2 = "credits2", // credit system + + AiCredits = "ai_credits", // AI credit system (models.dev pricing) + AiCredits2 = "ai_credits_2", // second AI credit system (for disambiguation tests) } export const getFeatures = ({ orgId }: { orgId: string }) => ({ @@ -121,4 +125,37 @@ export const getFeatures = ({ orgId }: { orgId: string }) => ({ }, ], }), + [TestFeature.AiCredits]: constructAiCreditSystem({ + featureId: TestFeature.AiCredits, + orgId, + env: AppEnv.Sandbox, + modelMarkups: { + "anthropic/claude-sonnet-4-20250514": { + markup: 0, + }, + "anthropic/claude-3-5-haiku-20241022": { + markup: 20, + }, + "custom/internal-model": { + markup: 0, + input_cost: 5, + output_cost: 15, + }, + "custom/marked-up-model": { + markup: 50, + input_cost: 10, + output_cost: 30, + }, + }, + }), + [TestFeature.AiCredits2]: constructAiCreditSystem({ + featureId: TestFeature.AiCredits2, + orgId, + env: AppEnv.Sandbox, + modelMarkups: { + "anthropic/claude-sonnet-4-20250514": { + markup: 10, + }, + }, + }), }); diff --git a/server/tests/utils/fixtures/db/entitlements.ts b/server/tests/utils/fixtures/db/entitlements.ts index f8c05e8e2..b8be6cafd 100644 --- a/server/tests/utils/fixtures/db/entitlements.ts +++ b/server/tests/utils/fixtures/db/entitlements.ts @@ -21,6 +21,7 @@ const create = ({ intervalCount = 1, entityFeatureId = null, rollover = null, + modelMarkups = null, }: { id?: string; featureId: string; @@ -33,6 +34,10 @@ const create = ({ intervalCount?: number; entityFeatureId?: string | null; rollover?: RolloverConfig | null; + modelMarkups?: Record< + string, + { markup: number; input_cost?: number; output_cost?: number } + > | null; }) => ({ id: id ?? `ent_${featureId}_${crypto.randomUUID().slice(0, 8)}`, created_at: Date.now(), @@ -54,6 +59,7 @@ const create = ({ name: featureName, type: featureType, config: featureConfig, + modelMarkups: modelMarkups ?? null, }), }); diff --git a/server/tests/utils/fixtures/db/features.ts b/server/tests/utils/fixtures/db/features.ts index 807712d42..695a5053a 100644 --- a/server/tests/utils/fixtures/db/features.ts +++ b/server/tests/utils/fixtures/db/features.ts @@ -9,12 +9,17 @@ const create = ({ name, type = FeatureType.Metered, config = {}, + modelMarkups = null, }: { id: string; internalId?: string; name: string; type?: FeatureType; config?: Record; + modelMarkups?: Record< + string, + { markup: number; input_cost?: number; output_cost?: number } + > | null; }) => ({ internal_id: internalId ?? `internal_${id}`, org_id: "org_test", @@ -27,6 +32,7 @@ const create = ({ display: null, archived: false, event_names: [], + model_markups: modelMarkups ?? null, }); // ═══════════════════════════════════════════════════════════════════ diff --git a/shared/api/balances/track/trackTokensParams.ts b/shared/api/balances/track/trackTokensParams.ts new file mode 100644 index 000000000..3d79830a5 --- /dev/null +++ b/shared/api/balances/track/trackTokensParams.ts @@ -0,0 +1,45 @@ +import { z } from "zod/v4"; +import { CustomerDataSchema } from "../../common/customerData"; +import { EntityDataSchema } from "../../common/entityData"; + +export const TrackTokensParamsSchema = z.object({ + customer_id: z.string().meta({ + description: "The ID of the customer.", + }), + entity_id: z.string().optional().meta({ + description: "The ID of the entity for entity-scoped balances.", + }), + feature_id: z.string().optional().meta({ + description: + "The ID of the AI credit system feature. Auto-detected if omitted.", + }), + model_id: z.string().meta({ + description: "The AI model name with provider prefix (e.g., 'anthropic/claude-opus-4-6').", + }), + input_tokens: z.number().int().nonnegative().meta({ + description: "Number of input tokens consumed.", + }), + output_tokens: z.number().int().nonnegative().meta({ + description: "Number of output tokens consumed.", + }), + properties: z.record(z.string(), z.any()).optional().meta({ + description: "Additional properties to attach to this usage event.", + }), + idempotency_key: z.string().optional().meta({ + internal: true, + }), + overage_behavior: z.enum(["cap", "reject"]).optional().meta({ + internal: true, + }), + customer_data: CustomerDataSchema.optional().meta({ + internal: true, + }), + entity_data: EntityDataSchema.optional().meta({ + internal: true, + }), + skip_event: z.boolean().optional().meta({ + internal: true, + }), +}); + +export type TrackTokensParams = z.infer; diff --git a/shared/api/features/apiFeatureV1.ts b/shared/api/features/apiFeatureV1.ts index a5ebac6c1..bc3b92cfc 100644 --- a/shared/api/features/apiFeatureV1.ts +++ b/shared/api/features/apiFeatureV1.ts @@ -1,4 +1,5 @@ import { z } from "zod/v4"; +import { ModelMarkupsSchema } from "../../models/featureModels/featureConfig/creditConfig"; import { FeatureType } from "../../models/featureModels/featureEnums"; export const ApiFeatureV1Schema = z.object({ @@ -12,7 +13,7 @@ export const ApiFeatureV1Schema = z.object({ }), type: z.enum(FeatureType).meta({ description: - "Feature type: 'boolean' for on/off access, 'metered' for usage-tracked features, 'credit_system' for unified credit pools.", + "Feature type: 'boolean' for on/off access, 'metered' for usage-tracked features, 'credit_system' for unified credit pools, 'ai_credit_system' for model-based token pricing.", }), consumable: z.boolean().meta({ @@ -42,6 +43,10 @@ export const ApiFeatureV1Schema = z.object({ "For credit_system features: maps metered features to their credit costs.", }), + model_markups: ModelMarkupsSchema.optional().meta({ + description: "Per-model markup percentages for AI credit systems.", + }), + display: z .object({ singular: z.string().nullish().meta({ diff --git a/shared/api/features/changes/V1.2_FeatureChange.ts b/shared/api/features/changes/V1.2_FeatureChange.ts index e08a1e2d1..101e9cc57 100644 --- a/shared/api/features/changes/V1.2_FeatureChange.ts +++ b/shared/api/features/changes/V1.2_FeatureChange.ts @@ -83,7 +83,10 @@ export const V1_2_FeatureChange = defineVersionChange({ plural: input.display.plural || "", } : null, - credit_schema: input.credit_schema || null, + credit_schema: input.credit_schema?.map((item) => ({ + metered_feature_id: item.metered_feature_id, + credit_cost: item.credit_cost, + })) || null, archived: input.archived, } satisfies z.infer; }, diff --git a/shared/api/features/crud/common/baseFeatureParamsV1.ts b/shared/api/features/crud/common/baseFeatureParamsV1.ts index e3aa3fcd5..69fee0f97 100644 --- a/shared/api/features/crud/common/baseFeatureParamsV1.ts +++ b/shared/api/features/crud/common/baseFeatureParamsV1.ts @@ -1,4 +1,5 @@ import { z } from "zod/v4"; +import { ModelMarkupsSchema } from "../../../../models/featureModels/featureConfig/creditConfig"; import { FeatureType } from "../../../../models/featureModels/featureEnums"; import { idRegex } from "../../../../utils/utils"; @@ -43,8 +44,13 @@ export const BaseFeatureV1ParamsSchema = z.object({ .optional() .meta({ description: - "A schema that maps 'single_use' feature IDs to credit costs. Applicable only for 'credit_system' features.", + "A schema that maps 'single_use' feature IDs to credit costs. For classic credit systems only — AI credit systems use model_markups instead.", }), + model_markups: ModelMarkupsSchema.optional().meta({ + description: + "Per-model markup percentages for AI credit systems. Maps model IDs to their markup configuration.", + }), + event_names: z.array(z.string()).optional(), }); diff --git a/shared/api/models.ts b/shared/api/models.ts index 23271c21c..289bd8d40 100644 --- a/shared/api/models.ts +++ b/shared/api/models.ts @@ -47,6 +47,7 @@ export * from "./balances/index.js"; export * from "./balances/prevVersions/legacyUpdateBalanceModels.js"; export * from "./balances/track/prevVersions/trackResponseV1.js"; export * from "./balances/track/trackParams.js"; +export * from "./balances/track/trackTokensParams.js"; export * from "./balances/track/trackResponseV2.js"; export * from "./balances/track/trackResponseV3.js"; export * from "./balances/update/updateBalanceParams.js"; diff --git a/shared/index.ts b/shared/index.ts index 437e4bf13..b93810450 100644 --- a/shared/index.ts +++ b/shared/index.ts @@ -195,6 +195,8 @@ export * from "./models/rewardModels/rewardProgramModels/rewardProgramTable"; export * from "./models/scheduleModels/scheduleTable"; export * from "./models/subModels/subModels"; export * from "./models/subModels/subTable"; +// AI Models +export * from "./models/aiModels/modelsDevTypes"; export * from "./types"; // Agent Types (for pricing agent AI) export * from "./utils/agentTypes"; diff --git a/shared/models/aiModels/modelsDevTypes.ts b/shared/models/aiModels/modelsDevTypes.ts new file mode 100644 index 000000000..35f7bb30e --- /dev/null +++ b/shared/models/aiModels/modelsDevTypes.ts @@ -0,0 +1,17 @@ +/** Shape of a single model from the models.dev API */ +export interface ModelsDevModel { + id: string; + name: string; + release_date?: string; + cost: { + input: number; + output: number; + }; +} + +/** Shape of a provider from the models.dev API */ +export interface ModelsDevProvider { + id: string; + name: string; + models: Record; +} diff --git a/shared/models/featureModels/featureConfig/creditConfig.ts b/shared/models/featureModels/featureConfig/creditConfig.ts index c75cc8265..5bc5b9b51 100644 --- a/shared/models/featureModels/featureConfig/creditConfig.ts +++ b/shared/models/featureModels/featureConfig/creditConfig.ts @@ -11,12 +11,23 @@ export const CreditSystemConfigSchema = z.object({ schema: z.array( z.object({ metered_feature_id: z.string(), - // feature_amount: z.number(), credit_amount: z.number(), }), ), usage_type: z.nativeEnum(FeatureUsageType), }); +export const ModelMarkupsSchema = z + .record( + z.string(), // Represents the model name in "provider/model" format, e.g. "anthropic/claude-2" + z.object({ + markup: z.number().min(0), // percentage markup, e.g. 20 for 20% + input_cost: z.number().min(0).optional(), // $/M tokens, required for custom/ models + output_cost: z.number().min(0).optional(), // $/M tokens, required for custom/ models + }), + ) + .nullish(); + export type CreditSystemConfig = z.infer; export type CreditSchemaItem = z.infer; +export type ModelMarkups = z.infer; diff --git a/shared/models/featureModels/featureEnums.ts b/shared/models/featureModels/featureEnums.ts index 430d87a62..c75480114 100644 --- a/shared/models/featureModels/featureEnums.ts +++ b/shared/models/featureModels/featureEnums.ts @@ -2,6 +2,7 @@ export enum FeatureType { Boolean = "boolean", Metered = "metered", CreditSystem = "credit_system", + AiCreditSystem = "ai_credit_system", } export enum AggregateType { diff --git a/shared/models/featureModels/featureModels.ts b/shared/models/featureModels/featureModels.ts index 87fd02e7b..fd9afeff7 100644 --- a/shared/models/featureModels/featureModels.ts +++ b/shared/models/featureModels/featureModels.ts @@ -1,3 +1,4 @@ +import { ModelMarkupsSchema } from "@models/featureModels/featureConfig/creditConfig"; import { z } from "zod/v4"; import { AppEnv } from "../genModels/genEnums"; import { FeatureType } from "./featureEnums"; @@ -22,6 +23,7 @@ export const FeatureSchema = z.object({ .nullish(), archived: z.boolean(), event_names: z.array(z.string()).default([]), + model_markups: ModelMarkupsSchema.nullish(), }); export const CreateFeatureSchema = FeatureSchema.pick({ @@ -31,6 +33,7 @@ export const CreateFeatureSchema = FeatureSchema.pick({ config: true, display: true, event_names: true, + model_markups: true, }); export const MinFeatureSchema = z.object({ diff --git a/shared/models/featureModels/featureTable.ts b/shared/models/featureModels/featureTable.ts index 07cb69d00..4b9d7e81a 100644 --- a/shared/models/featureModels/featureTable.ts +++ b/shared/models/featureModels/featureTable.ts @@ -10,7 +10,7 @@ import { } from "drizzle-orm/pg-core"; import { collatePgColumn } from "../../db/utils"; import { organizations } from "../orgModels/orgTable"; -import type { CreditSystemConfig } from "./featureConfig/creditConfig"; +import type { CreditSystemConfig, ModelMarkups } from "./featureConfig/creditConfig"; import type { MeteredConfig } from "./featureConfig/meteredConfig"; type FeatureDisplay = { @@ -33,6 +33,7 @@ export const features = pgTable( display: jsonb().default(sql`null`).$type(), archived: boolean("archived").notNull().default(false), event_names: text("event_names").array().default([]), + model_markups: jsonb().$type().default(sql`null`), }, (table) => [ foreignKey({ diff --git a/shared/utils/agentTypes.ts b/shared/utils/agentTypes.ts index b2d0abfdb..69e0ca6ad 100644 --- a/shared/utils/agentTypes.ts +++ b/shared/utils/agentTypes.ts @@ -27,7 +27,8 @@ export type AgentFeatureType = | "boolean" | "single_use" | "continuous_use" - | "credit_system"; + | "credit_system" + | "ai_credit_system"; export interface AgentFeature { id: string; @@ -41,6 +42,14 @@ export interface AgentFeature { metered_feature_id: string; credit_cost: number; }> | null; + model_markups?: Record< + string, + { + markup: number; + input_cost?: number; + output_cost?: number; + } + > | null; } export interface AgentProductItem { @@ -84,6 +93,8 @@ function mapAgentTypeToFeatureType(agentType: AgentFeatureType): FeatureType { return FeatureType.Boolean; case "credit_system": return FeatureType.CreditSystem; + case "ai_credit_system": + return FeatureType.AiCreditSystem; default: return FeatureType.Metered; } @@ -129,6 +140,7 @@ export function agentFeatureToFeature(agentFeature: AgentFeature): Feature { display: agentFeature.display ?? undefined, archived: false, event_names: [], + model_markups: agentFeature.model_markups ?? null, }; } @@ -217,6 +229,9 @@ export function featureToAgentFeature(feature: Feature): AgentFeature { }), ); } + if (feature.type === FeatureType.AiCreditSystem) { + agentFeature.model_markups = feature.model_markups; + } return agentFeature; } diff --git a/shared/utils/featureUtils.ts b/shared/utils/featureUtils.ts index de6d0db2e..0448e8073 100644 --- a/shared/utils/featureUtils.ts +++ b/shared/utils/featureUtils.ts @@ -18,7 +18,7 @@ export const toApiFeature = ({ feature }: { feature: Feature }) => { } let creditSchema: CreditSchemaItem[] | undefined; - if (feature.type === FeatureType.CreditSystem) { + if (feature.type === FeatureType.CreditSystem && feature.config?.schema) { creditSchema = feature.config.schema.map((s: CreditSchemaItem) => ({ metered_feature_id: s.metered_feature_id, credit_cost: s.credit_amount, diff --git a/shared/utils/featureUtils/apiFeatureToDbFeature.ts b/shared/utils/featureUtils/apiFeatureToDbFeature.ts index aa1c5bf1a..dff59220a 100644 --- a/shared/utils/featureUtils/apiFeatureToDbFeature.ts +++ b/shared/utils/featureUtils/apiFeatureToDbFeature.ts @@ -72,6 +72,7 @@ export const apiFeatureToDbFeature = ({ config: newConfig, archived: apiFeature.archived ?? originalFeature?.archived ?? false, event_names: [], + model_markups: null, } satisfies Feature; }; @@ -150,6 +151,9 @@ export const featureV1ToDbFeature = ({ ); } + const modelMarkups = + apiFeature.model_markups ?? originalFeature?.model_markups ?? null; + return { internal_id: originalFeature?.internal_id ?? "", org_id: originalFeature?.org_id ?? "", @@ -165,6 +169,7 @@ export const featureV1ToDbFeature = ({ ? apiFeature.archived : (originalFeature?.archived ?? false), event_names: eventNames ?? [], + model_markups: modelMarkups, } satisfies Feature; }; @@ -192,6 +197,7 @@ export const dbToApiFeatureV1 = ({ type: dbFeature.type, consumable: dbFeature.type === FeatureType.CreditSystem || + dbFeature.type === FeatureType.AiCreditSystem || dbFeature.config?.usage_type === FeatureUsageType.Single, credit_schema: Array.isArray(dbFeature.config?.schema) @@ -200,6 +206,7 @@ export const dbToApiFeatureV1 = ({ credit_cost: schema.credit_amount, })) : undefined, + model_markups: dbFeature.model_markups ?? undefined, event_names: Array.isArray(dbFeature.event_names) ? dbFeature.event_names : [], diff --git a/shared/utils/featureUtils/creditSystemUtils.ts b/shared/utils/featureUtils/creditSystemUtils.ts index 47c0095a8..039dbffde 100644 --- a/shared/utils/featureUtils/creditSystemUtils.ts +++ b/shared/utils/featureUtils/creditSystemUtils.ts @@ -12,7 +12,8 @@ export const creditSystemContainsFeature = ({ if (creditSystem.type !== FeatureType.CreditSystem) { return false; } - const schema: CreditSchemaItem[] = creditSystem.config.schema; + const schema: CreditSchemaItem[] | undefined = creditSystem.config?.schema; + if (!schema) return false; for (const schemaItem of schema) { if (schemaItem.metered_feature_id === meteredFeatureId) { diff --git a/shared/utils/productDisplayUtils.ts b/shared/utils/productDisplayUtils.ts index 1cba7d5a5..7d32f4dc2 100644 --- a/shared/utils/productDisplayUtils.ts +++ b/shared/utils/productDisplayUtils.ts @@ -54,7 +54,9 @@ const getIncludedUsageText = (item: ProductItem, feature: Feature): string => { if (item.included_usage === Infinite) { return `Unlimited ${featureName}`; } - + if (feature.type === FeatureType.AiCreditSystem) { + return `$${numberWithCommas(item.included_usage ?? 0)} of ${featureName}`; + } if (nullish(item.included_usage) || item.included_usage === 0) { return `0 ${featureName}`; } @@ -219,14 +221,20 @@ export const getFeaturePriceItemDisplay = ({ // Build included usage string (e.g., "100 credits") const includedUsage = item.included_usage as number | null; const hasIncludedUsage = notNullish(includedUsage) && includedUsage > 0; + const isAiCreditSystem = feature.type === FeatureType.AiCreditSystem; const includedFeatureName = getFeatureName({ feature, units: item.included_usage, }); - const includedUsageStr = hasIncludedUsage - ? `${numberWithCommas(includedUsage)} ${includedFeatureName}` - : ""; + let includedUsageStr = ""; + if (hasIncludedUsage) { + if (isAiCreditSystem) { + includedUsageStr = `$${numberWithCommas(includedUsage)} of ${includedFeatureName}`; + } else { + includedUsageStr = `${numberWithCommas(includedUsage)} ${includedFeatureName}`; + } + } const volumeFlatAmount = isVolumeFlatAmountItem(item); @@ -249,10 +257,18 @@ export const getFeaturePriceItemDisplay = ({ feature, units: billingUnits, }); - const perUnitStr = - billingUnits > 1 - ? `${numberWithCommas(billingUnits)} ${billingFeatureName}` - : billingFeatureName; + let perUnitStr: string; + if (isAiCreditSystem) { + perUnitStr = + billingUnits > 1 + ? `$${numberWithCommas(billingUnits)} of ${billingFeatureName}` + : `$1 of ${billingFeatureName}`; + } else { + perUnitStr = + billingUnits > 1 + ? `${numberWithCommas(billingUnits)} ${billingFeatureName}` + : billingFeatureName; + } // Build interval string const showInterval = isMainPrice || fullDisplay; @@ -267,6 +283,13 @@ export const getFeaturePriceItemDisplay = ({ } // Format output based on what we have + if (isAiCreditSystem) { + return { + primary_text: includedUsageStr || "$0 included", + secondary_text: "then charged based on model usage", + }; + } + if (hasIncludedUsage) { if (volumeFlatAmount) { const featureName = getFeatureName({ feature, units: 2 }); diff --git a/vite/src/hooks/queries/useAiModelsQuery.ts b/vite/src/hooks/queries/useAiModelsQuery.ts new file mode 100644 index 000000000..15de1a42c --- /dev/null +++ b/vite/src/hooks/queries/useAiModelsQuery.ts @@ -0,0 +1,24 @@ +import type { ModelsDevProvider } from "@autumn/shared"; +import { useQuery } from "@tanstack/react-query"; +import { useAxiosInstance } from "@/services/useAxiosInstance"; + +export const useModelsDevPricing = () => { + const axiosInstance = useAxiosInstance(); + + const { data, isLoading, error } = useQuery({ + queryKey: ["models-dev-pricing"], + queryFn: async () => { + const { data } = await axiosInstance.get< + Record + >("/v1/features/ai/model_pricing"); + return data; + }, + staleTime: 1000 * 60 * 10, + }); + + return { + providers: data ?? {}, + isLoading, + error, + }; +}; diff --git a/vite/src/utils/product/product-item/formatProductItem.ts b/vite/src/utils/product/product-item/formatProductItem.ts index 212d9647e..b29ae767c 100644 --- a/vite/src/utils/product/product-item/formatProductItem.ts +++ b/vite/src/utils/product/product-item/formatProductItem.ts @@ -140,13 +140,15 @@ const getFeatureString = ({ item: ProductItem; features: Feature[]; }) => { - const feature = features.find((f: Feature) => f.id == item.feature_id); + // This function isn't exported or used + // Just updated it to account for the new AI Credit System features, but it may be unused and can probably be deleted + const feature = features.find((f: Feature) => f.id === item.feature_id); if (feature?.type === FeatureType.Boolean) { return `${feature.name}`; } - if (item.included_usage == Infinite) { + if (item.included_usage === Infinite) { return `Unlimited ${feature?.name}`; } @@ -155,6 +157,14 @@ const getFeatureString = ({ intervalCount: item.interval_count ?? undefined, }); + const isAiCreditSystem = feature?.type === FeatureType.AiCreditSystem; + if (isAiCreditSystem) { + const amount = item.included_usage ?? 0; + const formattedAmount = + amount === 0 ? "$0.00" : `$${Number(amount).toFixed(2)}`; + return `${formattedAmount} of ${feature?.name}${notNullish(item.interval) ? ` ${intervalStr}` : ""}`; + } + return `${item.included_usage ?? 0} ${feature?.name}${item.entity_feature_id ? ` per ${getFeature(item.entity_feature_id, features)?.name}` : ""}${notNullish(item.interval) ? ` ${intervalStr}` : ""}`; }; @@ -171,13 +181,13 @@ export const formatProductItemText = ({ const itemType = getItemType(item); - if (itemType == ProductItemType.FeaturePrice) { + if (itemType === ProductItemType.FeaturePrice) { return getPaidFeatureString({ item, currency: org?.default_currency, features, }); - } else if (itemType == ProductItemType.Price) { + } else if (itemType === ProductItemType.Price) { return getFixedPriceString({ item, currency: org?.default_currency }); } }; diff --git a/vite/src/views/products/features/components/CreateFeatureSheet.tsx b/vite/src/views/products/features/components/CreateFeatureSheet.tsx index fed83c894..527f6a0ea 100644 --- a/vite/src/views/products/features/components/CreateFeatureSheet.tsx +++ b/vite/src/views/products/features/components/CreateFeatureSheet.tsx @@ -70,6 +70,8 @@ function CreateFeatureSheet({ setLoading(false); } else { try { + const isAiCreditSystem = feature.type === FeatureType.AiCreditSystem; + const { data: createdFeature } = await FeatureService.createFeature( axiosInstance, { @@ -77,12 +79,13 @@ function CreateFeatureSheet({ id: feature.id, type: feature.type, consumable: feature.config?.usage_type === FeatureUsageType.Single, - credit_schema: feature.config?.schema?.map( - (x: CreditSchemaItem) => ({ - metered_feature_id: x.metered_feature_id, - credit_cost: x.credit_amount, - }), - ), + model_markups: feature.model_markups ?? undefined, + credit_schema: isAiCreditSystem + ? undefined + : feature.config?.schema?.map((x: CreditSchemaItem) => ({ + metered_feature_id: x.metered_feature_id, + credit_cost: x.credit_amount, + })), event_names: feature.event_names, }, ); @@ -119,13 +122,6 @@ function CreateFeatureSheet({ return ( - {/* {!isControlled && ( - - - - )} */} ({ + metered_feature_id: item.metered_feature_id, + credit_cost: item.credit_amount, + })), }); await refetch(); diff --git a/vite/src/views/products/features/credit-systems/components/AiCreditSchema.tsx b/vite/src/views/products/features/credit-systems/components/AiCreditSchema.tsx new file mode 100644 index 000000000..10bb29b0e --- /dev/null +++ b/vite/src/views/products/features/credit-systems/components/AiCreditSchema.tsx @@ -0,0 +1,191 @@ +import type { CreateFeature } from "@autumn/shared"; +import { PlusIcon, X } from "lucide-react"; +import { IconButton } from "@/components/v2/buttons/IconButton"; +import { FormLabel } from "@/components/v2/form/FormLabel"; +import { Input } from "@/components/v2/inputs/Input"; +import { SearchableSelect } from "@/components/v2/selects/SearchableSelect"; +import { useAiCreditSchema } from "../hooks/useAiCreditSchema"; +import { AiCreditSchemaRow } from "./AiCreditSchemaRow"; + +interface AiCreditSchemaProps { + creditSystem: CreateFeature; + setCreditSystem: ( + creditSystem: CreateFeature | ((prev: CreateFeature) => CreateFeature), + ) => void; +} + +export function AiCreditSchema({ + creditSystem, + setCreditSystem, +}: AiCreditSchemaProps) { + const { + providers, + modelsLoading, + modelMarkups, + defaultMarkup, + providerGroups, + activeProviderKeys, + availableProviders, + handleModelChange, + handleMarkupChange, + handleDefaultMarkupChange, + handleCostChange, + handleRemoveModel, + handleRemoveProvider, + addProvider, + addModelToProvider, + } = useAiCreditSchema({ creditSystem, setCreditSystem }); + + return ( +
+
+ Default Markup % + + handleDefaultMarkupChange(Number(e.target.value) || 0) + } + onBlur={(e) => handleDefaultMarkupChange(Number(e.target.value) || 0)} + placeholder="0" + className="w-24" + /> +
+
+ {activeProviderKeys.map((providerKey) => { + const provider = providers[providerKey]; + const modelFullIds = providerGroups[providerKey] ?? []; + const providerName = + provider?.name ?? + providerKey.charAt(0).toUpperCase() + providerKey.slice(1); + + return ( +
+
+ + {providerName} + {providerKey !== "custom" && ( + {providerName} + )} + + } + onClick={() => handleRemoveProvider(providerKey)} + /> +
+ +
+ {providerKey === "custom" && ( +

+ In your API tracking, use the format{" "} + + custom/{"modelId"} + +

+ )} + +
+
+ Model +
+
+ {providerKey === "custom" ? "In $/M" : "Cost In"} +
+
+ {providerKey === "custom" ? "Out $/M" : "Cost Out"} +
+
+ Markup % +
+
+
+ + {modelFullIds.map((fullId) => { + const [, ...parts] = fullId.split("/"); + const modelKey = parts.join("/"); + const isCustom = providerKey === "custom"; + return ( + + handleModelChange(providerKey, oldKey, newKey) + } + onMarkupChange={(key, newMarkup) => + handleMarkupChange(providerKey, key, newMarkup) + } + onCostChange={(key, field, value) => + handleCostChange(providerKey, key, field, value) + } + onRemove={(key) => handleRemoveModel(providerKey, key)} + /> + ); + })} + + addModelToProvider(providerKey)} + className="w-fit mt-0.5" + icon={} + disabled={ + providerKey === "custom" + ? false + : Object.keys(provider?.models ?? {}).length === + modelFullIds.length + } + > + Add model + +
+
+ ); + })} +
+ +

All prices in $/M tokens

+ +
e.stopPropagation()}> + provider.id} + getOptionLabel={(provider) => provider.name} + renderValue={() => ( + + + Add provider + + )} + placeholder="Add provider" + searchable + searchPlaceholder="Search providers..." + emptyText="No providers available" + disabled={modelsLoading} + /> +
+
+ ); +} diff --git a/vite/src/views/products/features/credit-systems/components/AiCreditSchemaRow.tsx b/vite/src/views/products/features/credit-systems/components/AiCreditSchemaRow.tsx new file mode 100644 index 000000000..b0b2c7b29 --- /dev/null +++ b/vite/src/views/products/features/credit-systems/components/AiCreditSchemaRow.tsx @@ -0,0 +1,172 @@ +import type { ModelsDevProvider } from "@autumn/shared"; +import { X } from "lucide-react"; +import { useState } from "react"; +import { IconButton } from "@/components/v2/buttons/IconButton"; +import { Input } from "@/components/v2/inputs/Input"; +import { AiModelSelectDropdown } from "./AiModelSelectDropdown"; + +interface AiCreditSchemaRowProps { + modelKey: string; + markup: number; + provider: ModelsDevProvider; + isLoading: boolean; + isCustom?: boolean; + inputCost?: number; + outputCost?: number; + onModelChange: (oldModelKey: string, newModelKey: string) => void; + onMarkupChange: (modelKey: string, markup: number) => void; + onCostChange?: ( + modelKey: string, + field: "input_cost" | "output_cost", + value: number, + ) => void; + onRemove: (modelKey: string) => void; +} + +function formatCost(value: number | null | undefined): string { + if (value == null) return "–"; + return value.toFixed(2); +} + +export function AiCreditSchemaRow({ + modelKey, + markup, + provider, + isLoading, + isCustom, + inputCost, + outputCost, + onModelChange, + onMarkupChange, + onCostChange, + onRemove, +}: AiCreditSchemaRowProps) { + const model = provider.models[modelKey]; + + const actualInput = isCustom + ? (inputCost ?? 0) + : model + ? (model.cost.input ?? 0) + : null; + const actualOutput = isCustom + ? (outputCost ?? 0) + : model + ? (model.cost.output ?? 0) + : null; + const multiplier = 1 + markup / 100; + const userInput = actualInput != null ? actualInput * multiplier : null; + const userOutput = actualOutput != null ? actualOutput * multiplier : null; + + const [localModelName, setLocalModelName] = useState(modelKey); + + return ( +
+
+ {/* Model Name */} +
+ {isCustom ? ( + setLocalModelName(e.target.value)} + onBlur={() => { + if (localModelName !== modelKey) { + onModelChange(modelKey, localModelName); + } + }} + placeholder="my-model-id" + className="w-full" + /> + ) : ( + + onModelChange(modelKey, newModelKey) + } + provider={provider} + isLoading={isLoading} + /> + )} +
+ + {/* Input Cost */} +
+ {isCustom ? ( + + onCostChange?.( + modelKey, + "input_cost", + Number(e.target.value) || 0, + ) + } + placeholder="0" + className="w-full" + /> + ) : ( +
+ {formatCost(actualInput)} +
+ )} +
+ + {/* Output Cost */} +
+ {isCustom ? ( + + onCostChange?.( + modelKey, + "output_cost", + Number(e.target.value) || 0, + ) + } + placeholder="0" + className="w-full" + /> + ) : ( +
+ {formatCost(actualOutput)} +
+ )} +
+ + {/* Markup % */} +
+ + onMarkupChange(modelKey, Number(e.target.value) || 0) + } + placeholder="0" + className="w-full" + /> +
+ + {/* Remove Button */} + } + onClick={() => onRemove(modelKey)} + className="shrink-0" + /> +
+ + {/* User Pays Info */} + {(userInput != null || userOutput != null) && ( +
+ User pays: ${formatCost(userInput)} in / ${formatCost(userOutput)} out + $/M +
+ )} +
+ ); +} diff --git a/vite/src/views/products/features/credit-systems/components/AiModelSelectDropdown.tsx b/vite/src/views/products/features/credit-systems/components/AiModelSelectDropdown.tsx new file mode 100644 index 000000000..6512dae37 --- /dev/null +++ b/vite/src/views/products/features/credit-systems/components/AiModelSelectDropdown.tsx @@ -0,0 +1,53 @@ +import type { ModelsDevModel, ModelsDevProvider } from "@autumn/shared"; +import { useMemo } from "react"; +import { SearchableSelect } from "@/components/v2/selects/SearchableSelect"; + +interface AiModelSelectDropdownProps { + value: string; + onValueChange: (modelKey: string) => void; + provider: ModelsDevProvider; + isLoading: boolean; +} + +export function AiModelSelectDropdown({ + value, + onValueChange, + provider, + isLoading, +}: AiModelSelectDropdownProps) { + const models: ModelsDevModel[] = useMemo( + () => Object.values(provider.models), + [provider], + ); + + return ( +
e.stopPropagation()}> + model.id} + getOptionLabel={(model) => model.name} + renderValue={(option) => + option ? ( + {option.name} + ) : provider.models[value]?.name ? ( + {provider.models[value].name} + ) : value ? ( + {value} + ) : ( + + {isLoading ? "Loading models..." : "Select model"} + + ) + } + placeholder={isLoading ? "Loading models..." : "Select model"} + searchable + searchPlaceholder="Search models..." + emptyText="No models found" + disabled={isLoading} + /> +
+ ); +} diff --git a/vite/src/views/products/features/credit-systems/components/ClassicCreditSchema.tsx b/vite/src/views/products/features/credit-systems/components/ClassicCreditSchema.tsx new file mode 100644 index 000000000..025640ea2 --- /dev/null +++ b/vite/src/views/products/features/credit-systems/components/ClassicCreditSchema.tsx @@ -0,0 +1,161 @@ +import type { CreateFeature, CreditSchemaItem, Feature } from "@autumn/shared"; +import { FeatureType } from "@autumn/shared"; +import { PlusIcon } from "@phosphor-icons/react"; +import { X } from "lucide-react"; +import { useEffect, useMemo, useRef } from "react"; +import { toast } from "sonner"; +import { IconButton } from "@/components/v2/buttons/IconButton"; +import { FormLabel } from "@/components/v2/form/FormLabel"; +import { Input } from "@/components/v2/inputs/Input"; +import { useFeaturesQuery } from "@/hooks/queries/useFeaturesQuery"; +import { FeatureSelectDropdown } from "./FeatureSelectDropdown"; + +interface ClassicCreditSchemaProps { + creditSystem: CreateFeature; + setCreditSystem: (creditSystem: CreateFeature) => void; +} + +export function ClassicCreditSchema({ + creditSystem, + setCreditSystem, +}: ClassicCreditSchemaProps) { + const { features } = useFeaturesQuery(); + const schema = creditSystem.config?.schema || []; + const schemaKeysRef = useRef([]); + const schemaKeys = useMemo(() => { + const nextKeys = [...schemaKeysRef.current]; + while (nextKeys.length < schema.length) { + nextKeys.push(crypto.randomUUID()); + } + while (nextKeys.length > schema.length) { + nextKeys.pop(); + } + return nextKeys; + }, [schema.length]); + + useEffect(() => { + schemaKeysRef.current = schemaKeys; + }, [schemaKeys]); + + const allMeteredFeatures = features.filter( + (feature: Feature) => feature.type === FeatureType.Metered, + ); + + const handleSchemaChange = ( + index: number, + key: keyof CreditSchemaItem, + value: string | number, + ) => { + const newSchema = [...schema]; + newSchema[index] = { ...newSchema[index], [key]: value }; + setCreditSystem({ + ...creditSystem, + config: { ...creditSystem.config, schema: newSchema }, + }); + }; + + const addSchemaItem = () => { + schemaKeysRef.current = [...schemaKeysRef.current, crypto.randomUUID()]; + const newSchema = [ + ...schema, + { + metered_feature_id: "", + feature_amount: 1, + credit_amount: 0, + }, + ]; + setCreditSystem({ + ...creditSystem, + config: { ...creditSystem.config, schema: newSchema }, + }); + }; + + const removeSchemaItem = (index: number) => { + if (schema.length === 1) { + toast.error("There must be at least one item in the credit system"); + return; + } + const nextKeys = [...schemaKeysRef.current]; + nextKeys.splice(index, 1); + schemaKeysRef.current = nextKeys; + const newSchema = [...schema]; + newSchema.splice(index, 1); + setCreditSystem({ + ...creditSystem, + config: { ...creditSystem.config, schema: newSchema }, + }); + }; + + return ( +
+
+ Metered Feature + Credit Cost +
+ +
+ {schema.map((item: CreditSchemaItem, index: number) => { + const availableFeatures = allMeteredFeatures.filter( + (feature: Feature) => + !schema.some( + (schemaItem: CreditSchemaItem) => + feature.id !== item.metered_feature_id && + schemaItem.metered_feature_id === feature.id, + ), + ); + + return ( +
+ + handleSchemaChange(index, "metered_feature_id", featureId) + } + availableFeatures={availableFeatures} + allFeatures={allMeteredFeatures} + /> + +
+ + handleSchemaChange(index, "credit_amount", e.target.value) + } + onBlur={(e) => + handleSchemaChange( + index, + "credit_amount", + Number(e.target.value) || 0, + ) + } + placeholder="eg. 10" + /> + } + onClick={() => removeSchemaItem(index)} + /> +
+
+ ); + })} +
+ + = allMeteredFeatures.length} + className="w-fit mt-4" + icon={} + > + Add + +
+ ); +} diff --git a/vite/src/views/products/features/credit-systems/components/CreditSystemSchema.tsx b/vite/src/views/products/features/credit-systems/components/CreditSystemSchema.tsx index 2cef9ff85..cb3a082fe 100644 --- a/vite/src/views/products/features/credit-systems/components/CreditSystemSchema.tsx +++ b/vite/src/views/products/features/credit-systems/components/CreditSystemSchema.tsx @@ -1,145 +1,136 @@ -import type { CreateFeature, CreditSchemaItem, Feature } from "@autumn/shared"; -import { FeatureType } from "@autumn/shared"; -import { PlusIcon } from "@phosphor-icons/react"; -import { X } from "lucide-react"; -import { toast } from "sonner"; -import { IconButton } from "@/components/v2/buttons/IconButton"; -import { FormLabel } from "@/components/v2/form/FormLabel"; -import { Input } from "@/components/v2/inputs/Input"; +import { type CreateFeature, FeatureType, type ModelsDevProvider } from "@autumn/shared"; +import { useMemo } from "react"; +import { GroupedTabButton } from "@/components/v2/buttons/GroupedTabButton"; import { SheetSection } from "@/components/v2/sheets/SharedSheetComponents"; -import { useFeaturesQuery } from "@/hooks/queries/useFeaturesQuery"; -import { FeatureSelectDropdown } from "@/views/products/features/credit-systems/components/FeatureSelectDropdown"; +import { useModelsDevPricing } from "@/hooks/queries/useAiModelsQuery"; +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, +): Record { + const result: Record = {}; + 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[`${providerKey}/${modelKey}`] = { + markup: 0, + }; + } + return result; +} interface CreditSystemSchemaProps { creditSystem: CreateFeature; setCreditSystem: (creditSystem: CreateFeature) => void; + disableModeSwitch?: boolean; } export function CreditSystemSchema({ creditSystem, setCreditSystem, + disableModeSwitch = false, }: CreditSystemSchemaProps) { - const { features } = useFeaturesQuery(); + const { providers } = useModelsDevPricing(); - const schema = creditSystem.config?.schema || []; + const mode: CreditSchemaMode = + creditSystem.type === FeatureType.AiCreditSystem ? "ai" : "classic"; - const handleSchemaChange = ( - index: number, - key: keyof CreditSchemaItem, - value: string | number, - ) => { - const newSchema = [...schema]; - newSchema[index] = { ...newSchema[index], [key]: value }; - setCreditSystem({ - ...creditSystem, - config: { ...creditSystem.config, schema: newSchema }, - }); - }; - - const addSchemaItem = () => { - const newSchema = [ - ...schema, - { - metered_feature_id: "", - feature_amount: 1, - credit_amount: 0, - }, - ]; - setCreditSystem({ - ...creditSystem, - config: { ...creditSystem.config, schema: newSchema }, - }); - }; - - const removeSchemaItem = (index: number) => { - if (schema.length === 1) { - toast.error("There must be at least one feature in the credit system"); - return; + const handleModeChange = (newMode: string) => { + if (newMode === "ai") { + const modelMarkups = getDefaultModelMarkups(providers); + setCreditSystem({ + ...creditSystem, + type: FeatureType.AiCreditSystem, + config: { ...creditSystem.config, schema: [] }, + model_markups: Object.keys(modelMarkups).length > 0 ? modelMarkups : {}, + }); + } else { + setCreditSystem({ + ...creditSystem, + type: FeatureType.CreditSystem, + config: { + ...creditSystem.config, + schema: [ + { metered_feature_id: "", feature_amount: 1, credit_amount: 0 }, + ], + }, + model_markups: null, + }); } - const newSchema = [...schema]; - newSchema.splice(index, 1); - setCreditSystem({ - ...creditSystem, - config: { ...creditSystem.config, schema: newSchema }, - }); }; - const allMeteredFeatures = features.filter( - (feature: Feature) => feature.type === FeatureType.Metered, + const modeOptions = useMemo( + () => [ + { value: "classic", label: "Classic" }, + { value: "ai", label: "AI" }, + ], + [], ); return ( -
-
- Metered Feature - Credit Cost -
+
+ {!disableModeSwitch && ( + + )} -
- {schema.map((item: CreditSchemaItem, index: number) => { - const availableFeatures = allMeteredFeatures.filter( - (feature: Feature) => - !schema.some( - (schemaItem: CreditSchemaItem) => - feature.id !== item.metered_feature_id && - schemaItem.metered_feature_id === feature.id, - ), - ); - - return ( -
- - handleSchemaChange(index, "metered_feature_id", featureId) - } - availableFeatures={availableFeatures} - allFeatures={allMeteredFeatures} - /> - -
- - handleSchemaChange(index, "credit_amount", e.target.value) - } - onBlur={(e) => - handleSchemaChange( - index, - "credit_amount", - Number(e.target.value) || 0, - ) - } - placeholder="eg. 10" - /> - } - onClick={() => removeSchemaItem(index)} - /> -
-
- ); - })} -
- - = allMeteredFeatures.length} - className="w-fit mt-4" - icon={} - > - Add - + {mode === "classic" ? ( + + ) : ( + + )}
); diff --git a/vite/src/views/products/features/credit-systems/components/UpdateCreditSystemSheet.tsx b/vite/src/views/products/features/credit-systems/components/UpdateCreditSystemSheet.tsx index 16d280699..de9b800bc 100644 --- a/vite/src/views/products/features/credit-systems/components/UpdateCreditSystemSheet.tsx +++ b/vite/src/views/products/features/credit-systems/components/UpdateCreditSystemSheet.tsx @@ -59,6 +59,7 @@ function UpdateCreditSystemSheet({ type: selectedCreditSystem.type, config: selectedCreditSystem.config, event_names: selectedCreditSystem.event_names, + model_markups: selectedCreditSystem.model_markups, }); } }, [open, selectedCreditSystem]); @@ -74,6 +75,8 @@ function UpdateCreditSystemSheet({ setLoading(true); try { + const isAiCreditSystem = creditSystem.type === FeatureType.AiCreditSystem; + await FeatureService.updateFeature( axiosInstance, selectedCreditSystem.id, @@ -81,12 +84,14 @@ function UpdateCreditSystemSheet({ id: creditSystem.id, name: creditSystem.name, type: creditSystem.type, - credit_schema: creditSystem.config?.schema?.map( - (x: CreditSchemaItem) => ({ - metered_feature_id: x.metered_feature_id, - credit_cost: Number(x.credit_amount), - }), - ), + model_markups: creditSystem.model_markups ?? undefined, + credit_schema: isAiCreditSystem + ? undefined + : creditSystem.config?.schema?.map((x: CreditSchemaItem) => ({ + metered_feature_id: x.metered_feature_id, + credit_cost: + x.credit_amount != null ? Number(x.credit_amount) : 0, + })), event_names: creditSystem.event_names, display: undefined, }, @@ -95,7 +100,6 @@ function UpdateCreditSystemSheet({ await refetch(); toast.success("Credit system updated successfully"); - // Call onSuccess with old and new IDs if (onSuccess) { onSuccess( selectedCreditSystem.id, @@ -105,7 +109,6 @@ function UpdateCreditSystemSheet({ setOpen(false); } catch (error: unknown) { - console.log(error); toast.error( getBackendErr(error as AxiosError, "Failed to update credit system"), ); @@ -134,6 +137,7 @@ function UpdateCreditSystemSheet({
diff --git a/vite/src/views/products/features/credit-systems/hooks/useAiCreditSchema.ts b/vite/src/views/products/features/credit-systems/hooks/useAiCreditSchema.ts new file mode 100644 index 000000000..0f3001eb0 --- /dev/null +++ b/vite/src/views/products/features/credit-systems/hooks/useAiCreditSchema.ts @@ -0,0 +1,259 @@ +import type { CreateFeature } from "@autumn/shared"; +import { useCallback, useEffect, useMemo, useRef, useState } from "react"; +import { toast } from "sonner"; +import { useModelsDevPricing } from "@/hooks/queries/useAiModelsQuery"; + +type ModelMarkupEntry = { + markup: number; + input_cost?: number; + output_cost?: number; +}; + +type ModelMarkups = Record; + +function groupByProvider(modelMarkups: ModelMarkups) { + const groups: Record = {}; + for (const fullId of Object.keys(modelMarkups)) { + const [providerKey] = fullId.split("/"); + if (!groups[providerKey]) groups[providerKey] = []; + groups[providerKey].push(fullId); + } + return groups; +} + +export function useAiCreditSchema({ + creditSystem, + setCreditSystem, +}: { + creditSystem: CreateFeature; + setCreditSystem: ( + creditSystem: CreateFeature | ((prev: CreateFeature) => CreateFeature), + ) => void; +}) { + const { + providers, + isLoading: modelsLoading, + error: modelsError, + } = useModelsDevPricing(); + + const modelMarkups = creditSystem.model_markups ?? {}; + + const [defaultMarkup, setDefaultMarkup] = useState(0); + const manuallyEditedModels = useRef>(new Set()); + + useEffect(() => { + if (modelsError) { + toast.error("Models.dev pricing is unavailable. Try again later."); + } + }, [modelsError]); + + const providerGroups = useMemo( + () => groupByProvider(modelMarkups), + [modelMarkups], + ); + const activeProviderKeys = Object.keys(providerGroups); + + const availableProviders = useMemo(() => { + const filtered = Object.values(providers).filter( + (provider) => !activeProviderKeys.includes(provider.id), + ); + if (!activeProviderKeys.includes("custom")) { + filtered.push({ id: "custom", name: "Custom", models: {} }); + } + return filtered; + }, [providers, activeProviderKeys]); + + const updateMarkups = useCallback( + (updatedMarkups: ModelMarkups) => { + setCreditSystem((prev) => ({ + ...prev, + model_markups: updatedMarkups, + })); + }, + [setCreditSystem], + ); + + const handleModelChange = useCallback( + (providerKey: string, oldModelKey: string, newModelKey: string) => { + const oldFullId = `${providerKey}/${oldModelKey}`; + const newFullId = `${providerKey}/${newModelKey}`; + if (oldFullId !== newFullId && newFullId in modelMarkups) return; + const updatedMarkups = { ...modelMarkups }; + const oldEntry = updatedMarkups[oldFullId]; + const markup = oldEntry?.markup ?? 0; + + if (manuallyEditedModels.current.has(oldFullId)) { + manuallyEditedModels.current.delete(oldFullId); + manuallyEditedModels.current.add(newFullId); + } + + delete updatedMarkups[oldFullId]; + if (providerKey === "custom") { + updatedMarkups[newFullId] = { + markup, + input_cost: oldEntry?.input_cost ?? 0, + output_cost: oldEntry?.output_cost ?? 0, + }; + } else { + updatedMarkups[newFullId] = { markup }; + } + updateMarkups(updatedMarkups); + }, + [modelMarkups, providers, updateMarkups], + ); + + const handleMarkupChange = useCallback( + (providerKey: string, modelKey: string, markup: number) => { + const fullId = `${providerKey}/${modelKey}`; + manuallyEditedModels.current.add(fullId); + updateMarkups({ + ...modelMarkups, + [fullId]: { ...modelMarkups[fullId], markup }, + }); + }, + [modelMarkups, updateMarkups], + ); + + const handleDefaultMarkupChange = useCallback( + (value: number) => { + setDefaultMarkup(value); + const updatedMarkups = { ...modelMarkups }; + for (const modelId of Object.keys(updatedMarkups)) { + if (!manuallyEditedModels.current.has(modelId)) { + updatedMarkups[modelId] = { + ...updatedMarkups[modelId], + markup: value, + }; + } + } + updateMarkups(updatedMarkups); + }, + [modelMarkups, updateMarkups], + ); + + const handleCostChange = useCallback( + ( + providerKey: string, + modelKey: string, + field: "input_cost" | "output_cost", + value: number, + ) => { + const fullId = `${providerKey}/${modelKey}`; + updateMarkups({ + ...modelMarkups, + [fullId]: { ...modelMarkups[fullId], [field]: value }, + }); + }, + [modelMarkups, updateMarkups], + ); + + const handleRemoveModel = useCallback( + (providerKey: string, modelKey: string) => { + const fullId = `${providerKey}/${modelKey}`; + manuallyEditedModels.current.delete(fullId); + const updatedMarkups = { ...modelMarkups }; + delete updatedMarkups[fullId]; + updateMarkups(updatedMarkups); + }, + [modelMarkups, updateMarkups], + ); + + const handleRemoveProvider = useCallback( + (providerKey: string) => { + const updatedMarkups = { ...modelMarkups }; + for (const fullId of providerGroups[providerKey] ?? []) { + manuallyEditedModels.current.delete(fullId); + delete updatedMarkups[fullId]; + } + updateMarkups(updatedMarkups); + }, + [modelMarkups, providerGroups, updateMarkups], + ); + + const addProvider = useCallback( + (providerKey: string) => { + if (providerKey === "custom") { + const existingKeys = (providerGroups.custom ?? []).map((fullId) => { + const [, ...parts] = fullId.split("/"); + return parts.join("/"); + }); + let counter = 1; + while (existingKeys.includes(`model-${counter}`)) counter++; + updateMarkups({ + ...modelMarkups, + [`custom/model-${counter}`]: { + markup: defaultMarkup, + input_cost: 0, + output_cost: 0, + }, + }); + return; + } + const provider = providers[providerKey]; + if (!provider) return; + const firstModelKey = Object.keys(provider.models)[0]; + if (!firstModelKey) return; + const fullId = `${providerKey}/${firstModelKey}`; + updateMarkups({ + ...modelMarkups, + [fullId]: { markup: defaultMarkup }, + }); + }, + [defaultMarkup, modelMarkups, providerGroups, providers, updateMarkups], + ); + + const addModelToProvider = useCallback( + (providerKey: string) => { + if (providerKey === "custom") { + const existingKeys = (providerGroups.custom ?? []).map((fullId) => { + const [, ...parts] = fullId.split("/"); + return parts.join("/"); + }); + let counter = 1; + while (existingKeys.includes(`model-${counter}`)) counter++; + const fullId = `custom/model-${counter}`; + updateMarkups({ + ...modelMarkups, + [fullId]: { markup: defaultMarkup, input_cost: 0, output_cost: 0 }, + }); + return; + } + const provider = providers[providerKey]; + if (!provider) return; + const usedModelKeys = new Set( + (providerGroups[providerKey] ?? []).map((fullId) => { + const [, ...parts] = fullId.split("/"); + return parts.join("/"); + }), + ); + const nextModelKey = Object.keys(provider.models).find( + (key) => !usedModelKeys.has(key), + ); + if (!nextModelKey) return; + const fullId = `${providerKey}/${nextModelKey}`; + updateMarkups({ + ...modelMarkups, + [fullId]: { markup: defaultMarkup }, + }); + }, + [defaultMarkup, modelMarkups, providerGroups, providers, updateMarkups], + ); + + return { + providers, + modelsLoading, + modelMarkups, + defaultMarkup, + providerGroups, + activeProviderKeys, + availableProviders, + handleModelChange, + handleMarkupChange, + handleDefaultMarkupChange, + handleCostChange, + handleRemoveModel, + handleRemoveProvider, + addProvider, + addModelToProvider, + }; +} diff --git a/vite/src/views/products/features/credit-systems/utils/validateCreditSystem.ts b/vite/src/views/products/features/credit-systems/utils/validateCreditSystem.ts index d470bc254..32c27e890 100644 --- a/vite/src/views/products/features/credit-systems/utils/validateCreditSystem.ts +++ b/vite/src/views/products/features/credit-systems/utils/validateCreditSystem.ts @@ -1,4 +1,4 @@ -import type { CreateFeature } from "@autumn/shared"; +import { type CreateFeature, FeatureType } from "@autumn/shared"; export const validateCreditSystem = ( creditSystem: CreateFeature, @@ -7,16 +7,36 @@ export const validateCreditSystem = ( return "Please fill in all fields"; } - if (creditSystem.config.schema.length === 0) { - return "Need at least one metered feature"; + const isAiCreditSystem = creditSystem.type === FeatureType.AiCreditSystem; + + if (isAiCreditSystem) { + if ( + !creditSystem.model_markups || + Object.keys(creditSystem.model_markups).length === 0 + ) + return "Add at least one model markup"; + + for (const [modelId, entry] of Object.entries(creditSystem.model_markups)) { + if (!modelId) return "Select a model for each row"; + if (modelId.startsWith("custom/")) { + const customModelKey = modelId.slice("custom/".length); + if (!customModelKey) return "Custom model ID cannot be empty"; + if (entry.input_cost == null || entry.output_cost == null) + return "Custom models require input and output costs"; + } + } + return null; + } + + if (!creditSystem.config?.schema || creditSystem.config.schema.length === 0) { + return "Need at least one item in the schema"; } for (const item of creditSystem.config.schema) { if (!item.metered_feature_id) { - return "Select a metered feature"; + return "Select a feature for each row"; } - - if (item.feature_amount <= 0 || item.credit_amount <= 0) { + if ((item.credit_amount ?? 0) <= 0) { return "Credit amount must be greater than 0"; } } diff --git a/vite/src/views/products/features/feature-list/CreditListColumns.tsx b/vite/src/views/products/features/feature-list/CreditListColumns.tsx index b6955c648..fa94fbc2f 100644 --- a/vite/src/views/products/features/feature-list/CreditListColumns.tsx +++ b/vite/src/views/products/features/feature-list/CreditListColumns.tsx @@ -1,11 +1,22 @@ -import type { Feature } from "@autumn/shared"; +import type { Feature, ModelsDevProvider } from "@autumn/shared"; import type { ColumnDef, Row } from "@tanstack/react-table"; import { AdminHover } from "@/components/general/AdminHover"; import { MiniCopyButton } from "@/components/v2/buttons/CopyButton"; import { getFeatureHoverTexts } from "@/views/admin/adminUtils"; import { FeatureListRowToolbar } from "./FeatureListRowToolbar"; -export const createCreditListColumns = (): ColumnDef[] => [ +function resolveModelName( + fullId: string, + providers: Record, +): string { + const [providerKey, ...modelParts] = fullId.split("/"); + const modelKey = modelParts.join("/"); + return providers[providerKey]?.models[modelKey]?.name ?? fullId; +} + +export const createCreditListColumns = ( + providers: Record, +): ColumnDef[] => [ { size: 150, header: "Name", @@ -43,13 +54,20 @@ export const createCreditListColumns = (): ColumnDef[] => [ accessorKey: "features", cell: ({ row }: { row: Row }) => { const creditSystem = row.original; + const modelMarkupEntries = creditSystem.model_markups + ? Object.entries(creditSystem.model_markups) + : null; const featureIds = - creditSystem.config?.schema - ?.map( - (schema: { metered_feature_id: string }) => - schema.metered_feature_id, - ) - .join(", ") || "—"; + modelMarkupEntries && modelMarkupEntries.length > 0 + ? modelMarkupEntries + .map(([fullId]) => resolveModelName(fullId, providers)) + .join(", ") + : creditSystem.config?.schema + ?.map( + (schema: { metered_feature_id: string }) => + schema.metered_feature_id, + ) + .join(", ") || "—"; return (
{featureIds}
); diff --git a/vite/src/views/products/features/feature-list/FeatureListTable.tsx b/vite/src/views/products/features/feature-list/FeatureListTable.tsx index a9aca0662..4bddbba2a 100644 --- a/vite/src/views/products/features/feature-list/FeatureListTable.tsx +++ b/vite/src/views/products/features/feature-list/FeatureListTable.tsx @@ -4,6 +4,7 @@ import { useMemo, useState } from "react"; import { Table } from "@/components/general/table"; import { IconButton } from "@/components/v2/buttons/IconButton"; import { EmptyState } from "@/components/v2/empty-states/EmptyState"; +import { useModelsDevPricing } from "@/hooks/queries/useAiModelsQuery"; import { useFeaturesQuery } from "@/hooks/queries/useFeaturesQuery"; import { useEnv } from "@/utils/envUtils"; import { useProductsQueryState } from "@/views/products/hooks/useProductsQueryState"; @@ -18,6 +19,7 @@ import { FeatureListMenuButton } from "./FeatureListMenuButton"; export function FeatureListTable() { const env = useEnv(); const { features } = useFeaturesQuery(); + const { providers } = useModelsDevPricing(); const { queryStates } = useProductsQueryState(); const [selectedFeature, setSelectedFeature] = useState(null); const [updateFeatureOpen, setUpdateFeatureOpen] = useState(false); @@ -52,7 +54,10 @@ export function FeatureListTable() { () => createFeatureListColumns({ showEventNames: hasEventNames }), [hasEventNames], ); - const creditColumns = useMemo(() => createCreditListColumns(), []); + const creditColumns = useMemo( + () => createCreditListColumns(providers), + [providers], + ); const featureTable = useProductTable({ data: regularFeatures || [], diff --git a/vite/src/views/products/plan/components/edit-plan-feature/BillingType.tsx b/vite/src/views/products/plan/components/edit-plan-feature/BillingType.tsx index 91d08183f..165307e2e 100644 --- a/vite/src/views/products/plan/components/edit-plan-feature/BillingType.tsx +++ b/vite/src/views/products/plan/components/edit-plan-feature/BillingType.tsx @@ -1,5 +1,6 @@ import { BillingInterval, + FeatureType, FeatureUsageType, getFeatureName, Infinite, @@ -73,6 +74,7 @@ export function BillingType() { }; const feature = features.find((f) => f.id === item.feature_id); + const isAiCreditSystem = feature?.type === FeatureType.AiCreditSystem; const featureName = getFeatureName({ feature, @@ -104,7 +106,9 @@ export function BillingType() {
Included
- {isConsumable + {isAiCreditSystem + ? "Set an included USD budget (eg, $10 per month)." + : isConsumable ? `Set an included usage limit (eg, 100 ${featureName} per month).` : isAllocated ? `Set a usage limit (eg, 5 ${featureName}).` @@ -124,7 +128,9 @@ export function BillingType() {
Priced
- {isConsumable + {isAiCreditSystem + ? "Bill model usage at the markup you set in USD after the included budget is used." + : isConsumable ? `Charge a price for usage (eg, $0.05 per ${singleFeatureName}).` : isAllocated ? `Charge a price based on usage (eg, $10 per ${singleFeatureName}).` diff --git a/vite/src/views/products/plan/components/edit-plan-feature/EditPlanFeatureSheet.tsx b/vite/src/views/products/plan/components/edit-plan-feature/EditPlanFeatureSheet.tsx index 95b2d427d..c6563308d 100644 --- a/vite/src/views/products/plan/components/edit-plan-feature/EditPlanFeatureSheet.tsx +++ b/vite/src/views/products/plan/components/edit-plan-feature/EditPlanFeatureSheet.tsx @@ -116,6 +116,7 @@ export function EditPlanFeatureSheet({ const feature = getFeature(item?.feature_id ?? "", features); const isFeaturePrice = isFeaturePriceItem(item); + const isAiCreditSystem = feature?.type === FeatureType.AiCreditSystem; // Allow confirming a priced feature that has a $0 tier (valid zero-price config) const isZeroPriceItem = @@ -166,7 +167,7 @@ export function EditPlanFeatureSheet({ - {isFeaturePrice && ( + {isFeaturePrice && !isAiCreditSystem && ( 1 ? ( diff --git a/vite/src/views/products/plan/components/edit-plan-feature/IncludedUsage.tsx b/vite/src/views/products/plan/components/edit-plan-feature/IncludedUsage.tsx index 1daaa82fd..8b4528277 100644 --- a/vite/src/views/products/plan/components/edit-plan-feature/IncludedUsage.tsx +++ b/vite/src/views/products/plan/components/edit-plan-feature/IncludedUsage.tsx @@ -1,18 +1,24 @@ +import { IconCheckbox } from "@/components/v2/checkboxes/IconCheckbox"; +import { Input } from "@/components/v2/inputs/Input"; +import { + InputGroup, + InputGroupInput, + InputGroupText, +} from "@/components/v2/inputs/InputGroup"; +import { useFeaturesQuery } from "@/hooks/queries/useFeaturesQuery"; +import { isFeaturePriceItem } from "@/utils/product/getItemType"; +import { useProductItemContext } from "@/views/products/product/product-item/ProductItemContext"; import { BillingInterval, billingToItemInterval, EntInterval, entToItemInterval, + FeatureType, getFeatureName, Infinite, isContUseItem, } from "@autumn/shared"; import { InfinityIcon } from "@phosphor-icons/react"; -import { IconCheckbox } from "@/components/v2/checkboxes/IconCheckbox"; -import { Input } from "@/components/v2/inputs/Input"; -import { useFeaturesQuery } from "@/hooks/queries/useFeaturesQuery"; -import { isFeaturePriceItem } from "@/utils/product/getItemType"; -import { useProductItemContext } from "@/views/products/product/product-item/ProductItemContext"; import { UsageReset } from "./UsageReset"; export function IncludedUsage() { @@ -24,6 +30,8 @@ export function IncludedUsage() { const includedUsage = item.included_usage; const isFeaturePrice = isFeaturePriceItem(item); + const feature = features.find((f) => f.id === item.feature_id); + const isAiCreditSystem = feature?.type === FeatureType.AiCreditSystem; // Helper function to get the display value for the input const getInputValue = () => { @@ -41,35 +49,71 @@ export function IncludedUsage() {
- Quantity of  - - {getFeatureName({ - feature: features.find((f) => f.id === item.feature_id), - plural: true, - })}{" "} - - {!isFeaturePrice ? " that can be used" : " granted before billing"} + {isAiCreditSystem ? ( + <> + USD budget{" "} + {!isFeaturePrice + ? "allocated to this plan" + : "granted before billing"} + + ) : ( + <> + Quantity of  + + {getFeatureName({ feature, plural: true })}{" "} + + {!isFeaturePrice + ? " that can be used" + : " granted before billing"} + + )}
- { - const value = e.target.value.trim(); + {isAiCreditSystem ? ( + // Couldn't find a disabled property but data-disabled is accounted for in CSS + + $ + { + const value = e.target.value.trim(); - if (value === "") { - setItem({ ...item, included_usage: null }); - } else { - const numValue = value; - if (!Number.isNaN(numValue)) { - setItem({ ...item, included_usage: Number(numValue) }); + if (value === "") { + setItem({ ...item, included_usage: null }); + } else { + const numValue = Number(value); + if (!Number.isNaN(numValue)) { + setItem({ ...item, included_usage: numValue }); + } + } + }} + disabled={includedUsage === Infinite} + type="number" + /> + + ) : ( + { + const value = e.target.value.trim(); + + if (value === "") { + setItem({ ...item, included_usage: null }); + } else { + const numValue = Number(value); + if (!Number.isNaN(numValue)) { + setItem({ ...item, included_usage: numValue }); + } } - } - }} - disabled={includedUsage === Infinite} - type="number" - /> + }} + disabled={includedUsage === Infinite} + type="number" + /> + )} } diff --git a/vite/src/views/products/plan/components/new-feature/NewFeatureBehaviour.tsx b/vite/src/views/products/plan/components/new-feature/NewFeatureBehaviour.tsx index a3cba3ac6..3626ed9e5 100644 --- a/vite/src/views/products/plan/components/new-feature/NewFeatureBehaviour.tsx +++ b/vite/src/views/products/plan/components/new-feature/NewFeatureBehaviour.tsx @@ -17,7 +17,10 @@ export function NewFeatureBehaviour({ }) { if (feature.type === FeatureType.CreditSystem) { return ( - + ); } diff --git a/vite/src/views/products/plan/components/plan-card/DummyPlanFeatureRow.tsx b/vite/src/views/products/plan/components/plan-card/DummyPlanFeatureRow.tsx index 13a045e96..7bc9cb3ae 100644 --- a/vite/src/views/products/plan/components/plan-card/DummyPlanFeatureRow.tsx +++ b/vite/src/views/products/plan/components/plan-card/DummyPlanFeatureRow.tsx @@ -44,10 +44,16 @@ export const DummyPlanFeatureRow = () => { return "Chat Messages"; }; + const isAiCreditSystem = feature.type === FeatureType.AiCreditSystem; + // Build display text based on feature type const getDisplayText = () => { const name = hasName ? featureName : getPlaceholderName(); + if (isAiCreditSystem) { + return { primary: `$10.00 of ${name}`, secondary: "" }; + } + if (featureType === FeatureType.CreditSystem) { return { primary: `100 ${name}`, secondary: "" }; } diff --git a/vite/src/views/products/plan/components/plan-card/PlanFeatureRow.tsx b/vite/src/views/products/plan/components/plan-card/PlanFeatureRow.tsx index 121e52f27..8961e88dc 100644 --- a/vite/src/views/products/plan/components/plan-card/PlanFeatureRow.tsx +++ b/vite/src/views/products/plan/components/plan-card/PlanFeatureRow.tsx @@ -1,6 +1,6 @@ /** biome-ignore-all lint/a11y/noStaticElementInteractions: needed */ /** biome-ignore-all lint/a11y/useSemanticElements: needed */ -import type { ProductItem } from "@autumn/shared"; +import { type ProductItem, FeatureType } from "@autumn/shared"; import { getProductItemDisplay } from "@autumn/shared"; import { TrashIcon } from "@phosphor-icons/react"; import { useEffect, useRef, useState } from "react"; @@ -68,6 +68,8 @@ export const PlanFeatureRow = ({ const feature = features.find((f) => f.id === item.feature_id); const hasFeatureName = feature?.name && feature.name.trim() !== ""; + const isAiCreditSystem = feature?.type === FeatureType.AiCreditSystem; + const displayText = hasFeatureName ? display.primary_text : "Name your feature"; @@ -209,7 +211,9 @@ export const PlanFeatureRow = ({ {displayText} - {display.secondary_text} + {!isAiCreditSystem && display.secondary_text && ( + {display.secondary_text} + )}

Date: Wed, 20 May 2026 10:07:54 +0100 Subject: [PATCH 02/46] chore: cleanup handleTrackTokens Co-authored-by: TheUntraceable <73362400+TheUntraceable@users.noreply.github.com> --- .../balances/handlers/handleTrackTokens.ts | 112 +---------------- .../track/utils/getTokenTrackParams.ts | 113 ++++++++++++++++++ 2 files changed, 118 insertions(+), 107 deletions(-) create mode 100644 server/src/internal/balances/track/utils/getTokenTrackParams.ts diff --git a/server/src/internal/balances/handlers/handleTrackTokens.ts b/server/src/internal/balances/handlers/handleTrackTokens.ts index b0f6df62f..024baf635 100644 --- a/server/src/internal/balances/handlers/handleTrackTokens.ts +++ b/server/src/internal/balances/handlers/handleTrackTokens.ts @@ -1,62 +1,7 @@ import { createRoute } from "@/honoMiddlewares/routeHandler.js"; import { runTrackV2 } from "@/internal/balances/track/runTrackV2.js"; -import { getCreditCost } from "@/internal/features/creditSystemUtils.js"; -import { - AffectedResource, - ErrCode, - type Feature, - FeatureType, - RecaseError, - Scopes, - type TrackParams, - TrackTokensParamsSchema, -} from "@autumn/shared"; -import type { FeatureDeduction } from "../utils/types/featureDeduction.js"; - -const resolveAiCreditFeature = ({ - features, - featureId, -}: { - features: Feature[]; - featureId?: string; -}): Feature => { - if (featureId) { - const candidate = features.find((f) => f.id === featureId); - if (!candidate) { - throw new RecaseError({ - message: `Feature ${featureId} not found`, - code: ErrCode.FeatureNotFound, - statusCode: 404, - }); - } - if (candidate.type !== FeatureType.AiCreditSystem) { - throw new RecaseError({ - message: `Feature ${featureId} is not an AI credit system feature`, - code: ErrCode.InvalidRequest, - statusCode: 400, - }); - } - return candidate; - } - - const matches = features.filter((f) => f.type === FeatureType.AiCreditSystem); - if (matches.length === 0) { - throw new RecaseError({ - message: "No AI credit system feature found for this organization", - code: ErrCode.FeatureNotFound, - statusCode: 404, - }); - } - if (matches.length > 1) { - throw new RecaseError({ - message: - "Multiple AI credit system features found. Please specify a feature_id to disambiguate.", - code: ErrCode.InvalidRequest, - statusCode: 400, - }); - } - return matches[0]; -}; +import { getTokenTrackParams } from "@/internal/balances/track/utils/getTokenTrackParams.js"; +import { AffectedResource, Scopes, TrackTokensParamsSchema } from "@autumn/shared"; export const handleTrackTokens = createRoute({ scopes: [Scopes.Balances.Write], @@ -66,58 +11,11 @@ export const handleTrackTokens = createRoute({ const body = c.req.valid("json"); const ctx = c.get("ctx"); - const aiCreditFeature = resolveAiCreditFeature({ - features: ctx.features, - featureId: body.feature_id, + const { body: trackBody, featureDeductions } = await getTokenTrackParams({ + ctx, + input: body, }); - const rawModelName = body.model_id; - - // Compute the dollar cost once and reuse it for both the response/event row - // and the deduction layer (avoids a second getCreditCost call per entitlement). - const cost = await getCreditCost({ - featureId: aiCreditFeature.id, - creditSystem: aiCreditFeature, - modelName: rawModelName, - tokens: { - input: body.input_tokens, - output: body.output_tokens, - }, - }); - - const featureDeductions: FeatureDeduction[] = [ - { - feature: aiCreditFeature, - deduction: 1, // multiplied by per-entitlement credit_cost in the deduction layer (Postgres) - tokenUsage: { - modelName: rawModelName, - inputTokens: body.input_tokens, - outputTokens: body.output_tokens, - }, - precomputedCreditCost: cost, - }, - ]; - - // Build TrackParams body — store model/tokens in properties for audit - const trackBody: TrackParams = { - customer_id: body.customer_id, - entity_id: body.entity_id, - feature_id: aiCreditFeature.id, - value: cost, - properties: { - ...body.properties, - model: rawModelName, - input_tokens: body.input_tokens, - output_tokens: body.output_tokens, - cost, - }, - idempotency_key: body.idempotency_key, - overage_behavior: body.overage_behavior, - customer_data: body.customer_data, - entity_data: body.entity_data, - skip_event: body.skip_event, - }; - return c.json( await runTrackV2({ ctx, diff --git a/server/src/internal/balances/track/utils/getTokenTrackParams.ts b/server/src/internal/balances/track/utils/getTokenTrackParams.ts new file mode 100644 index 000000000..4e9aef820 --- /dev/null +++ b/server/src/internal/balances/track/utils/getTokenTrackParams.ts @@ -0,0 +1,113 @@ +import { + ErrCode, + type Feature, + FeatureType, + RecaseError, + type TrackParams, + type TrackTokensParams, +} from "@autumn/shared"; +import type { AutumnContext } from "@/honoUtils/HonoEnv.js"; +import { getCreditCost } from "@/internal/features/creditSystemUtils.js"; +import type { FeatureDeduction } from "../../utils/types/featureDeduction.js"; + +const resolveAiCreditFeature = ({ + features, + featureId, +}: { + features: Feature[]; + featureId?: string; +}): Feature => { + if (featureId) { + const candidate = features.find((f) => f.id === featureId); + if (!candidate) { + throw new RecaseError({ + message: `Feature ${featureId} not found`, + code: ErrCode.FeatureNotFound, + statusCode: 404, + }); + } + if (candidate.type !== FeatureType.AiCreditSystem) { + throw new RecaseError({ + message: `Feature ${featureId} is not an AI credit system feature`, + code: ErrCode.InvalidRequest, + statusCode: 400, + }); + } + return candidate; + } + + const matches = features.filter((f) => f.type === FeatureType.AiCreditSystem); + if (matches.length === 0) { + throw new RecaseError({ + message: "No AI credit system feature found for this organization", + code: ErrCode.FeatureNotFound, + statusCode: 404, + }); + } + if (matches.length > 1) { + throw new RecaseError({ + message: + "Multiple AI credit system features found. Please specify a feature_id to disambiguate.", + code: ErrCode.InvalidRequest, + statusCode: 400, + }); + } + return matches[0]; +}; + +export const getTokenTrackParams = async ({ + ctx, + input, +}: { + ctx: AutumnContext; + input: TrackTokensParams; +}): Promise<{ body: TrackParams; featureDeductions: FeatureDeduction[] }> => { + const aiCreditFeature = resolveAiCreditFeature({ + features: ctx.features, + featureId: input.feature_id, + }); + + const cost = await getCreditCost({ + featureId: aiCreditFeature.id, + creditSystem: aiCreditFeature, + modelName: input.model_id, + tokens: { + input: input.input_tokens, + output: input.output_tokens, + }, + }); + + const featureDeductions: FeatureDeduction[] = [ + { + feature: aiCreditFeature, + deduction: 1, + tokenUsage: { + modelName: input.model_id, + inputTokens: input.input_tokens, + outputTokens: input.output_tokens, + }, + precomputedCreditCost: cost, + }, + ]; + + const body: TrackParams = { + customer_id: input.customer_id, + entity_id: input.entity_id, + feature_id: aiCreditFeature.id, + value: cost, + properties: { + ...input.properties, + model: input.model_id, + input_tokens: input.input_tokens, + output_tokens: input.output_tokens, + cost, + }, + idempotency_key: input.idempotency_key, + overage_behavior: input.overage_behavior, + customer_data: input.customer_data, + entity_data: input.entity_data, + skip_event: input.skip_event, + }; + + return { body, featureDeductions }; +}; From 09a82d37a51d2523918ed8488b9053bf507e520f Mon Sep 17 00:00:00 2001 From: Charlie Lamb Date: Wed, 20 May 2026 10:35:36 +0100 Subject: [PATCH 03/46] chore: credit system logic cleanup Co-authored-by: TheUntraceable <73362400+TheUntraceable@users.noreply.github.com> --- .../src/internal/balances/balancesRouter.ts | 4 +- .../utils/deduction/computeCreditCosts.ts | 46 +++++++++++++++++++ .../deduction/prepareFeatureDeduction.ts | 29 ++---------- .../deductionV2/prepareFeatureDeductionV2.ts | 33 +++---------- .../features/featureActions/createFeature.ts | 6 +-- .../features/featureActions/updateFeature.ts | 21 +++++---- server/src/internal/features/featureUtils.ts | 9 ++-- .../balances/track/basic/track-tokens.test.ts | 18 ++++---- 8 files changed, 86 insertions(+), 80 deletions(-) create mode 100644 server/src/internal/balances/utils/deduction/computeCreditCosts.ts diff --git a/server/src/internal/balances/balancesRouter.ts b/server/src/internal/balances/balancesRouter.ts index 95d6c54ce..fd6fa96e7 100644 --- a/server/src/internal/balances/balancesRouter.ts +++ b/server/src/internal/balances/balancesRouter.ts @@ -20,7 +20,7 @@ balancesRouter.post("/balances/update", ...handleUpdateBalance); // Track balancesRouter.post("/events", ...handleTrack); balancesRouter.post("/track", ...handleTrack); -balancesRouter.post("/trackTokens", ...handleTrackTokens); +balancesRouter.post("/track_tokens", ...handleTrackTokens); // Check balancesRouter.post("/entitled", ...handleCheck); @@ -35,6 +35,6 @@ balancesRpcRouter.post("/balances.update", ...handleUpdateBalance); balancesRpcRouter.post("/balances.delete", ...handleDeleteBalance); balancesRpcRouter.post("/balances.track", ...handleTrack); -balancesRpcRouter.post("/balances.trackTokens", ...handleTrackTokens); +balancesRpcRouter.post("/balances.track_tokens", ...handleTrackTokens); balancesRpcRouter.post("/balances.check", ...handleCheck); balancesRpcRouter.post("/balances.finalize", ...handleFinalizeLock); diff --git a/server/src/internal/balances/utils/deduction/computeCreditCosts.ts b/server/src/internal/balances/utils/deduction/computeCreditCosts.ts new file mode 100644 index 000000000..b6c203f48 --- /dev/null +++ b/server/src/internal/balances/utils/deduction/computeCreditCosts.ts @@ -0,0 +1,46 @@ +import type { FullCusEntWithFullCusProduct } from "@autumn/shared"; +import { getCreditCost } from "@/internal/features/creditSystemUtils.js"; +import type { FeatureDeduction } from "../types/featureDeduction.js"; + +const DEFAULT_CREDIT_COST = 1; + +export type CreditCostLookup = (entitlementId: string) => number; + +/** + * Computes the credit cost for each customer entitlement and returns a lookup + * function. Uses precomputedCreditCost when available (token tracking), + * otherwise calls getCreditCost per entitlement (credit system schema lookups). + */ +export const computeCreditCosts = async ({ + cusEnts, + deduction, +}: { + cusEnts: FullCusEntWithFullCusProduct[]; + deduction: FeatureDeduction; +}): Promise => { + const costMap = new Map(); + + if (deduction.precomputedCreditCost != null) { + const cost = deduction.precomputedCreditCost; + return () => cost; + } + + await Promise.all( + cusEnts.map(async (ce) => { + const creditCost = await getCreditCost({ + featureId: deduction.feature.id, + creditSystem: ce.entitlement.feature, + modelName: deduction.tokenUsage?.modelName, + tokens: deduction.tokenUsage + ? { + input: deduction.tokenUsage.inputTokens, + output: deduction.tokenUsage.outputTokens, + } + : undefined, + }); + costMap.set(ce.id, creditCost); + }), + ); + + return (entitlementId) => costMap.get(entitlementId) ?? DEFAULT_CREDIT_COST; +}; diff --git a/server/src/internal/balances/utils/deduction/prepareFeatureDeduction.ts b/server/src/internal/balances/utils/deduction/prepareFeatureDeduction.ts index 1d4cba8f3..a03e6f3b5 100644 --- a/server/src/internal/balances/utils/deduction/prepareFeatureDeduction.ts +++ b/server/src/internal/balances/utils/deduction/prepareFeatureDeduction.ts @@ -15,7 +15,7 @@ import { import type { AutumnContext } from "@/honoUtils/HonoEnv.js"; import { buildLockReceiptKey } from "@/internal/balances/utils/lock/buildLockReceiptKey.js"; import { getUnlimitedAndUsageAllowed } from "@/internal/customers/cusProducts/cusEnts/cusEntUtils.js"; -import { getCreditCost } from "@/internal/features/creditSystemUtils.js"; +import { computeCreditCosts } from "./computeCreditCosts.js"; import type { CustomerEntitlementDeduction, DeductionOptions, @@ -68,7 +68,7 @@ export const prepareFeatureDeduction = async ({ for (const rf of relevantFeatures) { const { unlimited: featureUnlimited } = getUnlimitedAndUsageAllowed({ cusEnts, - internalFeatureId: rf.internal_id!, + internalFeatureId: rf.internal_id, }); if (featureUnlimited) { @@ -101,31 +101,12 @@ export const prepareFeatureDeduction = async ({ .map((ce) => ce.entitlement.feature.id), ); - // Compute credit cost once per customer entitlement - const creditCostByEntitlementId = new Map(); - await Promise.all( - cusEnts.map(async (ce) => { - const creditCost = - deduction.precomputedCreditCost ?? - (await getCreditCost({ - featureId: feature.id, - creditSystem: ce.entitlement.feature, - modelName: deduction.tokenUsage?.modelName, - tokens: deduction.tokenUsage - ? { - input: deduction.tokenUsage.inputTokens, - output: deduction.tokenUsage.outputTokens, - } - : undefined, - })); - creditCostByEntitlementId.set(ce.id, creditCost); - }), - ); + const getCreditCostForEnt = await computeCreditCosts({ cusEnts, deduction }); // Build input for each customer entitlement const customerEntitlementDeductions: CustomerEntitlementDeduction[] = cusEnts.map((ce) => { - const creditCost = creditCostByEntitlementId.get(ce.id)!; + const creditCost = getCreditCostForEnt(ce.id); const maxOverage = getMaxOverage({ cusEnt: ce }); const isFreeAllocated = @@ -165,7 +146,7 @@ export const prepareFeatureDeduction = async ({ // Collect and sort rollovers by expires_at (oldest first), including credit_cost from parent entitlement const sortedRollovers = cusEnts .flatMap((ce) => { - const creditCost = creditCostByEntitlementId.get(ce.id)!; + const creditCost = getCreditCostForEnt(ce.id); return (ce.rollovers || []).map((r) => ({ ...r, credit_cost: creditCost, diff --git a/server/src/internal/balances/utils/deductionV2/prepareFeatureDeductionV2.ts b/server/src/internal/balances/utils/deductionV2/prepareFeatureDeductionV2.ts index 44147d7ba..1d09dd5a7 100644 --- a/server/src/internal/balances/utils/deductionV2/prepareFeatureDeductionV2.ts +++ b/server/src/internal/balances/utils/deductionV2/prepareFeatureDeductionV2.ts @@ -17,7 +17,7 @@ import { import type { AutumnContext } from "@/honoUtils/HonoEnv.js"; import { buildLockReceiptKey } from "@/internal/balances/utils/lock/buildLockReceiptKey.js"; import { getUnlimitedAndUsageAllowed } from "@/internal/customers/cusProducts/cusEnts/cusEntUtils.js"; -import { getCreditCost } from "@/internal/features/creditSystemUtils.js"; +import { computeCreditCosts } from "../deduction/computeCreditCosts.js"; import type { CustomerEntitlementDeduction, DeductionOptions, @@ -115,31 +115,14 @@ export const prepareFeatureDeductionV2 = async ({ .map((customerEntitlement) => customerEntitlement.entitlement.feature.id), ); - const creditCostByCustomerEntitlementId = new Map(); - await Promise.all( - customerEntitlements.map(async (customerEntitlement) => { - const creditCost = - deduction.precomputedCreditCost ?? - (await getCreditCost({ - featureId: feature.id, - creditSystem: customerEntitlement.entitlement.feature, - modelName: deduction.tokenUsage?.modelName, - tokens: deduction.tokenUsage - ? { - input: deduction.tokenUsage.inputTokens, - output: deduction.tokenUsage.outputTokens, - } - : undefined, - })); - creditCostByCustomerEntitlementId.set(customerEntitlement.id, creditCost); - }), - ); + const getCreditCostForEnt = await computeCreditCosts({ + cusEnts: customerEntitlements, + deduction, + }); const customerEntitlementDeductions: CustomerEntitlementDeduction[] = customerEntitlements.map((customerEntitlement) => { - const creditCost = creditCostByCustomerEntitlementId.get( - customerEntitlement.id, - )!; + const creditCost = getCreditCostForEnt(customerEntitlement.id); const maxOverage = getMaxOverage({ cusEnt: customerEntitlement, @@ -182,9 +165,7 @@ export const prepareFeatureDeductionV2 = async ({ }); const rolloverArrays = customerEntitlements.map((customerEntitlement) => { - const creditCost = creditCostByCustomerEntitlementId.get( - customerEntitlement.id, - )!; + const creditCost = getCreditCostForEnt(customerEntitlement.id); return (customerEntitlement.rollovers || []).map((rollover) => ({ ...rollover, credit_cost: creditCost, diff --git a/server/src/internal/features/featureActions/createFeature.ts b/server/src/internal/features/featureActions/createFeature.ts index 1a1bbc1b2..a2f9ae88a 100644 --- a/server/src/internal/features/featureActions/createFeature.ts +++ b/server/src/internal/features/featureActions/createFeature.ts @@ -14,10 +14,8 @@ const validateFeature = (data: any) => { let config = data.config; if (featureType === FeatureType.Metered) { config = validateMeteredConfig(config); - } else if (featureType === FeatureType.CreditSystem) { - config = validateCreditSystem(config, { isAiCreditSystem: false }); - } else if (featureType === FeatureType.AiCreditSystem) { - config = validateCreditSystem(config, { isAiCreditSystem: true }); + } else if (featureType === FeatureType.CreditSystem || featureType === FeatureType.AiCreditSystem) { + config = validateCreditSystem(config, featureType); } const parsedFeature = CreateFeatureSchema.parse({ ...data, config }); diff --git a/server/src/internal/features/featureActions/updateFeature.ts b/server/src/internal/features/featureActions/updateFeature.ts index 514d9b164..d1c48914b 100644 --- a/server/src/internal/features/featureActions/updateFeature.ts +++ b/server/src/internal/features/featureActions/updateFeature.ts @@ -176,16 +176,19 @@ export const updateFeature = async ({ } const effectiveType = updates.type ?? feature.type; - const isAiCreditSystem = effectiveType === FeatureType.AiCreditSystem; - const newConfig = - updates.config !== undefined - ? effectiveType === FeatureType.CreditSystem || isAiCreditSystem - ? validateCreditSystem(updates.config, { isAiCreditSystem }) - : effectiveType === FeatureType.Metered - ? validateMeteredConfig(updates.config) - : updates.config - : feature.config; + const newConfig = (() => { + if (updates.config === undefined) return feature.config; + switch (effectiveType) { + case FeatureType.AiCreditSystem: + case FeatureType.CreditSystem: + return validateCreditSystem(updates.config, effectiveType); + case FeatureType.Metered: + return validateMeteredConfig(updates.config); + default: + return updates.config; + } + })(); // Update the feature const updatedFeature = await FeatureService.update({ diff --git a/server/src/internal/features/featureUtils.ts b/server/src/internal/features/featureUtils.ts index 761b84e04..91a0118e1 100644 --- a/server/src/internal/features/featureUtils.ts +++ b/server/src/internal/features/featureUtils.ts @@ -43,9 +43,10 @@ export const validateMeteredConfig = (config: MeteredConfig) => { export const validateCreditSystem = ( config: CreditSystemConfig, - { isAiCreditSystem = false }: { isAiCreditSystem?: boolean } = {}, + featureType: FeatureType = FeatureType.CreditSystem, ) => { const schema = config.schema; + const isAiCreditSystem = featureType === FeatureType.AiCreditSystem; if (!isAiCreditSystem && (!schema || schema.length === 0)) { throw new RecaseError({ @@ -55,11 +56,6 @@ export const validateCreditSystem = ( }); } - if (isAiCreditSystem) { - return { ...config, usage_type: FeatureUsageType.Single }; - } - - // Check if multiple of the same feature const meteredFeatureIds = schema.map( (schemaItem) => schemaItem.metered_feature_id, ); @@ -92,6 +88,7 @@ export const validateCreditSystem = ( return newConfig; }; + const getCusFeatureType = ({ feature }: { feature: Feature }) => { if (feature.type === FeatureType.Boolean) { return ApiFeatureType.Static; diff --git a/server/tests/integration/balances/track/basic/track-tokens.test.ts b/server/tests/integration/balances/track/basic/track-tokens.test.ts index fe7b3bc8e..b9d15dedb 100644 --- a/server/tests/integration/balances/track/basic/track-tokens.test.ts +++ b/server/tests/integration/balances/track/basic/track-tokens.test.ts @@ -48,7 +48,7 @@ test.concurrent(`${chalk.yellowBright("track-tokens-1: basic trackTokens with mo tokens: { input: inputTokens, output: outputTokens }, }); - const trackRes: TrackResponseV2 = await autumnV2.post("/trackTokens", { + const trackRes: TrackResponseV2 = await autumnV2.post("/track_tokens", { customer_id: customerId, feature_id: TestFeature.AiCredits, model_id: modelId, @@ -93,7 +93,7 @@ test.concurrent(`${chalk.yellowBright("track-tokens-2: disambiguation error and // Without feature_id, should fail with disambiguation error let error: any; try { - await autumnV2.post("/trackTokens", { + await autumnV2.post("/track_tokens", { customer_id: customerId, model_id: "anthropic/claude-sonnet-4-20250514", input_tokens: 100, @@ -121,7 +121,7 @@ test.concurrent(`${chalk.yellowBright("track-tokens-2: disambiguation error and tokens: { input: inputTokens, output: outputTokens }, }); - const trackRes: TrackResponseV2 = await autumnV2.post("/trackTokens", { + const trackRes: TrackResponseV2 = await autumnV2.post("/track_tokens", { customer_id: customerId, feature_id: TestFeature.AiCredits, model_id: modelId, @@ -170,7 +170,7 @@ test.concurrent(`${chalk.yellowBright("track-tokens-3: custom model pricing with .div(1_000_000) .toNumber(); // 0.125 - const trackRes1: TrackResponseV2 = await autumnV2.post("/trackTokens", { + const trackRes1: TrackResponseV2 = await autumnV2.post("/track_tokens", { customer_id: customerId, feature_id: TestFeature.AiCredits, model_id: "custom/internal-model", @@ -189,7 +189,7 @@ test.concurrent(`${chalk.yellowBright("track-tokens-3: custom model pricing with .mul(new Decimal(1).add(new Decimal(50).div(100))) .toNumber(); // 0.21 - const trackRes2: TrackResponseV2 = await autumnV2.post("/trackTokens", { + const trackRes2: TrackResponseV2 = await autumnV2.post("/track_tokens", { customer_id: customerId, feature_id: TestFeature.AiCredits, model_id: "custom/marked-up-model", @@ -250,7 +250,7 @@ test.concurrent(`${chalk.yellowBright("track-tokens-4: models.dev markup and non tokens: { input: inputTokens, output: outputTokens }, }); - const trackRes: TrackResponseV2 = await autumnV2.post("/trackTokens", { + const trackRes: TrackResponseV2 = await autumnV2.post("/track_tokens", { customer_id: customerId, feature_id: TestFeature.AiCredits, model_id: modelId, @@ -269,7 +269,7 @@ test.concurrent(`${chalk.yellowBright("track-tokens-4: models.dev markup and non // Pointing at a regular credit system should fail let error: any; try { - await autumnV2.post("/trackTokens", { + await autumnV2.post("/track_tokens", { customer_id: customerId, feature_id: TestFeature.Credits, model_id: "anthropic/claude-sonnet-4-20250514", @@ -315,7 +315,7 @@ test.concurrent(`${chalk.yellowBright("track-tokens-5: multiple tracks accumulat tokens: { input: 5000, output: 2000 }, }); - await autumnV2.post("/trackTokens", { + await autumnV2.post("/track_tokens", { customer_id: customerId, feature_id: TestFeature.AiCredits, model_id: "custom/internal-model", @@ -331,7 +331,7 @@ test.concurrent(`${chalk.yellowBright("track-tokens-5: multiple tracks accumulat tokens: { input: 3000, output: 1000 }, }); - await autumnV2.post("/trackTokens", { + await autumnV2.post("/track_tokens", { customer_id: customerId, feature_id: TestFeature.AiCredits, model_id: "custom/marked-up-model", From 2b9b5c7b2b7816b75734a629f9051aa568d5ff32 Mon Sep 17 00:00:00 2001 From: Charlie Lamb Date: Wed, 20 May 2026 10:49:43 +0100 Subject: [PATCH 04/46] chore: cache/type cleanup Co-authored-by: TheUntraceable <73362400+TheUntraceable@users.noreply.github.com> --- .../balances/utils/types/featureDeduction.ts | 16 +++--- .../features/featureActions/createFeature.ts | 11 +--- .../handleUpdateFeatureV1.ts | 3 +- .../features/utils/getModelPricing.ts | 53 +++++++++---------- 4 files changed, 37 insertions(+), 46 deletions(-) diff --git a/server/src/internal/balances/utils/types/featureDeduction.ts b/server/src/internal/balances/utils/types/featureDeduction.ts index eb6d3b0eb..ae6b7755d 100644 --- a/server/src/internal/balances/utils/types/featureDeduction.ts +++ b/server/src/internal/balances/utils/types/featureDeduction.ts @@ -1,19 +1,19 @@ import type { Feature, LockParams } from "@autumn/shared"; import type { LockReceipt } from "../lock/fetchLockReceipt.js"; + +export type TokenUsage = { + modelName: string; + inputTokens: number; + outputTokens: number; +}; + export type FeatureDeduction = { feature: Feature; deduction: number; targetBalance?: number; - - tokenUsage?: { - modelName: string; - inputTokens: number; - outputTokens: number; - }; - + tokenUsage?: TokenUsage; /** Pre-computed dollar cost; if set, the deduction layer skips its own getCreditCost call. */ precomputedCreditCost?: number; - lock?: LockParams; lockReceipt?: LockReceipt; lockReceiptKey?: string; diff --git a/server/src/internal/features/featureActions/createFeature.ts b/server/src/internal/features/featureActions/createFeature.ts index a2f9ae88a..4f11bffa3 100644 --- a/server/src/internal/features/featureActions/createFeature.ts +++ b/server/src/internal/features/featureActions/createFeature.ts @@ -1,4 +1,4 @@ -import { CreateFeatureSchema, type Feature, FeatureType } from "@autumn/shared"; +import { CreateFeatureSchema, type Feature, FeatureType, type ModelMarkups } from "@autumn/shared"; import type { AutumnContext } from "@/honoUtils/HonoEnv.js"; import { workflows } from "@/queue/workflows.js"; import { generateId } from "@/utils/genUtils.js"; @@ -30,14 +30,7 @@ interface CreateFeatureParams { type: string; config?: any; event_names?: string[]; - model_markups?: Record< - string, - { - markup: number; - input_cost?: number; - output_cost?: number; - } - > | null; + model_markups?: ModelMarkups; }; skipGenerateDisplay?: boolean; } diff --git a/server/src/internal/features/handlers/handleUpdateFeature/handleUpdateFeatureV1.ts b/server/src/internal/features/handlers/handleUpdateFeature/handleUpdateFeatureV1.ts index 2c4bd4263..3a0bb0631 100644 --- a/server/src/internal/features/handlers/handleUpdateFeature/handleUpdateFeatureV1.ts +++ b/server/src/internal/features/handlers/handleUpdateFeature/handleUpdateFeatureV1.ts @@ -58,8 +58,7 @@ export const handleUpdateFeatureV1 = createRoute({ archived: body.archived, event_names: body.event_names, display: body.display, - model_markups: - body.model_markups === undefined ? undefined : body.model_markups, + model_markups: body.model_markups, }, }); diff --git a/server/src/internal/features/utils/getModelPricing.ts b/server/src/internal/features/utils/getModelPricing.ts index ac61bb8a3..d45d23cd6 100644 --- a/server/src/internal/features/utils/getModelPricing.ts +++ b/server/src/internal/features/utils/getModelPricing.ts @@ -1,36 +1,35 @@ -import { CacheManager } from "@/utils/cacheUtils/CacheManager"; +import { CacheManager } from "@/utils/cacheUtils/CacheManager.js"; import { ErrCode, InternalError, type ModelsDevProvider } from "@autumn/shared"; -const MODELS_DEV_CACHE_KEY = "models_dev_pricing"; +type ModelPricingData = Record; + +const CACHE_KEY = "models_dev_pricing"; +const STALE_KEY = `${CACHE_KEY}_stale`; +const TTL_PRIMARY = 60 * 60 * 3; +const TTL_STALE = 60 * 60 * 24 * 3; + +const fetchFromSource = async (): Promise => { + const response = await fetch("https://models.dev/api.json"); + if (!response.ok) { + throw new InternalError({ + message: `models.dev returned ${response.status}`, + code: ErrCode.InternalError, + }); + } + return response.json(); +}; + +export const getModelsDevPricing = async (): Promise => { + const cached = await CacheManager.getJson(CACHE_KEY); + if (cached) return cached; -export const getModelsDevPricing = async () => { try { - const cached = - await CacheManager.getJson>( - MODELS_DEV_CACHE_KEY, - ); - if (cached) return cached; - const response = await fetch("https://models.dev/api.json"); - if (!response.ok) - throw new InternalError({ - message: `models.dev returned ${response.status}`, - code: ErrCode.InternalError, - }); - - const data: Record = await response.json(); - Promise.all([ - CacheManager.setJson(MODELS_DEV_CACHE_KEY, data, 60 * 60 * 3), - CacheManager.setJson( - `${MODELS_DEV_CACHE_KEY}_stale`, - data, - 60 * 60 * 24 * 3, - ), - ]).catch(() => {}); + const data = await fetchFromSource(); + CacheManager.setJson(CACHE_KEY, data, TTL_PRIMARY).catch(() => {}); + CacheManager.setJson(STALE_KEY, data, TTL_STALE).catch(() => {}); return data; } catch { - const stale = await CacheManager.getJson>( - `${MODELS_DEV_CACHE_KEY}_stale`, - ); + const stale = await CacheManager.getJson(STALE_KEY); if (stale) return stale; throw new InternalError({ message: "Failed to fetch models.dev pricing and no cache available", From 455e2216531df9a3085ccbdccab933e23b1b125f Mon Sep 17 00:00:00 2001 From: Charlie Lamb Date: Wed, 20 May 2026 11:41:56 +0100 Subject: [PATCH 05/46] chore: clean up types Co-authored-by: TheUntraceable <73362400+TheUntraceable@users.noreply.github.com> --- .../balances/utils/deduction/computeCreditCosts.ts | 11 +++++------ server/src/internal/features/creditSystemUtils.ts | 13 +++++-------- .../features/featureActions/updateFeature.ts | 3 ++- server/tests/utils/fixtures/db/entitlements.ts | 6 ++---- server/tests/utils/fixtures/db/features.ts | 7 ++----- shared/utils/agentTypes.ts | 10 ++-------- 6 files changed, 18 insertions(+), 32 deletions(-) diff --git a/server/src/internal/balances/utils/deduction/computeCreditCosts.ts b/server/src/internal/balances/utils/deduction/computeCreditCosts.ts index b6c203f48..05889e041 100644 --- a/server/src/internal/balances/utils/deduction/computeCreditCosts.ts +++ b/server/src/internal/balances/utils/deduction/computeCreditCosts.ts @@ -25,18 +25,17 @@ export const computeCreditCosts = async ({ return () => cost; } + const tokens = deduction.tokenUsage + ? { input: deduction.tokenUsage.inputTokens, output: deduction.tokenUsage.outputTokens } + : undefined; + await Promise.all( cusEnts.map(async (ce) => { const creditCost = await getCreditCost({ featureId: deduction.feature.id, creditSystem: ce.entitlement.feature, modelName: deduction.tokenUsage?.modelName, - tokens: deduction.tokenUsage - ? { - input: deduction.tokenUsage.inputTokens, - output: deduction.tokenUsage.outputTokens, - } - : undefined, + tokens, }); costMap.set(ce.id, creditCost); }), diff --git a/server/src/internal/features/creditSystemUtils.ts b/server/src/internal/features/creditSystemUtils.ts index 367af65e6..3b47a3009 100644 --- a/server/src/internal/features/creditSystemUtils.ts +++ b/server/src/internal/features/creditSystemUtils.ts @@ -7,7 +7,9 @@ import { RecaseError, } from "@autumn/shared"; import { Decimal } from "decimal.js"; -import { getModelsDevPricing } from "@/internal/features/utils/getModelPricing"; +import { getModelsDevPricing } from "@/internal/features/utils/getModelPricing.js"; + +type TokenInput = { input: number; output: number }; const creditSystemContainsFeature = ({ creditSystem, @@ -103,9 +105,7 @@ const getModelCreditCost = async ({ }: { modelName: string; creditSystem: Feature; - input: number; - output: number; -}) => { +} & TokenInput) => { const markups = creditSystem.model_markups || {}; const markupEntry = markups[modelName]; const { markup } = markupEntry ?? { markup: 0 }; @@ -168,10 +168,7 @@ export const getCreditCost = async ({ creditSystem: Feature; amount?: number; modelName?: string; - tokens?: { - input: number; - output: number; - }; + tokens?: TokenInput; }) => { if (creditSystem.type !== FeatureType.CreditSystem && creditSystem.type !== FeatureType.AiCreditSystem) { return amount; diff --git a/server/src/internal/features/featureActions/updateFeature.ts b/server/src/internal/features/featureActions/updateFeature.ts index d1c48914b..909c4ec4c 100644 --- a/server/src/internal/features/featureActions/updateFeature.ts +++ b/server/src/internal/features/featureActions/updateFeature.ts @@ -217,7 +217,8 @@ export const updateFeature = async ({ } // Queue cache clear for credit system if schema or model markups changed - if ((feature.type === FeatureType.CreditSystem || feature.type === FeatureType.AiCreditSystem) && updatedFeature) { + const isCreditSystem = feature.type === FeatureType.CreditSystem || feature.type === FeatureType.AiCreditSystem; + if (isCreditSystem && updatedFeature) { const schemaChanged = updates.config != null && hasCreditSchemaChanged({ diff --git a/server/tests/utils/fixtures/db/entitlements.ts b/server/tests/utils/fixtures/db/entitlements.ts index b8be6cafd..7bc7dcf0c 100644 --- a/server/tests/utils/fixtures/db/entitlements.ts +++ b/server/tests/utils/fixtures/db/entitlements.ts @@ -2,6 +2,7 @@ import { AllowanceType, type EntInterval, FeatureType, + type ModelMarkups, type RolloverConfig, } from "@autumn/shared"; import { features } from "./features"; @@ -34,10 +35,7 @@ const create = ({ intervalCount?: number; entityFeatureId?: string | null; rollover?: RolloverConfig | null; - modelMarkups?: Record< - string, - { markup: number; input_cost?: number; output_cost?: number } - > | null; + modelMarkups?: ModelMarkups; }) => ({ id: id ?? `ent_${featureId}_${crypto.randomUUID().slice(0, 8)}`, created_at: Date.now(), diff --git a/server/tests/utils/fixtures/db/features.ts b/server/tests/utils/fixtures/db/features.ts index 695a5053a..05dcc2df6 100644 --- a/server/tests/utils/fixtures/db/features.ts +++ b/server/tests/utils/fixtures/db/features.ts @@ -1,4 +1,4 @@ -import { AppEnv, FeatureType } from "@autumn/shared"; +import { AppEnv, FeatureType, type ModelMarkups } from "@autumn/shared"; /** * Create a feature fixture @@ -16,10 +16,7 @@ const create = ({ name: string; type?: FeatureType; config?: Record; - modelMarkups?: Record< - string, - { markup: number; input_cost?: number; output_cost?: number } - > | null; + modelMarkups?: ModelMarkups; }) => ({ internal_id: internalId ?? `internal_${id}`, org_id: "org_test", diff --git a/shared/utils/agentTypes.ts b/shared/utils/agentTypes.ts index 69e0ca6ad..5eea84691 100644 --- a/shared/utils/agentTypes.ts +++ b/shared/utils/agentTypes.ts @@ -10,6 +10,7 @@ * - Converters: AgentFeature ↔ Feature, AgentProduct ↔ ProductV2 */ +import type { ModelMarkups } from "../models/featureModels/featureConfig/creditConfig.js"; import { FeatureType, FeatureUsageType, @@ -42,14 +43,7 @@ export interface AgentFeature { metered_feature_id: string; credit_cost: number; }> | null; - model_markups?: Record< - string, - { - markup: number; - input_cost?: number; - output_cost?: number; - } - > | null; + model_markups?: ModelMarkups; } export interface AgentProductItem { From 3e6ffa54f66cb0de301f818b5a93ea97b2545f11 Mon Sep 17 00:00:00 2001 From: Charlie Lamb Date: Wed, 20 May 2026 15:33:52 +0100 Subject: [PATCH 06/46] chore: ai credit system frontend overhaul Co-authored-by: TheUntraceable <73362400+TheUntraceable@users.noreply.github.com> --- .../features/changes/V1.2_FeatureChange.ts | 4 +- .../api/features/prevVersions/apiFeatureV0.ts | 1 + .../components/ai-elements/prompt-input.tsx | 8 +- .../table/table-content-virtualized.tsx | 12 +- .../product/product-item/formatProductItem.ts | 2 - .../customer/analytics/AnalyticsView.tsx | 16 +- .../components/AuthorizedApps.tsx | 71 ++-- .../org-dropdown/components/OrgLogo.tsx | 12 +- .../components/AiCreditSchema.tsx | 232 ++++------- .../components/AiCreditSchemaRow.tsx | 172 -------- .../components/AiCreditSchemaTable.tsx | 369 ++++++++++++++++++ .../components/AiModelSelectDropdown.tsx | 14 +- .../components/ClassicCreditSchema.tsx | 44 +-- .../components/CreditSystemDetails.tsx | 39 +- .../components/CreditSystemSchema.tsx | 63 ++- .../components/UpdateCreditSystemSheet.tsx | 136 +++---- .../credit-systems/hooks/useAiCreditSchema.ts | 259 ------------ .../hooks/useCreditSystemForm.ts | 41 ++ .../feature-list/CreditListColumns.tsx | 28 +- .../feature-list/FeatureListTable.tsx | 8 +- .../features/utils/getFeatureIcon.tsx | 13 + .../edit-plan-feature/IncludedUsage.tsx | 5 +- .../new-feature/NewFeatureBehaviour.tsx | 46 ++- 23 files changed, 739 insertions(+), 856 deletions(-) delete mode 100644 vite/src/views/products/features/credit-systems/components/AiCreditSchemaRow.tsx create mode 100644 vite/src/views/products/features/credit-systems/components/AiCreditSchemaTable.tsx delete mode 100644 vite/src/views/products/features/credit-systems/hooks/useAiCreditSchema.ts create mode 100644 vite/src/views/products/features/credit-systems/hooks/useCreditSystemForm.ts diff --git a/shared/api/features/changes/V1.2_FeatureChange.ts b/shared/api/features/changes/V1.2_FeatureChange.ts index 101e9cc57..a03642547 100644 --- a/shared/api/features/changes/V1.2_FeatureChange.ts +++ b/shared/api/features/changes/V1.2_FeatureChange.ts @@ -63,13 +63,13 @@ export const V1_2_FeatureChange = defineVersionChange({ v0Type = ApiFeatureType.Boolean; } else if (input.type === FeatureType.CreditSystem) { v0Type = ApiFeatureType.CreditSystem; + } else if (input.type === FeatureType.AiCreditSystem) { + v0Type = ApiFeatureType.AiCreditSystem; } else if (input.type === FeatureType.Metered) { - // Use consumable flag to determine single_use vs continuous_use v0Type = input.consumable ? ApiFeatureType.SingleUsage : ApiFeatureType.ContinuousUse; } else { - // Fallback (should never happen) v0Type = ApiFeatureType.Boolean; } diff --git a/shared/api/features/prevVersions/apiFeatureV0.ts b/shared/api/features/prevVersions/apiFeatureV0.ts index 1ac487994..cf383d811 100644 --- a/shared/api/features/prevVersions/apiFeatureV0.ts +++ b/shared/api/features/prevVersions/apiFeatureV0.ts @@ -7,6 +7,7 @@ export enum ApiFeatureType { SingleUsage = "single_use", ContinuousUse = "continuous_use", CreditSystem = "credit_system", + AiCreditSystem = "ai_credit_system", } export const FEATURE_EXAMPLE = { diff --git a/vite/src/components/ai-elements/prompt-input.tsx b/vite/src/components/ai-elements/prompt-input.tsx index 4831f85c2..f43bcab65 100644 --- a/vite/src/components/ai-elements/prompt-input.tsx +++ b/vite/src/components/ai-elements/prompt-input.tsx @@ -301,7 +301,7 @@ export function PromptInputAttachment({
{isImage && ( -
+
{filename )} - {/* Fixed header table - scrolls horizontally in sync with body */} -
- {/* Column visibility toggle - only render if not in toolbar */} - {enableColumnVisibility && !columnVisibilityInToolbar && ( + {enableColumnVisibility && !columnVisibilityInToolbar && (
- {/* Scroll container for body only - key forces remount when columns change */} -
- {/* Clone children with key to force remount when columns change */} - {React.Children.map(children, (child) => + {React.Children.map(children, (child) => React.isValidElement(child) ? React.cloneElement(child, { key: visibleColumnKey }) : child, diff --git a/vite/src/utils/product/product-item/formatProductItem.ts b/vite/src/utils/product/product-item/formatProductItem.ts index b29ae767c..d3eff3e7d 100644 --- a/vite/src/utils/product/product-item/formatProductItem.ts +++ b/vite/src/utils/product/product-item/formatProductItem.ts @@ -140,8 +140,6 @@ const getFeatureString = ({ item: ProductItem; features: Feature[]; }) => { - // This function isn't exported or used - // Just updated it to account for the new AI Credit System features, but it may be unused and can probably be deleted const feature = features.find((f: Feature) => f.id === item.feature_id); if (feature?.type === FeatureType.Boolean) { diff --git a/vite/src/views/customers/customer/analytics/AnalyticsView.tsx b/vite/src/views/customers/customer/analytics/AnalyticsView.tsx index d21f97441..84763d418 100644 --- a/vite/src/views/customers/customer/analytics/AnalyticsView.tsx +++ b/vite/src/views/customers/customer/analytics/AnalyticsView.tsx @@ -156,14 +156,14 @@ export const AnalyticsView = () => { const trimmed = { ...transformed, data: nonEmptyData, rows: nonEmptyData.length }; const config = generateChartConfig({ - events: trimmed, - features, - groupBy, - originalColors: colors, - entityNames, - customerNames, - planNames, - }); + events: trimmed, + features, + groupBy, + originalColors: colors, + entityNames, + customerNames, + planNames, + }); return { chartData: trimmed, chartConfig: config }; }, [events, features, groupBy, groupFilter, planDeselected, entityNames, customerNames, planNames]); diff --git a/vite/src/views/main-sidebar/components/AuthorizedApps.tsx b/vite/src/views/main-sidebar/components/AuthorizedApps.tsx index a09afc5a6..323f3e67a 100644 --- a/vite/src/views/main-sidebar/components/AuthorizedApps.tsx +++ b/vite/src/views/main-sidebar/components/AuthorizedApps.tsx @@ -15,26 +15,25 @@ import { import { getBackendErr } from "@/utils/genUtils"; interface OAuthConsent { - id: string; - clientId: string; - scopes: string[]; - referenceId: string | null; - createdAt: string; - updatedAt: string; + readonly id: string; + readonly clientId: string; + readonly scopes: string[]; + readonly referenceId: string | null; + readonly createdAt: string; + readonly updatedAt: string; } interface ClientInfo { - client_id: string; - name: string; + readonly client_id: string; + readonly name: string; } interface ApiKeyPreview { - prefix: string; - name: string; - env: string; + readonly prefix: string; + readonly name: string; + readonly env: string; } -// Helper to render scope badges function renderScopeBadges(scopes: string[]): JSX.Element { const grouped = groupAndFormatScopes(scopes); @@ -64,8 +63,6 @@ export const AuthorizedApps = () => { const [clientNames, setClientNames] = useState>({}); const [isLoading, setIsLoading] = useState(true); const [revokingId, setRevokingId] = useState(null); - - // Revoke confirmation dialog state const [revokeDialogOpen, setRevokeDialogOpen] = useState(false); const [revokeTarget, setRevokeTarget] = useState<{ consentId: string; @@ -77,25 +74,18 @@ export const AuthorizedApps = () => { const fetchConsents = async () => { setIsLoading(true); try { - // Use our new org-level consent endpoint const response = await fetch( `${import.meta.env.VITE_BACKEND_URL}/consents`, - { - credentials: "include", - }, + { credentials: "include" }, ); - if (!response.ok) { const error = await response.json().catch(() => ({})); toast.error(error.message || "Failed to fetch authorized apps"); return; } - const data = await response.json(); const consentList = (data.consents as OAuthConsent[]) || []; setConsents(consentList); - - // Fetch client names for each consent const names: Record = {}; await Promise.all( consentList.map(async (consent) => { @@ -131,22 +121,17 @@ export const AuthorizedApps = () => { setRevokeDialogOpen(true); setLoadingApiKeys(true); setLinkedApiKeys([]); - - // Fetch API keys linked to this consent try { const response = await fetch( `${import.meta.env.VITE_BACKEND_URL}/consents/${consentId}/api-keys`, - { - credentials: "include", - }, + { credentials: "include" }, ); - if (response.ok) { const data = await response.json(); setLinkedApiKeys(data.apiKeys || []); } } catch { - // If we can't fetch API keys, just show the basic dialog + // Silent — dialog still usable without API key preview } finally { setLoadingApiKeys(false); } @@ -154,23 +139,17 @@ export const AuthorizedApps = () => { const handleConfirmRevoke = async () => { if (!revokeTarget) return; - setRevokingId(revokeTarget.consentId); try { const response = await fetch( `${import.meta.env.VITE_BACKEND_URL}/consents/${revokeTarget.consentId}`, - { - method: "DELETE", - credentials: "include", - }, + { method: "DELETE", credentials: "include" }, ); - if (!response.ok) { const error = await response.json().catch(() => ({})); toast.error(error.message || "Failed to revoke access"); return; } - const result = await response.json(); toast.success( `Access revoked for ${revokeTarget.clientName}${result.deletedApiKeys > 0 ? ` (${result.deletedApiKeys} API key${result.deletedApiKeys > 1 ? "s" : ""} deleted)` : ""}`, @@ -226,20 +205,16 @@ export const AuthorizedApps = () => {

- {/* Scopes - Grouped by Resource */} -
- {renderScopeBadges(consent.scopes)} -
- - {/* Date */} -
+
+ {renderScopeBadges(consent.scopes)} +
+
Authorized on {formatDate(consent.createdAt)}
- {/* Revoke Button */} -
- {/* Revoke Confirmation Dialog */} - + Revoke Access @@ -274,8 +248,7 @@ export const AuthorizedApps = () => { - {/* Show linked API keys that will be deleted */} - {loadingApiKeys ? ( + {loadingApiKeys ? (
Checking for linked API keys...
diff --git a/vite/src/views/main-sidebar/org-dropdown/components/OrgLogo.tsx b/vite/src/views/main-sidebar/org-dropdown/components/OrgLogo.tsx index 1bdbcc254..faf599de2 100644 --- a/vite/src/views/main-sidebar/org-dropdown/components/OrgLogo.tsx +++ b/vite/src/views/main-sidebar/org-dropdown/components/OrgLogo.tsx @@ -1,21 +1,13 @@ import type { FrontendOrg } from "@autumn/shared"; -import { cn } from "@/lib/utils"; + export const OrgLogo = ({ org }: { org: FrontendOrg }) => { const firstLetter = org?.name?.charAt(0).toUpperCase() || "A"; return ( -
- {/* {org.logo ? ( - {org.name} - ) : ( */} +
{firstLetter} - {/* )} */}
); }; diff --git a/vite/src/views/products/features/credit-systems/components/AiCreditSchema.tsx b/vite/src/views/products/features/credit-systems/components/AiCreditSchema.tsx index 10bb29b0e..624d818f3 100644 --- a/vite/src/views/products/features/credit-systems/components/AiCreditSchema.tsx +++ b/vite/src/views/products/features/credit-systems/components/AiCreditSchema.tsx @@ -1,189 +1,115 @@ -import type { CreateFeature } from "@autumn/shared"; -import { PlusIcon, X } from "lucide-react"; -import { IconButton } from "@/components/v2/buttons/IconButton"; +import { PlusIcon } from "lucide-react"; +import { useMemo } from "react"; +import { useStore } from "@tanstack/react-form"; import { FormLabel } from "@/components/v2/form/FormLabel"; import { Input } from "@/components/v2/inputs/Input"; import { SearchableSelect } from "@/components/v2/selects/SearchableSelect"; -import { useAiCreditSchema } from "../hooks/useAiCreditSchema"; -import { AiCreditSchemaRow } from "./AiCreditSchemaRow"; +import { useModelsDevPricing } from "@/hooks/queries/useAiModelsQuery"; +import type { CreditSystemFormInstance } from "../hooks/useCreditSystemForm"; +import { AiCreditSchemaTable } from "./AiCreditSchemaTable"; interface AiCreditSchemaProps { - creditSystem: CreateFeature; - setCreditSystem: ( - creditSystem: CreateFeature | ((prev: CreateFeature) => CreateFeature), - ) => void; + form: CreditSystemFormInstance; } -export function AiCreditSchema({ - creditSystem, - setCreditSystem, -}: AiCreditSchemaProps) { - const { - providers, - modelsLoading, - modelMarkups, - defaultMarkup, - providerGroups, - activeProviderKeys, - availableProviders, - handleModelChange, - handleMarkupChange, - handleDefaultMarkupChange, - handleCostChange, - handleRemoveModel, - handleRemoveProvider, - addProvider, - addModelToProvider, - } = useAiCreditSchema({ creditSystem, setCreditSystem }); +function groupByProvider(markups: Record) { + const groups: Record = {}; + for (const fullId of Object.keys(markups)) { + const [provider] = fullId.split("/"); + (groups[provider] ??= []).push(fullId); + } + return groups; +} + +export function AiCreditSchema({ form }: AiCreditSchemaProps) { + const { providers, isLoading } = useModelsDevPricing(); + const modelMarkups = useStore(form.store, (s) => s.values.model_markups); + const defaultMarkup = useStore(form.store, (s) => s.values.defaultMarkup); + + const providerGroups = useMemo(() => groupByProvider(modelMarkups), [modelMarkups]); + const activeProviderKeys = Object.keys(providerGroups); + + const availableProviders = useMemo(() => { + const filtered = Object.values(providers).filter( + (p) => !activeProviderKeys.includes(p.id), + ); + if (!activeProviderKeys.includes("custom")) { + filtered.push({ id: "custom", name: "Custom", models: {} }); + } + return filtered; + }, [providers, activeProviderKeys]); + + const addProvider = (providerKey: string) => { + form.setFieldValue("model_markups", (prev) => { + if (providerKey === "custom") { + const existing = Object.keys(prev).filter((k) => k.startsWith("custom/")); + let i = 1; + while (existing.includes(`custom/model-${i}`)) i++; + return { ...prev, [`custom/model-${i}`]: { input_cost: 0, output_cost: 0 } }; + } + const provider = providers[providerKey]; + if (!provider) return prev; + const firstKey = Object.keys(provider.models)[0]; + if (!firstKey) return prev; + return { ...prev, [`${providerKey}/${firstKey}`]: {} }; + }); + }; return ( -
-
- Default Markup % +
+
+ Default Markup % - handleDefaultMarkupChange(Number(e.target.value) || 0) - } - onBlur={(e) => handleDefaultMarkupChange(Number(e.target.value) || 0)} + type="text" + inputMode="numeric" + value={defaultMarkup === 0 ? "" : String(defaultMarkup)} + onChange={(e) => { + const raw = e.target.value; + if (raw === "" || /^-?\d*\.?\d*$/.test(raw)) { + form.setFieldValue("defaultMarkup", raw === "" ? 0 : Number(raw)); + } + }} placeholder="0" - className="w-24" />
-
+ +
{activeProviderKeys.map((providerKey) => { const provider = providers[providerKey]; const modelFullIds = providerGroups[providerKey] ?? []; const providerName = - provider?.name ?? - providerKey.charAt(0).toUpperCase() + providerKey.slice(1); + provider?.name ?? providerKey.charAt(0).toUpperCase() + providerKey.slice(1); return ( -
-
- - {providerName} - {providerKey !== "custom" && ( - {providerName} - )} - - } - onClick={() => handleRemoveProvider(providerKey)} - /> -
- -
- {providerKey === "custom" && ( -

- In your API tracking, use the format{" "} - - custom/{"modelId"} - -

- )} - -
-
- Model -
-
- {providerKey === "custom" ? "In $/M" : "Cost In"} -
-
- {providerKey === "custom" ? "Out $/M" : "Cost Out"} -
-
- Markup % -
-
-
- - {modelFullIds.map((fullId) => { - const [, ...parts] = fullId.split("/"); - const modelKey = parts.join("/"); - const isCustom = providerKey === "custom"; - return ( - - handleModelChange(providerKey, oldKey, newKey) - } - onMarkupChange={(key, newMarkup) => - handleMarkupChange(providerKey, key, newMarkup) - } - onCostChange={(key, field, value) => - handleCostChange(providerKey, key, field, value) - } - onRemove={(key) => handleRemoveModel(providerKey, key)} - /> - ); - })} - - addModelToProvider(providerKey)} - className="w-fit mt-0.5" - icon={} - disabled={ - providerKey === "custom" - ? false - : Object.keys(provider?.models ?? {}).length === - modelFullIds.length - } - > - Add model - -
-
+ form={form} + providerKey={providerKey} + providerName={providerName} + modelFullIds={modelFullIds} + provider={provider ?? { id: providerKey, name: providerKey, models: {} }} + isLoading={isLoading} + /> ); })}
-

All prices in $/M tokens

- -
e.stopPropagation()}> +
e.stopPropagation()}> + Add Provider provider.id} - getOptionLabel={(provider) => provider.name} + getOptionValue={(p) => p.id} + getOptionLabel={(p) => p.name} renderValue={() => ( - - - Add provider - + Select provider )} - placeholder="Add provider" + placeholder="Select provider" searchable searchPlaceholder="Search providers..." emptyText="No providers available" - disabled={modelsLoading} + disabled={isLoading} />
diff --git a/vite/src/views/products/features/credit-systems/components/AiCreditSchemaRow.tsx b/vite/src/views/products/features/credit-systems/components/AiCreditSchemaRow.tsx deleted file mode 100644 index b0b2c7b29..000000000 --- a/vite/src/views/products/features/credit-systems/components/AiCreditSchemaRow.tsx +++ /dev/null @@ -1,172 +0,0 @@ -import type { ModelsDevProvider } from "@autumn/shared"; -import { X } from "lucide-react"; -import { useState } from "react"; -import { IconButton } from "@/components/v2/buttons/IconButton"; -import { Input } from "@/components/v2/inputs/Input"; -import { AiModelSelectDropdown } from "./AiModelSelectDropdown"; - -interface AiCreditSchemaRowProps { - modelKey: string; - markup: number; - provider: ModelsDevProvider; - isLoading: boolean; - isCustom?: boolean; - inputCost?: number; - outputCost?: number; - onModelChange: (oldModelKey: string, newModelKey: string) => void; - onMarkupChange: (modelKey: string, markup: number) => void; - onCostChange?: ( - modelKey: string, - field: "input_cost" | "output_cost", - value: number, - ) => void; - onRemove: (modelKey: string) => void; -} - -function formatCost(value: number | null | undefined): string { - if (value == null) return "–"; - return value.toFixed(2); -} - -export function AiCreditSchemaRow({ - modelKey, - markup, - provider, - isLoading, - isCustom, - inputCost, - outputCost, - onModelChange, - onMarkupChange, - onCostChange, - onRemove, -}: AiCreditSchemaRowProps) { - const model = provider.models[modelKey]; - - const actualInput = isCustom - ? (inputCost ?? 0) - : model - ? (model.cost.input ?? 0) - : null; - const actualOutput = isCustom - ? (outputCost ?? 0) - : model - ? (model.cost.output ?? 0) - : null; - const multiplier = 1 + markup / 100; - const userInput = actualInput != null ? actualInput * multiplier : null; - const userOutput = actualOutput != null ? actualOutput * multiplier : null; - - const [localModelName, setLocalModelName] = useState(modelKey); - - return ( -
-
- {/* Model Name */} -
- {isCustom ? ( - setLocalModelName(e.target.value)} - onBlur={() => { - if (localModelName !== modelKey) { - onModelChange(modelKey, localModelName); - } - }} - placeholder="my-model-id" - className="w-full" - /> - ) : ( - - onModelChange(modelKey, newModelKey) - } - provider={provider} - isLoading={isLoading} - /> - )} -
- - {/* Input Cost */} -
- {isCustom ? ( - - onCostChange?.( - modelKey, - "input_cost", - Number(e.target.value) || 0, - ) - } - placeholder="0" - className="w-full" - /> - ) : ( -
- {formatCost(actualInput)} -
- )} -
- - {/* Output Cost */} -
- {isCustom ? ( - - onCostChange?.( - modelKey, - "output_cost", - Number(e.target.value) || 0, - ) - } - placeholder="0" - className="w-full" - /> - ) : ( -
- {formatCost(actualOutput)} -
- )} -
- - {/* Markup % */} -
- - onMarkupChange(modelKey, Number(e.target.value) || 0) - } - placeholder="0" - className="w-full" - /> -
- - {/* Remove Button */} - } - onClick={() => onRemove(modelKey)} - className="shrink-0" - /> -
- - {/* User Pays Info */} - {(userInput != null || userOutput != null) && ( -
- User pays: ${formatCost(userInput)} in / ${formatCost(userOutput)} out - $/M -
- )} -
- ); -} diff --git a/vite/src/views/products/features/credit-systems/components/AiCreditSchemaTable.tsx b/vite/src/views/products/features/credit-systems/components/AiCreditSchemaTable.tsx new file mode 100644 index 000000000..7fe1b5a43 --- /dev/null +++ b/vite/src/views/products/features/credit-systems/components/AiCreditSchemaTable.tsx @@ -0,0 +1,369 @@ +import type { ModelsDevProvider } from "@autumn/shared"; +import type { ColumnDef, Row } from "@tanstack/react-table"; +import { useStore } from "@tanstack/react-form"; +import { InfoIcon, PlusIcon, X } from "lucide-react"; +import { useMemo, useState } from "react"; +import { Table } from "@/components/general/table"; +import { IconButton } from "@/components/v2/buttons/IconButton"; +import { Input } from "@/components/v2/inputs/Input"; +import { + Tooltip, + TooltipContent, + TooltipTrigger, +} from "@/components/v2/tooltips/Tooltip"; +import { useProductTable } from "@/views/products/hooks/useProductTable"; +import type { CreditSystemFormInstance } from "../hooks/useCreditSystemForm"; +import { AiModelSelectDropdown } from "./AiModelSelectDropdown"; + +interface ModelRow { + fullId: string; + modelKey: string; +} + +interface AiCreditSchemaTableProps { + form: CreditSystemFormInstance; + providerKey: string; + providerName: string; + modelFullIds: string[]; + provider: ModelsDevProvider; + isLoading: boolean; +} + +function formatCost(value: number | null | undefined): string { + if (value == null) return "–"; + return value.toFixed(2); +} + +export function AiCreditSchemaTable({ + form, + providerKey, + providerName, + modelFullIds, + provider, + isLoading, +}: AiCreditSchemaTableProps) { + const isCustom = providerKey === "custom"; + + + const removeKeys = (keys: string[]) => + form.setFieldValue("model_markups", (prev) => { + const updated = { ...prev }; + for (const k of keys) delete updated[k]; + return updated; + }); + + const renameKey = (oldKey: string, newKey: string) => + form.setFieldValue("model_markups", (prev) => { + if (newKey in prev) return prev; + const updated = { ...prev }; + const entry = updated[oldKey]; + delete updated[oldKey]; + updated[newKey] = { ...entry }; + return updated; + }); + + const data: ModelRow[] = useMemo( + () => + modelFullIds.map((fullId) => { + const [, ...parts] = fullId.split("/"); + return { fullId, modelKey: parts.join("/") }; + }), + [modelFullIds.join(",")], + ); + + const columns: ColumnDef[] = useMemo( + () => [ + { + header: "Model", + accessorKey: "modelKey", + size: 200, + cell: ({ row }: { row: Row }) => { + const { modelKey } = row.original; + if (isCustom) { + return ( + + renameKey(`${providerKey}/${modelKey}`, `${providerKey}/${newKey}`) + } + /> + ); + } + return ( + + renameKey(`${providerKey}/${modelKey}`, `${providerKey}/${newKey}`) + } + provider={provider} + isLoading={isLoading} + /> + ); + }, + }, + { + header: isCustom ? "In $/M" : "Input", + id: "inputCost", + size: 80, + cell: ({ row }: { row: Row }) => { + const { fullId, modelKey } = row.original; + if (isCustom) { + return ( + + ); + } + const cost = provider.models[modelKey]?.cost?.input ?? null; + return ( + + {formatCost(cost)} + + ); + }, + }, + { + header: isCustom ? "Out $/M" : "Output", + id: "outputCost", + size: 80, + cell: ({ row }: { row: Row }) => { + const { fullId, modelKey } = row.original; + if (isCustom) { + return ( + + ); + } + const cost = provider.models[modelKey]?.cost?.output ?? null; + return ( + + {formatCost(cost)} + + ); + }, + }, + { + header: "Markup %", + id: "markup", + size: 80, + cell: ({ row }: { row: Row }) => ( + + ), + }, + { + header: "", + accessorKey: "actions", + size: 40, + enableSorting: false, + cell: ({ row }: { row: Row }) => ( +
e.stopPropagation()}> + } + onClick={() => removeKeys([row.original.fullId])} + className="!text-subtle hover:!text-foreground" + /> +
+ ), + }, + ], + [isCustom, provider, isLoading, providerKey, form], + ); + + const allModelsUsed = + !isCustom && Object.keys(provider.models).length === modelFullIds.length; + + const table = useProductTable({ + data, + columns, + options: { getRowId: (row) => row.fullId }, + }); + + return ( +
+
+ + {providerName} + {!isCustom && ( + {providerName} + )} + {isCustom && ( + + + + + + Use format custom/modelId in API tracking + + + )} + + } + onClick={() => removeKeys(modelFullIds)} + className="!text-subtle hover:!text-foreground" + /> +
+ +
+ + + + + + + + + + {!allModelsUsed && ( + + )} +
+
+ ); +} + +function CustomModelInput({ + modelKey, + onRename, +}: { + modelKey: string; + onRename: (newKey: string) => void; +}) { + const [local, setLocal] = useState(modelKey); + return ( + setLocal(e.target.value)} + onBlur={() => { + if (local !== modelKey) onRename(local); + }} + placeholder="my-model-id" + className="text-sm" + /> + ); +} + +function EditableNumberCell({ + form, + fullId, + field, + useDefaultAsPlaceholder = false, + allowUndefined = false, +}: { + form: CreditSystemFormInstance; + fullId: string; + field: "markup" | "input_cost" | "output_cost"; + useDefaultAsPlaceholder?: boolean; + allowUndefined?: boolean; +}) { + const currentValue = useStore( + form.store, + (s) => s.values.model_markups[fullId]?.[field], + ); + const placeholder = useStore(form.store, (s) => + useDefaultAsPlaceholder ? String(s.values.defaultMarkup) : "0", + ); + const [local, setLocal] = useState(""); + const [focused, setFocused] = useState(false); + + const hasValue = allowUndefined ? currentValue != null && currentValue !== 0 : currentValue != null; + const displayed = focused ? local : hasValue ? String(currentValue) : ""; + + return ( + { + const raw = e.target.value; + if (raw === "" || /^-?\d*\.?\d*$/.test(raw)) { + setLocal(raw); + if (raw === "" && allowUndefined) { + form.setFieldValue("model_markups", (prev) => { + const entry = { ...prev[fullId] }; + delete entry[field]; + return { ...prev, [fullId]: entry }; + }); + } else if (raw !== "") { + const parsed = Number(raw); + if (!Number.isNaN(parsed)) { + form.setFieldValue("model_markups", (prev) => ({ + ...prev, + [fullId]: { ...prev[fullId], [field]: parsed }, + })); + } + } + } + }} + onFocus={() => { + setLocal(hasValue ? String(currentValue) : ""); + setFocused(true); + }} + onBlur={() => { + setFocused(false); + if (local === "" && !allowUndefined) { + form.setFieldValue("model_markups", (prev) => ({ + ...prev, + [fullId]: { ...prev[fullId], [field]: 0 }, + })); + } + }} + placeholder={placeholder} + className="text-sm" + /> + ); +} diff --git a/vite/src/views/products/features/credit-systems/components/AiModelSelectDropdown.tsx b/vite/src/views/products/features/credit-systems/components/AiModelSelectDropdown.tsx index 6512dae37..e333cd314 100644 --- a/vite/src/views/products/features/credit-systems/components/AiModelSelectDropdown.tsx +++ b/vite/src/views/products/features/credit-systems/components/AiModelSelectDropdown.tsx @@ -23,7 +23,7 @@ export function AiModelSelectDropdown({ return (
e.stopPropagation()}> model.name} renderValue={(option) => option ? ( - {option.name} + {option.name} ) : provider.models[value]?.name ? ( - {provider.models[value].name} + {provider.models[value].name} ) : value ? ( - {value} + {value} ) : ( - - {isLoading ? "Loading models..." : "Select model"} + + {isLoading ? "Loading..." : "Select model"} ) } - placeholder={isLoading ? "Loading models..." : "Select model"} + placeholder={isLoading ? "Loading..." : "Select model"} searchable searchPlaceholder="Search models..." emptyText="No models found" diff --git a/vite/src/views/products/features/credit-systems/components/ClassicCreditSchema.tsx b/vite/src/views/products/features/credit-systems/components/ClassicCreditSchema.tsx index 025640ea2..4542b62c2 100644 --- a/vite/src/views/products/features/credit-systems/components/ClassicCreditSchema.tsx +++ b/vite/src/views/products/features/credit-systems/components/ClassicCreditSchema.tsx @@ -1,26 +1,26 @@ -import type { CreateFeature, CreditSchemaItem, Feature } from "@autumn/shared"; +import type { CreditSchemaItem, Feature } from "@autumn/shared"; import { FeatureType } from "@autumn/shared"; import { PlusIcon } from "@phosphor-icons/react"; +import { useStore } from "@tanstack/react-form"; import { X } from "lucide-react"; -import { useEffect, useMemo, useRef } from "react"; +import { useMemo, useRef } from "react"; import { toast } from "sonner"; import { IconButton } from "@/components/v2/buttons/IconButton"; import { FormLabel } from "@/components/v2/form/FormLabel"; import { Input } from "@/components/v2/inputs/Input"; import { useFeaturesQuery } from "@/hooks/queries/useFeaturesQuery"; +import type { CreditSystemFormInstance } from "../hooks/useCreditSystemForm"; import { FeatureSelectDropdown } from "./FeatureSelectDropdown"; interface ClassicCreditSchemaProps { - creditSystem: CreateFeature; - setCreditSystem: (creditSystem: CreateFeature) => void; + form: CreditSystemFormInstance; } -export function ClassicCreditSchema({ - creditSystem, - setCreditSystem, -}: ClassicCreditSchemaProps) { +export function ClassicCreditSchema({ form }: ClassicCreditSchemaProps) { const { features } = useFeaturesQuery(); - const schema = creditSystem.config?.schema || []; + const config = useStore(form.store, (s) => s.values.config); + const schema: CreditSchemaItem[] = config?.schema || []; + const schemaKeysRef = useRef([]); const schemaKeys = useMemo(() => { const nextKeys = [...schemaKeysRef.current]; @@ -30,13 +30,10 @@ export function ClassicCreditSchema({ while (nextKeys.length > schema.length) { nextKeys.pop(); } + schemaKeysRef.current = nextKeys; return nextKeys; }, [schema.length]); - useEffect(() => { - schemaKeysRef.current = schemaKeys; - }, [schemaKeys]); - const allMeteredFeatures = features.filter( (feature: Feature) => feature.type === FeatureType.Metered, ); @@ -48,26 +45,16 @@ export function ClassicCreditSchema({ ) => { const newSchema = [...schema]; newSchema[index] = { ...newSchema[index], [key]: value }; - setCreditSystem({ - ...creditSystem, - config: { ...creditSystem.config, schema: newSchema }, - }); + form.setFieldValue("config", { ...config, schema: newSchema }); }; const addSchemaItem = () => { schemaKeysRef.current = [...schemaKeysRef.current, crypto.randomUUID()]; const newSchema = [ ...schema, - { - metered_feature_id: "", - feature_amount: 1, - credit_amount: 0, - }, + { metered_feature_id: "", feature_amount: 1, credit_amount: 0 }, ]; - setCreditSystem({ - ...creditSystem, - config: { ...creditSystem.config, schema: newSchema }, - }); + form.setFieldValue("config", { ...config, schema: newSchema }); }; const removeSchemaItem = (index: number) => { @@ -80,10 +67,7 @@ export function ClassicCreditSchema({ schemaKeysRef.current = nextKeys; const newSchema = [...schema]; newSchema.splice(index, 1); - setCreditSystem({ - ...creditSystem, - config: { ...creditSystem.config, schema: newSchema }, - }); + form.setFieldValue("config", { ...config, schema: newSchema }); }; return ( diff --git a/vite/src/views/products/features/credit-systems/components/CreditSystemDetails.tsx b/vite/src/views/products/features/credit-systems/components/CreditSystemDetails.tsx index b1c834f36..1eb66b454 100644 --- a/vite/src/views/products/features/credit-systems/components/CreditSystemDetails.tsx +++ b/vite/src/views/products/features/credit-systems/components/CreditSystemDetails.tsx @@ -1,23 +1,16 @@ -import type { CreateFeature } from "@autumn/shared"; +import { useStore } from "@tanstack/react-form"; import { FormLabel } from "@/components/v2/form/FormLabel"; import { Input } from "@/components/v2/inputs/Input"; import { SheetSection } from "@/components/v2/sheets/SharedSheetComponents"; -import { useAutoSlug } from "@/hooks/common/useAutoSlug"; +import type { CreditSystemFormInstance } from "../hooks/useCreditSystemForm"; interface CreditSystemDetailsProps { - creditSystem: CreateFeature; - setCreditSystem: (creditSystem: CreateFeature) => void; + form: CreditSystemFormInstance; } -export function CreditSystemDetails({ - creditSystem, - setCreditSystem, -}: CreditSystemDetailsProps) { - const { setSource, setTarget } = useAutoSlug({ - setState: setCreditSystem, - sourceKey: "name", - targetKey: "id", - }); +export function CreditSystemDetails({ form }: CreditSystemDetailsProps) { + const name = useStore(form.store, (s) => s.values.name); + const id = useStore(form.store, (s) => s.values.id); return ( @@ -26,19 +19,31 @@ export function CreditSystemDetails({ Name setSource(e.target.value)} + value={name} + onChange={(e) => { + form.setFieldValue("name", e.target.value); + if (!id || id === slugify(name)) { + form.setFieldValue("id", slugify(e.target.value)); + } + }} />
ID setTarget(e.target.value)} + value={id} + onChange={(e) => form.setFieldValue("id", e.target.value)} />
); } + +function slugify(text: string): string { + return text + .toLowerCase() + .replace(/[^a-z0-9]+/g, "_") + .replace(/^_+|_+$/g, ""); +} diff --git a/vite/src/views/products/features/credit-systems/components/CreditSystemSchema.tsx b/vite/src/views/products/features/credit-systems/components/CreditSystemSchema.tsx index cb3a082fe..86136ab90 100644 --- a/vite/src/views/products/features/credit-systems/components/CreditSystemSchema.tsx +++ b/vite/src/views/products/features/credit-systems/components/CreditSystemSchema.tsx @@ -1,8 +1,10 @@ -import { type CreateFeature, FeatureType, type ModelsDevProvider } from "@autumn/shared"; +import { FeatureType, type ModelsDevProvider } 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"; @@ -18,8 +20,8 @@ const getReleaseDateMs = (releaseDate?: string) => { function getDefaultModelMarkups( providers: Record, -): Record { - const result: Record = {}; +): Record { + const result: Record = {}; const preferredProvider = providers["openrouter"] ?? Object.values(providers)[0]; if (!preferredProvider) return result; @@ -34,60 +36,50 @@ function getDefaultModelMarkups( [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[`${providerKey}/${modelKey}`] = { - markup: 0, - }; + result[`${providerKey}/${modelKey}`] = {}; } return result; } interface CreditSystemSchemaProps { - creditSystem: CreateFeature; - setCreditSystem: (creditSystem: CreateFeature) => void; + form: CreditSystemFormInstance; disableModeSwitch?: boolean; } export function CreditSystemSchema({ - creditSystem, - setCreditSystem, + form, disableModeSwitch = false, }: CreditSystemSchemaProps) { const { providers } = useModelsDevPricing(); + const type = useStore(form.store, (s) => s.values.type); const mode: CreditSchemaMode = - creditSystem.type === FeatureType.AiCreditSystem ? "ai" : "classic"; + type === FeatureType.AiCreditSystem ? "ai" : "classic"; const handleModeChange = (newMode: string) => { if (newMode === "ai") { const modelMarkups = getDefaultModelMarkups(providers); - setCreditSystem({ - ...creditSystem, - type: FeatureType.AiCreditSystem, - config: { ...creditSystem.config, schema: [] }, - model_markups: Object.keys(modelMarkups).length > 0 ? modelMarkups : {}, - }); + form.setFieldValue("type", FeatureType.AiCreditSystem); + form.setFieldValue("config", { ...form.state.values.config, schema: [] }); + form.setFieldValue( + "model_markups", + Object.keys(modelMarkups).length > 0 ? modelMarkups : {}, + ); } else { - setCreditSystem({ - ...creditSystem, - type: FeatureType.CreditSystem, - config: { - ...creditSystem.config, - schema: [ - { metered_feature_id: "", feature_amount: 1, credit_amount: 0 }, - ], - }, - model_markups: null, + form.setFieldValue("type", FeatureType.CreditSystem); + form.setFieldValue("config", { + ...form.state.values.config, + schema: [{ metered_feature_id: "", feature_amount: 1, credit_amount: 0 }], }); + form.setFieldValue("model_markups", {}); } }; @@ -105,7 +97,7 @@ export function CreditSystemSchema({ withSeparator={false} description={ mode === "ai" - ? "Select AI models and set a markup on top of their base pricing" + ? "Select AI models and set a markup on top of their base pricing. All prices in $/M tokens." : "When you track usage for these features, the value will be multiplied by the credit cost, then deducted from the balance" } > @@ -120,16 +112,9 @@ export function CreditSystemSchema({ )} {mode === "classic" ? ( - + ) : ( - + )}
diff --git a/vite/src/views/products/features/credit-systems/components/UpdateCreditSystemSheet.tsx b/vite/src/views/products/features/credit-systems/components/UpdateCreditSystemSheet.tsx index de9b800bc..504c827f9 100644 --- a/vite/src/views/products/features/credit-systems/components/UpdateCreditSystemSheet.tsx +++ b/vite/src/views/products/features/credit-systems/components/UpdateCreditSystemSheet.tsx @@ -1,7 +1,7 @@ -import type { CreateFeature, CreditSchemaItem, Feature } from "@autumn/shared"; +import type { CreditSchemaItem, Feature } from "@autumn/shared"; import { FeatureType } from "@autumn/shared"; +import { useStore } from "@tanstack/react-form"; import type { AxiosError } from "axios"; -import { useEffect, useState } from "react"; import { toast } from "sonner"; import { ShortcutButton } from "@/components/v2/buttons/ShortcutButton"; import { @@ -13,6 +13,7 @@ import { useFeaturesQuery } from "@/hooks/queries/useFeaturesQuery"; import { FeatureService } from "@/services/FeatureService"; import { useAxiosInstance } from "@/services/useAxiosInstance"; import { getBackendErr } from "@/utils/genUtils"; +import { useCreditSystemForm } from "../hooks/useCreditSystemForm"; import { validateCreditSystem } from "../utils/validateCreditSystem"; import { CreditSystemDetails } from "./CreditSystemDetails"; import { CreditSystemSchema } from "./CreditSystemSchema"; @@ -30,131 +31,102 @@ function UpdateCreditSystemSheet({ selectedCreditSystem, onSuccess, }: UpdateCreditSystemSheetProps) { - const [loading, setLoading] = useState(false); - const [creditSystem, setCreditSystem] = useState({ - name: "", - id: "", - type: FeatureType.CreditSystem, - config: { - schema: [ - { - metered_feature_id: "", - feature_amount: 1, - credit_amount: 0, - }, - ], - }, - event_names: [], - }); - const axiosInstance = useAxiosInstance(); const { refetch } = useFeaturesQuery(); - // Initialize credit system when selectedCreditSystem changes - useEffect(() => { - if (open && selectedCreditSystem) { - setCreditSystem({ - name: selectedCreditSystem.name, - id: selectedCreditSystem.id, - type: selectedCreditSystem.type, - config: selectedCreditSystem.config, - event_names: selectedCreditSystem.event_names, - model_markups: selectedCreditSystem.model_markups, - }); - } - }, [open, selectedCreditSystem]); + const form = useCreditSystemForm({ + feature: open ? selectedCreditSystem : null, + onSubmit: async (values) => { + if (!selectedCreditSystem) return; - const handleUpdateCreditSystem = async () => { - if (!selectedCreditSystem) return; + const creditSystem = { + name: values.name, + id: values.id, + type: values.type, + config: values.config, + event_names: values.event_names, + model_markups: values.model_markups, + }; - const validationError = validateCreditSystem(creditSystem); - if (validationError) { - toast.error(validationError); - return; - } + const validationError = validateCreditSystem(creditSystem); + if (validationError) { + toast.error(validationError); + return; + } - setLoading(true); - try { - const isAiCreditSystem = creditSystem.type === FeatureType.AiCreditSystem; + const isAiCreditSystem = values.type === FeatureType.AiCreditSystem; + + const finalMarkups = { ...values.model_markups }; + if (isAiCreditSystem) { + for (const [key, entry] of Object.entries(finalMarkups)) { + if (entry.markup == null) { + finalMarkups[key] = { ...entry, markup: values.defaultMarkup }; + } + } + } await FeatureService.updateFeature( axiosInstance, selectedCreditSystem.id, { - id: creditSystem.id, - name: creditSystem.name, - type: creditSystem.type, - model_markups: creditSystem.model_markups ?? undefined, + id: values.id, + name: values.name, + type: values.type, + model_markups: isAiCreditSystem ? finalMarkups : undefined, credit_schema: isAiCreditSystem ? undefined - : creditSystem.config?.schema?.map((x: CreditSchemaItem) => ({ + : values.config?.schema?.map((x: CreditSchemaItem) => ({ metered_feature_id: x.metered_feature_id, - credit_cost: - x.credit_amount != null ? Number(x.credit_amount) : 0, + credit_cost: x.credit_amount != null ? Number(x.credit_amount) : 0, })), - event_names: creditSystem.event_names, + event_names: values.event_names, display: undefined, }, ); await refetch(); toast.success("Credit system updated successfully"); - - if (onSuccess) { - onSuccess( - selectedCreditSystem.id, - creditSystem.id || selectedCreditSystem.id, - ); - } - + onSuccess?.(selectedCreditSystem.id, values.id || selectedCreditSystem.id); setOpen(false); - } catch (error: unknown) { - toast.error( - getBackendErr(error as AxiosError, "Failed to update credit system"), - ); - } finally { - setLoading(false); - } - }; + }, + }); - const handleCancel = () => { - setOpen(false); - }; + const isSubmitting = useStore(form.store, (s) => s.isSubmitting); return ( - +
- - + +
setOpen(false)} singleShortcut="escape" > Cancel + form.handleSubmit().catch((err: AxiosError) => { + toast.error(getBackendErr(err, "Failed to update credit system")); + }) + } metaShortcut="enter" - isLoading={loading} + isLoading={isSubmitting} > Update credit system diff --git a/vite/src/views/products/features/credit-systems/hooks/useAiCreditSchema.ts b/vite/src/views/products/features/credit-systems/hooks/useAiCreditSchema.ts deleted file mode 100644 index 0f3001eb0..000000000 --- a/vite/src/views/products/features/credit-systems/hooks/useAiCreditSchema.ts +++ /dev/null @@ -1,259 +0,0 @@ -import type { CreateFeature } from "@autumn/shared"; -import { useCallback, useEffect, useMemo, useRef, useState } from "react"; -import { toast } from "sonner"; -import { useModelsDevPricing } from "@/hooks/queries/useAiModelsQuery"; - -type ModelMarkupEntry = { - markup: number; - input_cost?: number; - output_cost?: number; -}; - -type ModelMarkups = Record; - -function groupByProvider(modelMarkups: ModelMarkups) { - const groups: Record = {}; - for (const fullId of Object.keys(modelMarkups)) { - const [providerKey] = fullId.split("/"); - if (!groups[providerKey]) groups[providerKey] = []; - groups[providerKey].push(fullId); - } - return groups; -} - -export function useAiCreditSchema({ - creditSystem, - setCreditSystem, -}: { - creditSystem: CreateFeature; - setCreditSystem: ( - creditSystem: CreateFeature | ((prev: CreateFeature) => CreateFeature), - ) => void; -}) { - const { - providers, - isLoading: modelsLoading, - error: modelsError, - } = useModelsDevPricing(); - - const modelMarkups = creditSystem.model_markups ?? {}; - - const [defaultMarkup, setDefaultMarkup] = useState(0); - const manuallyEditedModels = useRef>(new Set()); - - useEffect(() => { - if (modelsError) { - toast.error("Models.dev pricing is unavailable. Try again later."); - } - }, [modelsError]); - - const providerGroups = useMemo( - () => groupByProvider(modelMarkups), - [modelMarkups], - ); - const activeProviderKeys = Object.keys(providerGroups); - - const availableProviders = useMemo(() => { - const filtered = Object.values(providers).filter( - (provider) => !activeProviderKeys.includes(provider.id), - ); - if (!activeProviderKeys.includes("custom")) { - filtered.push({ id: "custom", name: "Custom", models: {} }); - } - return filtered; - }, [providers, activeProviderKeys]); - - const updateMarkups = useCallback( - (updatedMarkups: ModelMarkups) => { - setCreditSystem((prev) => ({ - ...prev, - model_markups: updatedMarkups, - })); - }, - [setCreditSystem], - ); - - const handleModelChange = useCallback( - (providerKey: string, oldModelKey: string, newModelKey: string) => { - const oldFullId = `${providerKey}/${oldModelKey}`; - const newFullId = `${providerKey}/${newModelKey}`; - if (oldFullId !== newFullId && newFullId in modelMarkups) return; - const updatedMarkups = { ...modelMarkups }; - const oldEntry = updatedMarkups[oldFullId]; - const markup = oldEntry?.markup ?? 0; - - if (manuallyEditedModels.current.has(oldFullId)) { - manuallyEditedModels.current.delete(oldFullId); - manuallyEditedModels.current.add(newFullId); - } - - delete updatedMarkups[oldFullId]; - if (providerKey === "custom") { - updatedMarkups[newFullId] = { - markup, - input_cost: oldEntry?.input_cost ?? 0, - output_cost: oldEntry?.output_cost ?? 0, - }; - } else { - updatedMarkups[newFullId] = { markup }; - } - updateMarkups(updatedMarkups); - }, - [modelMarkups, providers, updateMarkups], - ); - - const handleMarkupChange = useCallback( - (providerKey: string, modelKey: string, markup: number) => { - const fullId = `${providerKey}/${modelKey}`; - manuallyEditedModels.current.add(fullId); - updateMarkups({ - ...modelMarkups, - [fullId]: { ...modelMarkups[fullId], markup }, - }); - }, - [modelMarkups, updateMarkups], - ); - - const handleDefaultMarkupChange = useCallback( - (value: number) => { - setDefaultMarkup(value); - const updatedMarkups = { ...modelMarkups }; - for (const modelId of Object.keys(updatedMarkups)) { - if (!manuallyEditedModels.current.has(modelId)) { - updatedMarkups[modelId] = { - ...updatedMarkups[modelId], - markup: value, - }; - } - } - updateMarkups(updatedMarkups); - }, - [modelMarkups, updateMarkups], - ); - - const handleCostChange = useCallback( - ( - providerKey: string, - modelKey: string, - field: "input_cost" | "output_cost", - value: number, - ) => { - const fullId = `${providerKey}/${modelKey}`; - updateMarkups({ - ...modelMarkups, - [fullId]: { ...modelMarkups[fullId], [field]: value }, - }); - }, - [modelMarkups, updateMarkups], - ); - - const handleRemoveModel = useCallback( - (providerKey: string, modelKey: string) => { - const fullId = `${providerKey}/${modelKey}`; - manuallyEditedModels.current.delete(fullId); - const updatedMarkups = { ...modelMarkups }; - delete updatedMarkups[fullId]; - updateMarkups(updatedMarkups); - }, - [modelMarkups, updateMarkups], - ); - - const handleRemoveProvider = useCallback( - (providerKey: string) => { - const updatedMarkups = { ...modelMarkups }; - for (const fullId of providerGroups[providerKey] ?? []) { - manuallyEditedModels.current.delete(fullId); - delete updatedMarkups[fullId]; - } - updateMarkups(updatedMarkups); - }, - [modelMarkups, providerGroups, updateMarkups], - ); - - const addProvider = useCallback( - (providerKey: string) => { - if (providerKey === "custom") { - const existingKeys = (providerGroups.custom ?? []).map((fullId) => { - const [, ...parts] = fullId.split("/"); - return parts.join("/"); - }); - let counter = 1; - while (existingKeys.includes(`model-${counter}`)) counter++; - updateMarkups({ - ...modelMarkups, - [`custom/model-${counter}`]: { - markup: defaultMarkup, - input_cost: 0, - output_cost: 0, - }, - }); - return; - } - const provider = providers[providerKey]; - if (!provider) return; - const firstModelKey = Object.keys(provider.models)[0]; - if (!firstModelKey) return; - const fullId = `${providerKey}/${firstModelKey}`; - updateMarkups({ - ...modelMarkups, - [fullId]: { markup: defaultMarkup }, - }); - }, - [defaultMarkup, modelMarkups, providerGroups, providers, updateMarkups], - ); - - const addModelToProvider = useCallback( - (providerKey: string) => { - if (providerKey === "custom") { - const existingKeys = (providerGroups.custom ?? []).map((fullId) => { - const [, ...parts] = fullId.split("/"); - return parts.join("/"); - }); - let counter = 1; - while (existingKeys.includes(`model-${counter}`)) counter++; - const fullId = `custom/model-${counter}`; - updateMarkups({ - ...modelMarkups, - [fullId]: { markup: defaultMarkup, input_cost: 0, output_cost: 0 }, - }); - return; - } - const provider = providers[providerKey]; - if (!provider) return; - const usedModelKeys = new Set( - (providerGroups[providerKey] ?? []).map((fullId) => { - const [, ...parts] = fullId.split("/"); - return parts.join("/"); - }), - ); - const nextModelKey = Object.keys(provider.models).find( - (key) => !usedModelKeys.has(key), - ); - if (!nextModelKey) return; - const fullId = `${providerKey}/${nextModelKey}`; - updateMarkups({ - ...modelMarkups, - [fullId]: { markup: defaultMarkup }, - }); - }, - [defaultMarkup, modelMarkups, providerGroups, providers, updateMarkups], - ); - - return { - providers, - modelsLoading, - modelMarkups, - defaultMarkup, - providerGroups, - activeProviderKeys, - availableProviders, - handleModelChange, - handleMarkupChange, - handleDefaultMarkupChange, - handleCostChange, - handleRemoveModel, - handleRemoveProvider, - addProvider, - addModelToProvider, - }; -} diff --git a/vite/src/views/products/features/credit-systems/hooks/useCreditSystemForm.ts b/vite/src/views/products/features/credit-systems/hooks/useCreditSystemForm.ts new file mode 100644 index 000000000..478546aff --- /dev/null +++ b/vite/src/views/products/features/credit-systems/hooks/useCreditSystemForm.ts @@ -0,0 +1,41 @@ +import type { Feature, ModelMarkups } from "@autumn/shared"; +import { FeatureType } from "@autumn/shared"; +import { useAppForm } from "@/hooks/form/form"; + +export interface CreditSystemFormValues { + name: string; + id: string; + type: FeatureType; + config: Record; + event_names: string[]; + model_markups: NonNullable; + defaultMarkup: number; +} + +export function useCreditSystemForm({ + feature, + onSubmit, + onChange, +}: { + feature: Feature | null; + onSubmit?: (values: CreditSystemFormValues) => Promise; + onChange?: (values: CreditSystemFormValues) => void; +}) { + return useAppForm({ + defaultValues: { + name: feature?.name ?? "", + id: feature?.id ?? "", + type: feature?.type ?? FeatureType.CreditSystem, + config: feature?.config ?? { schema: [{ metered_feature_id: "", feature_amount: 1, credit_amount: 0 }] }, + event_names: feature?.event_names ?? [], + model_markups: (feature?.model_markups as CreditSystemFormValues["model_markups"]) ?? {}, + defaultMarkup: 0, + } satisfies CreditSystemFormValues, + onSubmit: onSubmit ? ({ value }) => onSubmit(value) : undefined, + listeners: onChange + ? { onChange: ({ formApi }) => onChange(formApi.state.values) } + : undefined, + }); +} + +export type CreditSystemFormInstance = ReturnType; diff --git a/vite/src/views/products/features/feature-list/CreditListColumns.tsx b/vite/src/views/products/features/feature-list/CreditListColumns.tsx index fa94fbc2f..34ba83739 100644 --- a/vite/src/views/products/features/feature-list/CreditListColumns.tsx +++ b/vite/src/views/products/features/feature-list/CreditListColumns.tsx @@ -1,4 +1,5 @@ -import type { Feature, ModelsDevProvider } from "@autumn/shared"; +import { type Feature, FeatureType, type ModelsDevProvider } from "@autumn/shared"; +import { CoinsIcon, CpuIcon } from "@phosphor-icons/react"; import type { ColumnDef, Row } from "@tanstack/react-table"; import { AdminHover } from "@/components/general/AdminHover"; import { MiniCopyButton } from "@/components/v2/buttons/CopyButton"; @@ -48,9 +49,32 @@ export const createCreditListColumns = ( ); }, }, + { + header: "Type", + size: 120, + accessorKey: "type", + cell: ({ row }: { row: Row }) => { + const isAi = row.original.type === FeatureType.AiCreditSystem; + return ( +
+ {isAi ? ( + <> + + AI + + ) : ( + <> + + Standard + + )} +
+ ); + }, + }, { header: "Features", - size: 250, + size: 200, accessorKey: "features", cell: ({ row }: { row: Row }) => { const creditSystem = row.original; diff --git a/vite/src/views/products/features/feature-list/FeatureListTable.tsx b/vite/src/views/products/features/feature-list/FeatureListTable.tsx index 4bddbba2a..a34d65026 100644 --- a/vite/src/views/products/features/feature-list/FeatureListTable.tsx +++ b/vite/src/views/products/features/feature-list/FeatureListTable.tsx @@ -29,15 +29,19 @@ export function FeatureListTable() { // Filter features and credit systems based on archived state const { regularFeatures, creditSystems, hasEventNames } = useMemo(() => { + const isCreditType = (type: string) => + type === FeatureType.CreditSystem || + type === FeatureType.AiCreditSystem; + const regularFeatures = features?.filter((feature) => { - if (feature.type === FeatureType.CreditSystem) return false; + if (isCreditType(feature.type)) return false; return queryStates.showArchivedFeatures ? feature.archived : !feature.archived; }); const creditSystems = features?.filter((feature) => { - if (feature.type !== FeatureType.CreditSystem) return false; + if (!isCreditType(feature.type)) return false; return queryStates.showArchivedFeatures ? feature.archived : !feature.archived; diff --git a/vite/src/views/products/features/utils/getFeatureIcon.tsx b/vite/src/views/products/features/utils/getFeatureIcon.tsx index 264e4c86c..6b5ce42a6 100644 --- a/vite/src/views/products/features/utils/getFeatureIcon.tsx +++ b/vite/src/views/products/features/utils/getFeatureIcon.tsx @@ -7,6 +7,7 @@ import { import { BatteryHighIcon, CoinsIcon, + CpuIcon, TicketIcon, ToggleRightIcon, } from "@phosphor-icons/react"; @@ -68,6 +69,18 @@ export const getFeatureIconConfig = ( }; } + // Handle AI credit system + if ( + typeStr === FeatureType.AiCreditSystem || + typeStr === "ai_credit_system" + ) { + return { + icon: , + color: "text-yellow-500", + label: "AI Credit System", + }; + } + // Handle credit system if (typeStr === FeatureType.CreditSystem || typeStr === "credit_system") { return { diff --git a/vite/src/views/products/plan/components/edit-plan-feature/IncludedUsage.tsx b/vite/src/views/products/plan/components/edit-plan-feature/IncludedUsage.tsx index 8b4528277..4b391abd6 100644 --- a/vite/src/views/products/plan/components/edit-plan-feature/IncludedUsage.tsx +++ b/vite/src/views/products/plan/components/edit-plan-feature/IncludedUsage.tsx @@ -69,9 +69,8 @@ export function IncludedUsage() { )}
- {isAiCreditSystem ? ( - // Couldn't find a disabled property but data-disabled is accounted for in CSS - + {isAiCreditSystem ? ( + $ void; +}) { + const form = useCreditSystemForm({ + feature: { + internal_id: "", + org_id: "", + created_at: 0, + env: "sandbox" as any, + id: feature.id ?? "", + name: feature.name ?? "", + type: feature.type, + config: feature.config ?? {}, + archived: false, + event_names: feature.event_names ?? [], + model_markups: feature.model_markups ?? null, + }, + onChange: (values) => + setFeature({ + ...feature, + type: values.type, + config: values.config, + model_markups: values.model_markups, + }), + }); + + return ; +} export function NewFeatureBehaviour({ feature, @@ -15,13 +49,11 @@ export function NewFeatureBehaviour({ feature: CreateFeature; setFeature: (feature: CreateFeature) => void; }) { - if (feature.type === FeatureType.CreditSystem) { - return ( - - ); + if ( + feature.type === FeatureType.CreditSystem || + feature.type === FeatureType.AiCreditSystem + ) { + return ; } if (feature.type === FeatureType.Metered) { From 072459b963d0304c8d947e96a93b0d88bc1d27e6 Mon Sep 17 00:00:00 2001 From: Charlie Lamb Date: Wed, 20 May 2026 16:11:58 +0100 Subject: [PATCH 07/46] chore: resolve conflicts --- .../components/ai-elements/prompt-input.tsx | 30 +-- .../table/table-content-virtualized.tsx | 27 ++- .../components/AuthorizedApps.tsx | 212 ++++++++++-------- .../org-dropdown/components/OrgLogo.tsx | 12 +- 4 files changed, 160 insertions(+), 121 deletions(-) diff --git a/vite/src/components/ai-elements/prompt-input.tsx b/vite/src/components/ai-elements/prompt-input.tsx index 73d24bb4f..baee45e1d 100644 --- a/vite/src/components/ai-elements/prompt-input.tsx +++ b/vite/src/components/ai-elements/prompt-input.tsx @@ -968,21 +968,15 @@ export const PromptInputButton = ({ className, size = "icon", ...props -<<<<<<< HEAD -}: PromptInputButtonProps) => { - const newSize = - size ?? (Children.count(props.children) > 1 ? "sm" : "icon-sm"); - - return ( - - ); -}; +}: PromptInputButtonProps) => ( + -
- )) - )} -
-
- - + Revoke Access @@ -247,15 +277,14 @@ export const AuthorizedApps = () => { ? This app will no longer be able to access your organization. - - {loadingApiKeys ? ( + {loadingApiKeys ? (
Checking for linked API keys...
) : linkedApiKeys.length > 0 ? (
- + The following API keys will also be deleted:
@@ -268,7 +297,9 @@ export const AuthorizedApps = () => { {key.prefix}...
- {key.name} + + {key.name} + {
) : null} -
-
+ ); }; diff --git a/vite/src/views/main-sidebar/org-dropdown/components/OrgLogo.tsx b/vite/src/views/main-sidebar/org-dropdown/components/OrgLogo.tsx index faf599de2..b0b87dcc0 100644 --- a/vite/src/views/main-sidebar/org-dropdown/components/OrgLogo.tsx +++ b/vite/src/views/main-sidebar/org-dropdown/components/OrgLogo.tsx @@ -1,13 +1,21 @@ import type { FrontendOrg } from "@autumn/shared"; - +import { cn } from "@/lib/utils"; export const OrgLogo = ({ org }: { org: FrontendOrg }) => { const firstLetter = org?.name?.charAt(0).toUpperCase() || "A"; return ( -
+
+ {/* {org.logo ? ( + {org.name} + ) : ( */} {firstLetter} + {/* )} */}
); }; From fb9f2a1456b1a01a13c882f13e6a5218436da89c Mon Sep 17 00:00:00 2001 From: Charlie Lamb Date: Wed, 20 May 2026 16:26:01 +0100 Subject: [PATCH 08/46] chore: form/hook cleanup --- .../components/AiCreditSchema.tsx | 58 ++++---------- .../components/AiCreditSchemaTable.tsx | 22 +----- .../components/ClassicCreditSchema.tsx | 66 +++------------- .../components/CreditSystemDetails.tsx | 8 +- .../credit-systems/hooks/useAiProviders.ts | 78 +++++++++++++++++++ .../credit-systems/hooks/useCreditSchema.ts | 66 ++++++++++++++++ 6 files changed, 172 insertions(+), 126 deletions(-) create mode 100644 vite/src/views/products/features/credit-systems/hooks/useAiProviders.ts create mode 100644 vite/src/views/products/features/credit-systems/hooks/useCreditSchema.ts diff --git a/vite/src/views/products/features/credit-systems/components/AiCreditSchema.tsx b/vite/src/views/products/features/credit-systems/components/AiCreditSchema.tsx index 624d818f3..f45798d27 100644 --- a/vite/src/views/products/features/credit-systems/components/AiCreditSchema.tsx +++ b/vite/src/views/products/features/credit-systems/components/AiCreditSchema.tsx @@ -1,59 +1,27 @@ import { PlusIcon } from "lucide-react"; -import { useMemo } from "react"; -import { useStore } from "@tanstack/react-form"; import { FormLabel } from "@/components/v2/form/FormLabel"; import { Input } from "@/components/v2/inputs/Input"; import { SearchableSelect } from "@/components/v2/selects/SearchableSelect"; -import { useModelsDevPricing } from "@/hooks/queries/useAiModelsQuery"; import type { CreditSystemFormInstance } from "../hooks/useCreditSystemForm"; +import { useAiProviders } from "../hooks/useAiProviders"; import { AiCreditSchemaTable } from "./AiCreditSchemaTable"; interface AiCreditSchemaProps { form: CreditSystemFormInstance; } -function groupByProvider(markups: Record) { - const groups: Record = {}; - for (const fullId of Object.keys(markups)) { - const [provider] = fullId.split("/"); - (groups[provider] ??= []).push(fullId); - } - return groups; -} - export function AiCreditSchema({ form }: AiCreditSchemaProps) { - const { providers, isLoading } = useModelsDevPricing(); - const modelMarkups = useStore(form.store, (s) => s.values.model_markups); - const defaultMarkup = useStore(form.store, (s) => s.values.defaultMarkup); - - const providerGroups = useMemo(() => groupByProvider(modelMarkups), [modelMarkups]); - const activeProviderKeys = Object.keys(providerGroups); - - const availableProviders = useMemo(() => { - const filtered = Object.values(providers).filter( - (p) => !activeProviderKeys.includes(p.id), - ); - if (!activeProviderKeys.includes("custom")) { - filtered.push({ id: "custom", name: "Custom", models: {} }); - } - return filtered; - }, [providers, activeProviderKeys]); - - const addProvider = (providerKey: string) => { - form.setFieldValue("model_markups", (prev) => { - if (providerKey === "custom") { - const existing = Object.keys(prev).filter((k) => k.startsWith("custom/")); - let i = 1; - while (existing.includes(`custom/model-${i}`)) i++; - return { ...prev, [`custom/model-${i}`]: { input_cost: 0, output_cost: 0 } }; - } - const provider = providers[providerKey]; - if (!provider) return prev; - const firstKey = Object.keys(provider.models)[0]; - if (!firstKey) return prev; - return { ...prev, [`${providerKey}/${firstKey}`]: {} }; - }); - }; + const { + providers, + isLoading, + defaultMarkup, + providerGroups, + activeProviderKeys, + availableProviders, + addProvider, + removeKeys, + renameKey, + } = useAiProviders(form); return (
@@ -89,6 +57,8 @@ export function AiCreditSchema({ form }: AiCreditSchemaProps) { modelFullIds={modelFullIds} provider={provider ?? { id: providerKey, name: providerKey, models: {} }} isLoading={isLoading} + removeKeys={removeKeys} + renameKey={renameKey} /> ); })} diff --git a/vite/src/views/products/features/credit-systems/components/AiCreditSchemaTable.tsx b/vite/src/views/products/features/credit-systems/components/AiCreditSchemaTable.tsx index 7fe1b5a43..22de92d83 100644 --- a/vite/src/views/products/features/credit-systems/components/AiCreditSchemaTable.tsx +++ b/vite/src/views/products/features/credit-systems/components/AiCreditSchemaTable.tsx @@ -27,6 +27,8 @@ interface AiCreditSchemaTableProps { modelFullIds: string[]; provider: ModelsDevProvider; isLoading: boolean; + removeKeys: (keys: string[]) => void; + renameKey: (oldKey: string, newKey: string) => void; } function formatCost(value: number | null | undefined): string { @@ -41,27 +43,11 @@ export function AiCreditSchemaTable({ modelFullIds, provider, isLoading, + removeKeys, + renameKey, }: AiCreditSchemaTableProps) { const isCustom = providerKey === "custom"; - - const removeKeys = (keys: string[]) => - form.setFieldValue("model_markups", (prev) => { - const updated = { ...prev }; - for (const k of keys) delete updated[k]; - return updated; - }); - - const renameKey = (oldKey: string, newKey: string) => - form.setFieldValue("model_markups", (prev) => { - if (newKey in prev) return prev; - const updated = { ...prev }; - const entry = updated[oldKey]; - delete updated[oldKey]; - updated[newKey] = { ...entry }; - return updated; - }); - const data: ModelRow[] = useMemo( () => modelFullIds.map((fullId) => { diff --git a/vite/src/views/products/features/credit-systems/components/ClassicCreditSchema.tsx b/vite/src/views/products/features/credit-systems/components/ClassicCreditSchema.tsx index 4542b62c2..fa1a21dd7 100644 --- a/vite/src/views/products/features/credit-systems/components/ClassicCreditSchema.tsx +++ b/vite/src/views/products/features/credit-systems/components/ClassicCreditSchema.tsx @@ -1,15 +1,11 @@ import type { CreditSchemaItem, Feature } from "@autumn/shared"; -import { FeatureType } from "@autumn/shared"; import { PlusIcon } from "@phosphor-icons/react"; -import { useStore } from "@tanstack/react-form"; import { X } from "lucide-react"; -import { useMemo, useRef } from "react"; -import { toast } from "sonner"; import { IconButton } from "@/components/v2/buttons/IconButton"; import { FormLabel } from "@/components/v2/form/FormLabel"; import { Input } from "@/components/v2/inputs/Input"; -import { useFeaturesQuery } from "@/hooks/queries/useFeaturesQuery"; import type { CreditSystemFormInstance } from "../hooks/useCreditSystemForm"; +import { useCreditSchema } from "../hooks/useCreditSchema"; import { FeatureSelectDropdown } from "./FeatureSelectDropdown"; interface ClassicCreditSchemaProps { @@ -17,58 +13,14 @@ interface ClassicCreditSchemaProps { } export function ClassicCreditSchema({ form }: ClassicCreditSchemaProps) { - const { features } = useFeaturesQuery(); - const config = useStore(form.store, (s) => s.values.config); - const schema: CreditSchemaItem[] = config?.schema || []; - - const schemaKeysRef = useRef([]); - const schemaKeys = useMemo(() => { - const nextKeys = [...schemaKeysRef.current]; - while (nextKeys.length < schema.length) { - nextKeys.push(crypto.randomUUID()); - } - while (nextKeys.length > schema.length) { - nextKeys.pop(); - } - schemaKeysRef.current = nextKeys; - return nextKeys; - }, [schema.length]); - - const allMeteredFeatures = features.filter( - (feature: Feature) => feature.type === FeatureType.Metered, - ); - - const handleSchemaChange = ( - index: number, - key: keyof CreditSchemaItem, - value: string | number, - ) => { - const newSchema = [...schema]; - newSchema[index] = { ...newSchema[index], [key]: value }; - form.setFieldValue("config", { ...config, schema: newSchema }); - }; - - const addSchemaItem = () => { - schemaKeysRef.current = [...schemaKeysRef.current, crypto.randomUUID()]; - const newSchema = [ - ...schema, - { metered_feature_id: "", feature_amount: 1, credit_amount: 0 }, - ]; - form.setFieldValue("config", { ...config, schema: newSchema }); - }; - - const removeSchemaItem = (index: number) => { - if (schema.length === 1) { - toast.error("There must be at least one item in the credit system"); - return; - } - const nextKeys = [...schemaKeysRef.current]; - nextKeys.splice(index, 1); - schemaKeysRef.current = nextKeys; - const newSchema = [...schema]; - newSchema.splice(index, 1); - form.setFieldValue("config", { ...config, schema: newSchema }); - }; + const { + schema, + schemaKeys, + allMeteredFeatures, + handleSchemaChange, + addSchemaItem, + removeSchemaItem, + } = useCreditSchema(form); return (
diff --git a/vite/src/views/products/features/credit-systems/components/CreditSystemDetails.tsx b/vite/src/views/products/features/credit-systems/components/CreditSystemDetails.tsx index 1eb66b454..7e273c280 100644 --- a/vite/src/views/products/features/credit-systems/components/CreditSystemDetails.tsx +++ b/vite/src/views/products/features/credit-systems/components/CreditSystemDetails.tsx @@ -2,6 +2,7 @@ import { useStore } from "@tanstack/react-form"; import { FormLabel } from "@/components/v2/form/FormLabel"; import { Input } from "@/components/v2/inputs/Input"; import { SheetSection } from "@/components/v2/sheets/SharedSheetComponents"; +import { slugify } from "@/utils/formatUtils/formatTextUtils"; import type { CreditSystemFormInstance } from "../hooks/useCreditSystemForm"; interface CreditSystemDetailsProps { @@ -40,10 +41,3 @@ export function CreditSystemDetails({ form }: CreditSystemDetailsProps) { ); } - -function slugify(text: string): string { - return text - .toLowerCase() - .replace(/[^a-z0-9]+/g, "_") - .replace(/^_+|_+$/g, ""); -} diff --git a/vite/src/views/products/features/credit-systems/hooks/useAiProviders.ts b/vite/src/views/products/features/credit-systems/hooks/useAiProviders.ts new file mode 100644 index 000000000..30aaafd46 --- /dev/null +++ b/vite/src/views/products/features/credit-systems/hooks/useAiProviders.ts @@ -0,0 +1,78 @@ +import type { ModelsDevProvider } from "@autumn/shared"; +import { useStore } from "@tanstack/react-form"; +import { useMemo } from "react"; +import { useModelsDevPricing } from "@/hooks/queries/useAiModelsQuery"; +import type { CreditSystemFormInstance } from "./useCreditSystemForm"; + +function groupByProvider(markups: Record) { + const groups: Record = {}; + for (const fullId of Object.keys(markups)) { + const [provider] = fullId.split("/"); + (groups[provider] ??= []).push(fullId); + } + return groups; +} + +export function useAiProviders(form: CreditSystemFormInstance) { + const { providers, isLoading } = useModelsDevPricing(); + const modelMarkups = useStore(form.store, (s) => s.values.model_markups); + const defaultMarkup = useStore(form.store, (s) => s.values.defaultMarkup); + + const providerGroups = useMemo(() => groupByProvider(modelMarkups), [modelMarkups]); + const activeProviderKeys = Object.keys(providerGroups); + + const availableProviders = useMemo(() => { + const filtered = Object.values(providers).filter( + (p) => !activeProviderKeys.includes(p.id), + ); + if (!activeProviderKeys.includes("custom")) { + filtered.push({ id: "custom", name: "Custom", models: {} } as ModelsDevProvider); + } + return filtered; + }, [providers, activeProviderKeys]); + + const addProvider = (providerKey: string) => { + form.setFieldValue("model_markups", (prev) => { + if (providerKey === "custom") { + const existing = Object.keys(prev).filter((k) => k.startsWith("custom/")); + let i = 1; + while (existing.includes(`custom/model-${i}`)) i++; + return { ...prev, [`custom/model-${i}`]: { input_cost: 0, output_cost: 0 } }; + } + const provider = providers[providerKey]; + if (!provider) return prev; + const firstKey = Object.keys(provider.models)[0]; + if (!firstKey) return prev; + return { ...prev, [`${providerKey}/${firstKey}`]: {} }; + }); + }; + + const removeKeys = (keys: string[]) => + form.setFieldValue("model_markups", (prev) => { + const updated = { ...prev }; + for (const k of keys) delete updated[k]; + return updated; + }); + + const renameKey = (oldKey: string, newKey: string) => + form.setFieldValue("model_markups", (prev) => { + if (newKey in prev) return prev; + const updated = { ...prev }; + const entry = updated[oldKey]; + delete updated[oldKey]; + updated[newKey] = { ...entry }; + return updated; + }); + + return { + providers, + isLoading, + defaultMarkup, + providerGroups, + activeProviderKeys, + availableProviders, + addProvider, + removeKeys, + renameKey, + }; +} diff --git a/vite/src/views/products/features/credit-systems/hooks/useCreditSchema.ts b/vite/src/views/products/features/credit-systems/hooks/useCreditSchema.ts new file mode 100644 index 000000000..672cbb39e --- /dev/null +++ b/vite/src/views/products/features/credit-systems/hooks/useCreditSchema.ts @@ -0,0 +1,66 @@ +import type { CreditSchemaItem, Feature } from "@autumn/shared"; +import { FeatureType } from "@autumn/shared"; +import { useStore } from "@tanstack/react-form"; +import { useMemo, useRef } from "react"; +import { toast } from "sonner"; +import { useFeaturesQuery } from "@/hooks/queries/useFeaturesQuery"; +import type { CreditSystemFormInstance } from "./useCreditSystemForm"; + +export function useCreditSchema(form: CreditSystemFormInstance) { + const { features } = useFeaturesQuery(); + const config = useStore(form.store, (s) => s.values.config); + const schema: CreditSchemaItem[] = config?.schema || []; + + const schemaKeysRef = useRef([]); + const schemaKeys = useMemo(() => { + const nextKeys = [...schemaKeysRef.current]; + while (nextKeys.length < schema.length) nextKeys.push(crypto.randomUUID()); + while (nextKeys.length > schema.length) nextKeys.pop(); + schemaKeysRef.current = nextKeys; + return nextKeys; + }, [schema.length]); + + const allMeteredFeatures = features.filter( + (f: Feature) => f.type === FeatureType.Metered, + ); + + const handleSchemaChange = ( + index: number, + key: keyof CreditSchemaItem, + value: string | number, + ) => { + const newSchema = [...schema]; + newSchema[index] = { ...newSchema[index], [key]: value }; + form.setFieldValue("config", { ...config, schema: newSchema }); + }; + + const addSchemaItem = () => { + schemaKeysRef.current = [...schemaKeysRef.current, crypto.randomUUID()]; + form.setFieldValue("config", { + ...config, + schema: [...schema, { metered_feature_id: "", feature_amount: 1, credit_amount: 0 }], + }); + }; + + const removeSchemaItem = (index: number) => { + if (schema.length === 1) { + toast.error("There must be at least one item in the credit system"); + return; + } + const nextKeys = [...schemaKeysRef.current]; + nextKeys.splice(index, 1); + schemaKeysRef.current = nextKeys; + const newSchema = [...schema]; + newSchema.splice(index, 1); + form.setFieldValue("config", { ...config, schema: newSchema }); + }; + + return { + schema, + schemaKeys, + allMeteredFeatures, + handleSchemaChange, + addSchemaItem, + removeSchemaItem, + }; +} From 78dddfea01665c13e3fec51b13a49df98f965677 Mon Sep 17 00:00:00 2001 From: Charlie Lamb Date: Wed, 20 May 2026 16:34:50 +0100 Subject: [PATCH 09/46] chore: ai credit table cleanup --- .../components/AiCreditSchemaTable.tsx | 100 +----------------- .../components/CustomModelInput.tsx | 23 ++++ .../components/EditableNumberCell.tsx | 78 ++++++++++++++ .../edit-plan-feature/IncludedUsage.tsx | 22 +--- 4 files changed, 109 insertions(+), 114 deletions(-) create mode 100644 vite/src/views/products/features/credit-systems/components/CustomModelInput.tsx create mode 100644 vite/src/views/products/features/credit-systems/components/EditableNumberCell.tsx diff --git a/vite/src/views/products/features/credit-systems/components/AiCreditSchemaTable.tsx b/vite/src/views/products/features/credit-systems/components/AiCreditSchemaTable.tsx index 22de92d83..c5fef3d33 100644 --- a/vite/src/views/products/features/credit-systems/components/AiCreditSchemaTable.tsx +++ b/vite/src/views/products/features/credit-systems/components/AiCreditSchemaTable.tsx @@ -2,11 +2,11 @@ import type { ModelsDevProvider } from "@autumn/shared"; import type { ColumnDef, Row } from "@tanstack/react-table"; import { useStore } from "@tanstack/react-form"; import { InfoIcon, PlusIcon, X } from "lucide-react"; -import { useMemo, useState } from "react"; +import { useMemo } from "react"; import { Table } from "@/components/general/table"; import { IconButton } from "@/components/v2/buttons/IconButton"; -import { Input } from "@/components/v2/inputs/Input"; import { + Tooltip, TooltipContent, TooltipTrigger, @@ -14,6 +14,8 @@ import { import { useProductTable } from "@/views/products/hooks/useProductTable"; import type { CreditSystemFormInstance } from "../hooks/useCreditSystemForm"; import { AiModelSelectDropdown } from "./AiModelSelectDropdown"; +import { CustomModelInput } from "./CustomModelInput"; +import { EditableNumberCell } from "./EditableNumberCell"; interface ModelRow { fullId: string; @@ -259,97 +261,3 @@ export function AiCreditSchemaTable({
); } - -function CustomModelInput({ - modelKey, - onRename, -}: { - modelKey: string; - onRename: (newKey: string) => void; -}) { - const [local, setLocal] = useState(modelKey); - return ( - setLocal(e.target.value)} - onBlur={() => { - if (local !== modelKey) onRename(local); - }} - placeholder="my-model-id" - className="text-sm" - /> - ); -} - -function EditableNumberCell({ - form, - fullId, - field, - useDefaultAsPlaceholder = false, - allowUndefined = false, -}: { - form: CreditSystemFormInstance; - fullId: string; - field: "markup" | "input_cost" | "output_cost"; - useDefaultAsPlaceholder?: boolean; - allowUndefined?: boolean; -}) { - const currentValue = useStore( - form.store, - (s) => s.values.model_markups[fullId]?.[field], - ); - const placeholder = useStore(form.store, (s) => - useDefaultAsPlaceholder ? String(s.values.defaultMarkup) : "0", - ); - const [local, setLocal] = useState(""); - const [focused, setFocused] = useState(false); - - const hasValue = allowUndefined ? currentValue != null && currentValue !== 0 : currentValue != null; - const displayed = focused ? local : hasValue ? String(currentValue) : ""; - - return ( - { - const raw = e.target.value; - if (raw === "" || /^-?\d*\.?\d*$/.test(raw)) { - setLocal(raw); - if (raw === "" && allowUndefined) { - form.setFieldValue("model_markups", (prev) => { - const entry = { ...prev[fullId] }; - delete entry[field]; - return { ...prev, [fullId]: entry }; - }); - } else if (raw !== "") { - const parsed = Number(raw); - if (!Number.isNaN(parsed)) { - form.setFieldValue("model_markups", (prev) => ({ - ...prev, - [fullId]: { ...prev[fullId], [field]: parsed }, - })); - } - } - } - }} - onFocus={() => { - setLocal(hasValue ? String(currentValue) : ""); - setFocused(true); - }} - onBlur={() => { - setFocused(false); - if (local === "" && !allowUndefined) { - form.setFieldValue("model_markups", (prev) => ({ - ...prev, - [fullId]: { ...prev[fullId], [field]: 0 }, - })); - } - }} - placeholder={placeholder} - className="text-sm" - /> - ); -} diff --git a/vite/src/views/products/features/credit-systems/components/CustomModelInput.tsx b/vite/src/views/products/features/credit-systems/components/CustomModelInput.tsx new file mode 100644 index 000000000..968abfb4e --- /dev/null +++ b/vite/src/views/products/features/credit-systems/components/CustomModelInput.tsx @@ -0,0 +1,23 @@ +import { useState } from "react"; +import { Input } from "@/components/v2/inputs/Input"; + +interface CustomModelInputProps { + modelKey: string; + onRename: (newKey: string) => void; +} + +export function CustomModelInput({ modelKey, onRename }: CustomModelInputProps) { + const [local, setLocal] = useState(modelKey); + return ( + setLocal(e.target.value)} + onBlur={() => { + if (local !== modelKey) onRename(local); + }} + placeholder="my-model-id" + className="text-sm" + /> + ); +} diff --git a/vite/src/views/products/features/credit-systems/components/EditableNumberCell.tsx b/vite/src/views/products/features/credit-systems/components/EditableNumberCell.tsx new file mode 100644 index 000000000..bd5bd0a8a --- /dev/null +++ b/vite/src/views/products/features/credit-systems/components/EditableNumberCell.tsx @@ -0,0 +1,78 @@ +import { useStore } from "@tanstack/react-form"; +import { useState } from "react"; +import { Input } from "@/components/v2/inputs/Input"; +import type { CreditSystemFormInstance } from "../hooks/useCreditSystemForm"; + +interface EditableNumberCellProps { + form: CreditSystemFormInstance; + fullId: string; + field: "markup" | "input_cost" | "output_cost"; + useDefaultAsPlaceholder?: boolean; + allowUndefined?: boolean; +} + +export function EditableNumberCell({ + form, + fullId, + field, + useDefaultAsPlaceholder = false, + allowUndefined = false, +}: EditableNumberCellProps) { + const currentValue = useStore( + form.store, + (s) => s.values.model_markups[fullId]?.[field], + ); + const placeholder = useStore(form.store, (s) => + useDefaultAsPlaceholder ? String(s.values.defaultMarkup) : "0", + ); + const [local, setLocal] = useState(""); + const [focused, setFocused] = useState(false); + + const hasValue = allowUndefined ? currentValue != null && currentValue !== 0 : currentValue != null; + const displayed = focused ? local : hasValue ? String(currentValue) : ""; + + return ( + { + const raw = e.target.value; + if (raw === "" || /^-?\d*\.?\d*$/.test(raw)) { + setLocal(raw); + if (raw === "" && allowUndefined) { + form.setFieldValue("model_markups", (prev) => { + const entry = { ...prev[fullId] }; + delete entry[field]; + return { ...prev, [fullId]: entry }; + }); + } else if (raw !== "") { + const parsed = Number(raw); + if (!Number.isNaN(parsed)) { + form.setFieldValue("model_markups", (prev) => ({ + ...prev, + [fullId]: { ...prev[fullId], [field]: parsed }, + })); + } + } + } + }} + onFocus={() => { + setLocal(hasValue ? String(currentValue) : ""); + setFocused(true); + }} + onBlur={() => { + setFocused(false); + if (local === "" && !allowUndefined) { + form.setFieldValue("model_markups", (prev) => ({ + ...prev, + [fullId]: { ...prev[fullId], [field]: 0 }, + })); + } + }} + placeholder={placeholder} + className="text-sm" + /> + ); +} diff --git a/vite/src/views/products/plan/components/edit-plan-feature/IncludedUsage.tsx b/vite/src/views/products/plan/components/edit-plan-feature/IncludedUsage.tsx index 4b391abd6..57ed3b92e 100644 --- a/vite/src/views/products/plan/components/edit-plan-feature/IncludedUsage.tsx +++ b/vite/src/views/products/plan/components/edit-plan-feature/IncludedUsage.tsx @@ -49,24 +49,10 @@ export function IncludedUsage() {
- {isAiCreditSystem ? ( - <> - USD budget{" "} - {!isFeaturePrice - ? "allocated to this plan" - : "granted before billing"} - - ) : ( - <> - Quantity of  - - {getFeatureName({ feature, plural: true })}{" "} - - {!isFeaturePrice - ? " that can be used" - : " granted before billing"} - - )} + {isAiCreditSystem + ? `USD budget ${isFeaturePrice ? "granted before billing" : "allocated to this plan"}` + : <>Quantity of {getFeatureName({ feature, plural: true })}{isFeaturePrice ? " granted before billing" : " that can be used"} + }
{isAiCreditSystem ? ( From 0ad885c23e0e657b3634d4358bd4d7e9b85066a6 Mon Sep 17 00:00:00 2001 From: Ridhwan Hussain <73362400+TheUntraceable@users.noreply.github.com.> Date: Tue, 26 May 2026 12:37:58 +0100 Subject: [PATCH 10/46] fix: ai credit system feature config --- server/src/internal/features/featureUtils.ts | 6 +++--- .../handleUpdateFeature/handleUpdateFeatureV2.ts | 1 + shared/utils/featureUtils/apiFeatureToDbFeature.ts | 12 ++++++++++++ .../features/components/CreateFeatureSheet.tsx | 5 ++++- 4 files changed, 20 insertions(+), 4 deletions(-) diff --git a/server/src/internal/features/featureUtils.ts b/server/src/internal/features/featureUtils.ts index 91a0118e1..5683fa0c3 100644 --- a/server/src/internal/features/featureUtils.ts +++ b/server/src/internal/features/featureUtils.ts @@ -45,10 +45,10 @@ export const validateCreditSystem = ( config: CreditSystemConfig, featureType: FeatureType = FeatureType.CreditSystem, ) => { - const schema = config.schema; const isAiCreditSystem = featureType === FeatureType.AiCreditSystem; + const schema = Array.isArray(config?.schema) ? config.schema : []; - if (!isAiCreditSystem && (!schema || schema.length === 0)) { + if (!isAiCreditSystem && schema.length === 0) { throw new RecaseError({ message: `At least one metered feature is required for credit system`, code: ErrCode.InvalidFeature, @@ -68,7 +68,7 @@ export const validateCreditSystem = ( }); } - const newConfig = { ...config, usage_type: FeatureUsageType.Single }; + const newConfig = { ...config, schema, usage_type: FeatureUsageType.Single }; for (let i = 0; i < newConfig.schema.length; i++) { const creditAmount = parseFloat( newConfig.schema[i].credit_amount.toString(), diff --git a/server/src/internal/features/handlers/handleUpdateFeature/handleUpdateFeatureV2.ts b/server/src/internal/features/handlers/handleUpdateFeature/handleUpdateFeatureV2.ts index 80eae8d37..dba75e714 100644 --- a/server/src/internal/features/handlers/handleUpdateFeature/handleUpdateFeatureV2.ts +++ b/server/src/internal/features/handlers/handleUpdateFeature/handleUpdateFeatureV2.ts @@ -55,6 +55,7 @@ export const handleUpdateFeatureV2 = createRoute({ archived: body.archived, event_names: body.event_names, display: body.display, + model_markups: body.model_markups, }, }); diff --git a/shared/utils/featureUtils/apiFeatureToDbFeature.ts b/shared/utils/featureUtils/apiFeatureToDbFeature.ts index dff59220a..4c1d0cb6e 100644 --- a/shared/utils/featureUtils/apiFeatureToDbFeature.ts +++ b/shared/utils/featureUtils/apiFeatureToDbFeature.ts @@ -85,6 +85,13 @@ export const featureV1ToDbFeatureConfig = ({ }) => { const type = apiFeature.type || originalFeature.type; + if (apiFeature.type === FeatureType.AiCreditSystem) { + return { + schema: [], + usage_type: FeatureUsageType.Single, + }; + } + if (nullish(apiFeature.consumable) && nullish(apiFeature.credit_schema)) return; @@ -141,6 +148,11 @@ export const featureV1ToDbFeature = ({ : FeatureUsageType.Continuous; } + if (apiFeature.type === FeatureType.AiCreditSystem) { + newConfig.usage_type = FeatureUsageType.Single; + newConfig.schema = []; + } + if (apiFeature.credit_schema) { newConfig.usage_type = FeatureUsageType.Single; newConfig.schema = apiFeature.credit_schema.map( diff --git a/vite/src/views/products/features/components/CreateFeatureSheet.tsx b/vite/src/views/products/features/components/CreateFeatureSheet.tsx index 527f6a0ea..ab61c4be4 100644 --- a/vite/src/views/products/features/components/CreateFeatureSheet.tsx +++ b/vite/src/views/products/features/components/CreateFeatureSheet.tsx @@ -52,7 +52,10 @@ function CreateFeatureSheet({ const handleCreateFeature = async () => { // Validate credit system specific fields first - if (feature.type === FeatureType.CreditSystem) { + if ( + feature.type === FeatureType.CreditSystem || + feature.type === FeatureType.AiCreditSystem + ) { const validationError = validateCreditSystem(feature); if (validationError) { toast.error(validationError); From abd06a35c591e8bc552c9c3407c79cbf0db52cfe Mon Sep 17 00:00:00 2001 From: Ridhwan Hussain <73362400+TheUntraceable@users.noreply.github.com.> Date: Tue, 26 May 2026 12:39:37 +0100 Subject: [PATCH 11/46] fix: track tokens endpoint --- .../balances/trackTokens.mdx | 2 +- .../api-reference/balances/trackTokens.mdx | 4 +- .../customers/tracking-usage.mdx | 2 +- .../modelling-pricing/credit-systems.mdx | 2 +- .../v2.3/contracts/balancesContract.ts | 59 ++++++++ packages/openapi/v2.3/contracts/index.ts | 2 + .../openapi/v2.3/jsDocs/balancesJsDocs.ts | 29 +++- packages/openapi/v2.3/openapi2.3.ts | 2 + .../balances/handlers/handleTrackTokens.ts | 23 ++-- .../track/handle-track-tokens.test.ts | 126 ++++++++++++++++++ 10 files changed, 236 insertions(+), 15 deletions(-) create mode 100644 server/tests/unit/balances/track/handle-track-tokens.test.ts 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); + }); +}); From 57291890b340126dbf0b4aa0c623368a37ace305 Mon Sep 17 00:00:00 2001 From: Ridhwan Hussain <73362400+TheUntraceable@users.noreply.github.com.> Date: Tue, 26 May 2026 12:40:12 +0100 Subject: [PATCH 12/46] fix: ai sdk token usage fields --- bun.lock | 15 +++++++++++++++ packages/ai-sdk/src/index.ts | 28 +++++++++++++++++++++------- 2 files changed, 36 insertions(+), 7 deletions(-) diff --git a/bun.lock b/bun.lock index 3994fb7d5..1688a2071 100644 --- a/bun.lock +++ b/bun.lock @@ -141,6 +141,19 @@ "typescript": "^6.0.2", }, }, + "packages/ai-sdk": { + "name": "@useautumn/ai-sdk", + "version": "0.0.1", + "devDependencies": { + "@types/node": "^24.9.1", + "tsup": "^8.4.0", + "typescript": "^5.8.3", + }, + "peerDependencies": { + "ai": "^6.0.116", + "autumn-js": "*", + }, + }, "packages/atmn": { "name": "atmn", "version": "1.1.8", @@ -2584,6 +2597,8 @@ "@upstash/redis": ["@upstash/redis@1.38.0", "", { "dependencies": { "uncrypto": "^0.1.3" } }, "sha512-wu+dZBptlLy0+MCUEoHmzrY/TnmgDey3+c7EbIGwrLqAvkP8yi5MWZHYGIFtAygmL4Bkz2TdFu+eU0vFPncIcg=="], + "@useautumn/ai-sdk": ["@useautumn/ai-sdk@workspace:packages/ai-sdk"], + "@useautumn/sdk": ["@useautumn/sdk@workspace:packages/sdk"], "@vdemedes/prettier-config": ["@vdemedes/prettier-config@2.0.1", "", {}, "sha512-lcHyyLfS2ro282qsXKpxw+canUkOlFIGoanxt3BaNCm5K1NR8k4hGvYbFO54/+QWq12d0y/EYRz68yNQkqWFrw=="], diff --git a/packages/ai-sdk/src/index.ts b/packages/ai-sdk/src/index.ts index 544423e5a..2f32eb08f 100644 --- a/packages/ai-sdk/src/index.ts +++ b/packages/ai-sdk/src/index.ts @@ -7,8 +7,17 @@ import { import type { Autumn } from "autumn-js"; type TokenCount = - | LanguageModelV3Usage["inputTokens"] - | LanguageModelUsage["inputTokens"]; + | number + | { + total?: number | null; + } + | null + | undefined; + +type TokenUsage = (LanguageModelV3Usage | LanguageModelUsage) & { + promptTokens?: TokenCount; + completionTokens?: TokenCount; +}; export const withTokenTracking = ({ autumn, @@ -46,15 +55,20 @@ export const withTokenTracking = ({ return value; }; - const trackUsage = async ( - usage: LanguageModelV3Usage | LanguageModelUsage, - ) => { + const trackUsage = async (usage: TokenUsage) => { try { + // @ts-ignore trackTokens is generated from OpenAPI; local autumn-js types may not include it yet. await autumn.balances.trackTokens({ customerId, modelId: modelName, - inputTokens: resolveTokens(usage.inputTokens, "Input"), - outputTokens: resolveTokens(usage.outputTokens, "Output"), + inputTokens: resolveTokens( + usage.inputTokens ?? usage.promptTokens, + "Input", + ), + outputTokens: resolveTokens( + usage.outputTokens ?? usage.completionTokens, + "Output", + ), featureId, entityId, properties, From 4ea52186720355604f19214683f7a7ebbf2379bb Mon Sep 17 00:00:00 2001 From: Ridhwan Hussain <73362400+TheUntraceable@users.noreply.github.com.> Date: Sat, 30 May 2026 18:57:25 +0100 Subject: [PATCH 13/46] fix(credit-systems): sync AI form state to feature store on create --- .../hooks/useCreditSystemForm.ts | 20 +++++++++++++++---- .../new-feature/NewFeatureBehaviour.tsx | 19 +++++++++++++++--- 2 files changed, 32 insertions(+), 7 deletions(-) diff --git a/vite/src/views/products/features/credit-systems/hooks/useCreditSystemForm.ts b/vite/src/views/products/features/credit-systems/hooks/useCreditSystemForm.ts index 478546aff..55bb89b7c 100644 --- a/vite/src/views/products/features/credit-systems/hooks/useCreditSystemForm.ts +++ b/vite/src/views/products/features/credit-systems/hooks/useCreditSystemForm.ts @@ -1,5 +1,7 @@ import type { Feature, ModelMarkups } from "@autumn/shared"; import { FeatureType } from "@autumn/shared"; +import { useStore } from "@tanstack/react-form"; +import { useEffect, useRef } from "react"; import { useAppForm } from "@/hooks/form/form"; export interface CreditSystemFormValues { @@ -21,7 +23,7 @@ export function useCreditSystemForm({ onSubmit?: (values: CreditSystemFormValues) => Promise; onChange?: (values: CreditSystemFormValues) => void; }) { - return useAppForm({ + const form = useAppForm({ defaultValues: { name: feature?.name ?? "", id: feature?.id ?? "", @@ -32,10 +34,20 @@ export function useCreditSystemForm({ defaultMarkup: 0, } satisfies CreditSystemFormValues, onSubmit: onSubmit ? ({ value }) => onSubmit(value) : undefined, - listeners: onChange - ? { onChange: ({ formApi }) => onChange(formApi.state.values) } - : undefined, }); + + // Form-level `listeners.onChange` only fires when a FieldApi instance is + // registered for the changed field (see form-core FormApi.setFieldValue). + // None of these fields are mounted via , so we subscribe to the + // store directly and push value changes out to the caller. + const values = useStore(form.store, (s) => s.values); + const onChangeRef = useRef(onChange); + onChangeRef.current = onChange; + useEffect(() => { + onChangeRef.current?.(values); + }, [values]); + + return form; } export type CreditSystemFormInstance = ReturnType; diff --git a/vite/src/views/products/plan/components/new-feature/NewFeatureBehaviour.tsx b/vite/src/views/products/plan/components/new-feature/NewFeatureBehaviour.tsx index 3485f9497..928f19176 100644 --- a/vite/src/views/products/plan/components/new-feature/NewFeatureBehaviour.tsx +++ b/vite/src/views/products/plan/components/new-feature/NewFeatureBehaviour.tsx @@ -30,13 +30,26 @@ function NewFeatureCreditSchema({ event_names: feature.event_names ?? [], model_markups: feature.model_markups ?? null, }, - onChange: (values) => + onChange: (values) => { + const isAi = values.type === FeatureType.AiCreditSystem; + const materializedMarkups = isAi + ? Object.fromEntries( + Object.entries(values.model_markups ?? {}).map(([key, entry]) => [ + key, + entry?.markup == null + ? { ...entry, markup: values.defaultMarkup } + : entry, + ]), + ) + : values.model_markups; + setFeature({ ...feature, type: values.type, config: values.config, - model_markups: values.model_markups, - }), + model_markups: materializedMarkups, + }); + }, }); return ; From 5b0bda94e5aee2861a2986499c0f3a63c39fb5da Mon Sep 17 00:00:00 2001 From: Ridhwan Hussain <73362400+TheUntraceable@users.noreply.github.com.> Date: Sun, 31 May 2026 00:02:51 +0100 Subject: [PATCH 14/46] fix(ui): guard credit-system totals when subrows missing --- .../CustomerFeatureUsageColumns.tsx | 9 +++++++-- .../CustomerFeatureUsageDisplay.tsx | 17 +++++++++++------ 2 files changed, 18 insertions(+), 8 deletions(-) diff --git a/vite/src/views/customers2/components/table/customer-feature-usage/CustomerFeatureUsageColumns.tsx b/vite/src/views/customers2/components/table/customer-feature-usage/CustomerFeatureUsageColumns.tsx index b69ac65c5..05865724c 100644 --- a/vite/src/views/customers2/components/table/customer-feature-usage/CustomerFeatureUsageColumns.tsx +++ b/vite/src/views/customers2/components/table/customer-feature-usage/CustomerFeatureUsageColumns.tsx @@ -73,6 +73,7 @@ export const CustomerFeatureUsageColumns = [ if (!subRowData && featureType === FeatureType.CreditSystem) { const subRows = (cusEnt as FullCusEntWithSubRows).subRows || []; let totalSpent = 0; + let canCompute = true; for (const subRow of subRows) { if (!("isSubRow" in subRow) || !subRow.isSubRow) continue; @@ -90,11 +91,15 @@ export const CustomerFeatureUsageColumns = [ const subUsed = subTotal - subRemaining; totalSpent += subUsed * creditCost; } + } else { + canCompute = false; } } - const total = allowance * quantity; - balance = total - totalSpent; + if (canCompute) { + const total = allowance * quantity; + balance = total - totalSpent; + } } return ( diff --git a/vite/src/views/customers2/components/table/customer-feature-usage/CustomerFeatureUsageDisplay.tsx b/vite/src/views/customers2/components/table/customer-feature-usage/CustomerFeatureUsageDisplay.tsx index 3525768f2..d1d879782 100644 --- a/vite/src/views/customers2/components/table/customer-feature-usage/CustomerFeatureUsageDisplay.tsx +++ b/vite/src/views/customers2/components/table/customer-feature-usage/CustomerFeatureUsageDisplay.tsx @@ -51,6 +51,7 @@ export function CustomerFeatureUsageDisplay({ if (featureType === FeatureType.CreditSystem) { let totalSpent = 0; + let canCompute = true; for (const subRow of subRows) { // Only process CreditSystemSubRow items @@ -69,16 +70,20 @@ export function CustomerFeatureUsageDisplay({ }); totalSpent += used * creditCost; } + } else { + canCompute = false; } } - const total = allowance * quantity; + if (canCompute) { + const total = allowance * quantity; - return ( -
- {totalSpent}/{total} used -
- ); + return ( +
+ {totalSpent}/{total} used +
+ ); + } } const { total, used } = calculateUsageMetrics({ From 5a5a525df7872f8b3c41d2756bedf49d42b9711c Mon Sep 17 00:00:00 2001 From: Ridhwan Hussain <73362400+TheUntraceable@users.noreply.github.com.> Date: Sun, 31 May 2026 00:03:13 +0100 Subject: [PATCH 15/46] feat(credit-systems): allow ai credit leaves in schemas --- .../features/featureActions/createFeature.ts | 24 +++- .../features/featureActions/updateFeature.ts | 21 ++- server/src/internal/features/featureUtils.ts | 44 ++++++ .../validate-credit-system.test.ts | 125 ++++++++++++++++++ .../components/ClassicCreditSchema.tsx | 77 ++++++----- .../credit-systems/hooks/useCreditSchema.ts | 12 +- 6 files changed, 262 insertions(+), 41 deletions(-) create mode 100644 server/tests/advanced/creditSystems/validate-credit-system.test.ts diff --git a/server/src/internal/features/featureActions/createFeature.ts b/server/src/internal/features/featureActions/createFeature.ts index 4f11bffa3..5e073b6a5 100644 --- a/server/src/internal/features/featureActions/createFeature.ts +++ b/server/src/internal/features/featureActions/createFeature.ts @@ -1,21 +1,37 @@ -import { CreateFeatureSchema, type Feature, FeatureType, type ModelMarkups } from "@autumn/shared"; +import { + CreateFeatureSchema, + type Feature, + FeatureType, + type ModelMarkups, +} from "@autumn/shared"; import type { AutumnContext } from "@/honoUtils/HonoEnv.js"; import { workflows } from "@/queue/workflows.js"; import { generateId } from "@/utils/genUtils.js"; import { FeatureService } from "../FeatureService.js"; import { validateCreditSystem, + validateCreditSystemSchemaReferences, validateMeteredConfig, } from "../featureUtils.js"; -const validateFeature = (data: any) => { +const validateFeature = (data: any, allFeatures: Feature[]) => { const featureType = data.type as FeatureType; let config = data.config; if (featureType === FeatureType.Metered) { config = validateMeteredConfig(config); - } else if (featureType === FeatureType.CreditSystem || featureType === FeatureType.AiCreditSystem) { + } else if ( + featureType === FeatureType.CreditSystem || + featureType === FeatureType.AiCreditSystem + ) { config = validateCreditSystem(config, featureType); + if (featureType === FeatureType.CreditSystem) { + validateCreditSystemSchemaReferences({ + config, + allFeatures, + selfFeatureId: data.id, + }); + } } const parsedFeature = CreateFeatureSchema.parse({ ...data, config }); @@ -44,7 +60,7 @@ export const createFeature = async ({ data, skipGenerateDisplay = false, }: CreateFeatureParams): Promise => { - const parsedFeature = validateFeature(data); + const parsedFeature = validateFeature(data, ctx.features); const feature: Feature = { archived: false, diff --git a/server/src/internal/features/featureActions/updateFeature.ts b/server/src/internal/features/featureActions/updateFeature.ts index 909c4ec4c..c597bbf4b 100644 --- a/server/src/internal/features/featureActions/updateFeature.ts +++ b/server/src/internal/features/featureActions/updateFeature.ts @@ -14,6 +14,7 @@ import RecaseError from "@/utils/errorUtils.js"; import { FeatureService } from "../FeatureService.js"; import { validateCreditSystem, + validateCreditSystemSchemaReferences, validateMeteredConfig, } from "../featureUtils.js"; import { getObjectsUsingFeature } from "../utils/updateFeatureUtils/getObjectsUsingFeature.js"; @@ -181,8 +182,20 @@ export const updateFeature = async ({ if (updates.config === undefined) return feature.config; switch (effectiveType) { case FeatureType.AiCreditSystem: - case FeatureType.CreditSystem: - return validateCreditSystem(updates.config, effectiveType); + case FeatureType.CreditSystem: { + const validatedConfig = validateCreditSystem( + updates.config, + effectiveType, + ); + if (effectiveType === FeatureType.CreditSystem) { + validateCreditSystemSchemaReferences({ + config: validatedConfig, + allFeatures, + selfFeatureId: updates.id ?? feature.id, + }); + } + return validatedConfig; + } case FeatureType.Metered: return validateMeteredConfig(updates.config); default: @@ -217,7 +230,9 @@ export const updateFeature = async ({ } // Queue cache clear for credit system if schema or model markups changed - const isCreditSystem = feature.type === FeatureType.CreditSystem || feature.type === FeatureType.AiCreditSystem; + const isCreditSystem = + feature.type === FeatureType.CreditSystem || + feature.type === FeatureType.AiCreditSystem; if (isCreditSystem && updatedFeature) { const schemaChanged = updates.config != null && diff --git a/server/src/internal/features/featureUtils.ts b/server/src/internal/features/featureUtils.ts index 5683fa0c3..49e737982 100644 --- a/server/src/internal/features/featureUtils.ts +++ b/server/src/internal/features/featureUtils.ts @@ -56,6 +56,14 @@ export const validateCreditSystem = ( }); } + if (isAiCreditSystem && schema.length > 0) { + throw new RecaseError({ + message: `AI credit systems are leaf features and cannot define a schema. Model rates live in model_markups.`, + code: ErrCode.InvalidFeature, + statusCode: 400, + }); + } + const meteredFeatureIds = schema.map( (schemaItem) => schemaItem.metered_feature_id, ); @@ -88,6 +96,42 @@ export const validateCreditSystem = ( return newConfig; }; +/** + * Validates that every feature referenced by a credit system's schema is a + * Metered or AiCreditSystem feature. Rejects nesting one CreditSystem inside + * another — composition is capped at two levels (parent credit system → leaf). + * + * Self-references are tolerated (the create flow doesn't yet have the new id + * in the features list, and updates simply read back as the feature itself). + */ +export const validateCreditSystemSchemaReferences = ({ + config, + allFeatures, + selfFeatureId, +}: { + config: CreditSystemConfig; + allFeatures: Feature[]; + selfFeatureId?: string; +}) => { + const schema = Array.isArray(config?.schema) ? config.schema : []; + if (schema.length === 0) return; + + for (const item of schema) { + const referencedId = item.metered_feature_id; + if (!referencedId || referencedId === selfFeatureId) continue; + + const referenced = allFeatures.find((f) => f.id === referencedId); + if (!referenced) continue; + + if (referenced.type === FeatureType.CreditSystem) { + throw new RecaseError({ + message: `Credit system schema cannot reference another credit system (${referencedId}). Only metered or AI credit features are allowed.`, + code: ErrCode.InvalidFeature, + statusCode: 400, + }); + } + } +}; const getCusFeatureType = ({ feature }: { feature: Feature }) => { if (feature.type === FeatureType.Boolean) { diff --git a/server/tests/advanced/creditSystems/validate-credit-system.test.ts b/server/tests/advanced/creditSystems/validate-credit-system.test.ts new file mode 100644 index 000000000..be4a4e521 --- /dev/null +++ b/server/tests/advanced/creditSystems/validate-credit-system.test.ts @@ -0,0 +1,125 @@ +import { describe, expect, test } from "bun:test"; +import { type Feature, FeatureType, FeatureUsageType } from "@autumn/shared"; +import { + validateCreditSystem, + validateCreditSystemSchemaReferences, +} from "@/internal/features/featureUtils.js"; + +const makeFeature = (id: string, type: FeatureType): Feature => ({ + internal_id: `fe_${id}`, + org_id: "org_test", + created_at: Date.now(), + env: "sandbox" as Feature["env"], + id, + name: id, + type, + config: {}, + archived: false, + event_names: [], + model_markups: null, +}); + +describe("validateCreditSystem — AI credit system schema restrictions", () => { + test("rejects an AI credit system with a non-empty schema", () => { + expect(() => + validateCreditSystem( + { + schema: [{ metered_feature_id: "messages", credit_amount: 1 }] as any, + usage_type: FeatureUsageType.Single, + }, + FeatureType.AiCreditSystem, + ), + ).toThrow(/leaf features/); + }); + + test("allows an AI credit system with an empty schema", () => { + const result = validateCreditSystem( + { + schema: [], + usage_type: FeatureUsageType.Single, + }, + FeatureType.AiCreditSystem, + ); + expect(result.schema).toHaveLength(0); + }); + + test("rejects a regular credit system with empty schema", () => { + expect(() => + validateCreditSystem( + { schema: [], usage_type: FeatureUsageType.Single }, + FeatureType.CreditSystem, + ), + ).toThrow(/At least one metered feature/); + }); +}); + +describe("validateCreditSystemSchemaReferences — cross-feature restrictions", () => { + const metered = makeFeature("messages", FeatureType.Metered); + const aiCredit = makeFeature("ai_credits", FeatureType.AiCreditSystem); + const otherCreditSystem = makeFeature("orbs", FeatureType.CreditSystem); + + test("allows referencing a metered feature", () => { + expect(() => + validateCreditSystemSchemaReferences({ + config: { + schema: [{ metered_feature_id: "messages", credit_amount: 1 } as any], + usage_type: FeatureUsageType.Single, + }, + allFeatures: [metered], + }), + ).not.toThrow(); + }); + + test("allows referencing an AI credit system", () => { + expect(() => + validateCreditSystemSchemaReferences({ + config: { + schema: [ + { metered_feature_id: "ai_credits", credit_amount: 1000 } as any, + ], + usage_type: FeatureUsageType.Single, + }, + allFeatures: [aiCredit], + }), + ).not.toThrow(); + }); + + test("rejects referencing another credit system (prevents nesting)", () => { + expect(() => + validateCreditSystemSchemaReferences({ + config: { + schema: [{ metered_feature_id: "orbs", credit_amount: 1 } as any], + usage_type: FeatureUsageType.Single, + }, + allFeatures: [otherCreditSystem], + }), + ).toThrow(/cannot reference another credit system/); + }); + + test("self-reference (id matches selfFeatureId) is tolerated", () => { + expect(() => + validateCreditSystemSchemaReferences({ + config: { + schema: [{ metered_feature_id: "self_id", credit_amount: 1 } as any], + usage_type: FeatureUsageType.Single, + }, + allFeatures: [makeFeature("self_id", FeatureType.CreditSystem)], + selfFeatureId: "self_id", + }), + ).not.toThrow(); + }); + + test("dangling reference (id not in allFeatures) is tolerated", () => { + expect(() => + validateCreditSystemSchemaReferences({ + config: { + schema: [ + { metered_feature_id: "nonexistent", credit_amount: 1 } as any, + ], + usage_type: FeatureUsageType.Single, + }, + allFeatures: [], + }), + ).not.toThrow(); + }); +}); diff --git a/vite/src/views/products/features/credit-systems/components/ClassicCreditSchema.tsx b/vite/src/views/products/features/credit-systems/components/ClassicCreditSchema.tsx index fa1a21dd7..7a5936a95 100644 --- a/vite/src/views/products/features/credit-systems/components/ClassicCreditSchema.tsx +++ b/vite/src/views/products/features/credit-systems/components/ClassicCreditSchema.tsx @@ -1,11 +1,15 @@ -import type { CreditSchemaItem, Feature } from "@autumn/shared"; +import { + type CreditSchemaItem, + type Feature, + FeatureType, +} from "@autumn/shared"; import { PlusIcon } from "@phosphor-icons/react"; import { X } from "lucide-react"; import { IconButton } from "@/components/v2/buttons/IconButton"; import { FormLabel } from "@/components/v2/form/FormLabel"; import { Input } from "@/components/v2/inputs/Input"; -import type { CreditSystemFormInstance } from "../hooks/useCreditSystemForm"; import { useCreditSchema } from "../hooks/useCreditSchema"; +import type { CreditSystemFormInstance } from "../hooks/useCreditSystemForm"; import { FeatureSelectDropdown } from "./FeatureSelectDropdown"; interface ClassicCreditSchemaProps { @@ -16,7 +20,7 @@ export function ClassicCreditSchema({ form }: ClassicCreditSchemaProps) { const { schema, schemaKeys, - allMeteredFeatures, + allSchemaCandidateFeatures, handleSchemaChange, addSchemaItem, removeSchemaItem, @@ -25,13 +29,13 @@ export function ClassicCreditSchema({ form }: ClassicCreditSchemaProps) { return (
- Metered Feature + Feature Credit Cost
{schema.map((item: CreditSchemaItem, index: number) => { - const availableFeatures = allMeteredFeatures.filter( + const availableFeatures = allSchemaCandidateFeatures.filter( (feature: Feature) => !schema.some( (schemaItem: CreditSchemaItem) => @@ -40,6 +44,12 @@ export function ClassicCreditSchema({ form }: ClassicCreditSchemaProps) { ), ); + const selectedFeature = allSchemaCandidateFeatures.find( + (f: Feature) => f.id === item.metered_feature_id, + ); + const isAiChild = + selectedFeature?.type === FeatureType.AiCreditSystem; + return (
-
- - handleSchemaChange(index, "credit_amount", e.target.value) - } - onBlur={(e) => - handleSchemaChange( - index, - "credit_amount", - Number(e.target.value) || 0, - ) - } - placeholder="eg. 10" - /> - } - onClick={() => removeSchemaItem(index)} - /> +
+
+ + handleSchemaChange(index, "credit_amount", e.target.value) + } + onBlur={(e) => + handleSchemaChange( + index, + "credit_amount", + Number(e.target.value) || 0, + ) + } + placeholder="eg. 10" + /> + } + onClick={() => removeSchemaItem(index)} + /> +
+ {isAiChild && ( + + credits per $1 of AI usage + + )}
); @@ -86,7 +103,7 @@ export function ClassicCreditSchema({ form }: ClassicCreditSchemaProps) { = allMeteredFeatures.length} + disabled={schema.length >= allSchemaCandidateFeatures.length} className="w-fit mt-4" icon={} > diff --git a/vite/src/views/products/features/credit-systems/hooks/useCreditSchema.ts b/vite/src/views/products/features/credit-systems/hooks/useCreditSchema.ts index 672cbb39e..fa805ac04 100644 --- a/vite/src/views/products/features/credit-systems/hooks/useCreditSchema.ts +++ b/vite/src/views/products/features/credit-systems/hooks/useCreditSchema.ts @@ -20,8 +20,9 @@ export function useCreditSchema(form: CreditSystemFormInstance) { return nextKeys; }, [schema.length]); - const allMeteredFeatures = features.filter( - (f: Feature) => f.type === FeatureType.Metered, + const allSchemaCandidateFeatures = features.filter( + (f: Feature) => + f.type === FeatureType.Metered || f.type === FeatureType.AiCreditSystem, ); const handleSchemaChange = ( @@ -38,7 +39,10 @@ export function useCreditSchema(form: CreditSystemFormInstance) { schemaKeysRef.current = [...schemaKeysRef.current, crypto.randomUUID()]; form.setFieldValue("config", { ...config, - schema: [...schema, { metered_feature_id: "", feature_amount: 1, credit_amount: 0 }], + schema: [ + ...schema, + { metered_feature_id: "", feature_amount: 1, credit_amount: 0 }, + ], }); }; @@ -58,7 +62,7 @@ export function useCreditSchema(form: CreditSystemFormInstance) { return { schema, schemaKeys, - allMeteredFeatures, + allSchemaCandidateFeatures, handleSchemaChange, addSchemaItem, removeSchemaItem, From 0b9510d1b508e973f290e230a8ee2a98eaf742f7 Mon Sep 17 00:00:00 2001 From: Ridhwan Hussain <73362400+TheUntraceable@users.noreply.github.com.> Date: Sun, 31 May 2026 00:03:54 +0100 Subject: [PATCH 16/46] feat(credit-systems): map token costs to parent systems --- .../balances/track/utils/runRedisTrack.ts | 39 +++++++++ .../balances/track/v3/runRedisTrackV3.ts | 29 +++++++ .../utils/deduction/computeCreditCosts.ts | 23 ++++-- .../track/basic/track-tokens-orbs.test.ts | 80 +++++++++++++++++++ server/tests/setup/v2Features.ts | 13 +++ 5 files changed, 178 insertions(+), 6 deletions(-) create mode 100644 server/tests/integration/balances/track/basic/track-tokens-orbs.test.ts diff --git a/server/src/internal/balances/track/utils/runRedisTrack.ts b/server/src/internal/balances/track/utils/runRedisTrack.ts index babd09206..01455f1cd 100644 --- a/server/src/internal/balances/track/utils/runRedisTrack.ts +++ b/server/src/internal/balances/track/utils/runRedisTrack.ts @@ -17,6 +17,36 @@ import type { FeatureDeduction } from "../../utils/types/featureDeduction.js"; import type { RolloverUpdate } from "../../utils/types/rolloverUpdate.js"; import { handleRedisTrackError } from "./handleRedisTrackError.js"; +const buildAiCreditCostProperty = ({ + featureDeductions, + updates, + fullCustomer, +}: { + featureDeductions: FeatureDeduction[]; + updates: Record; + fullCustomer: FullCustomer; +}): Record | undefined => { + const aiDeduction = featureDeductions.find((d) => d.tokenUsage); + if (!aiDeduction) return; + + const cusEntIdToFeatureId = new Map(); + for (const cp of fullCustomer.customer_products) { + for (const ce of cp.customer_entitlements ?? []) { + cusEntIdToFeatureId.set(ce.id, ce.entitlement.feature.id); + } + } + + const creditCost: Record = {}; + for (const [cusEntId, update] of Object.entries(updates)) { + const featureId = cusEntIdToFeatureId.get(cusEntId); + if (!featureId || featureId === aiDeduction.feature.id) continue; + if (update.deducted === 0) continue; + creditCost[featureId] = (creditCost[featureId] ?? 0) + update.deducted; + } + + return Object.keys(creditCost).length > 0 ? creditCost : undefined; +}; + const queueSyncItem = ({ ctx, body, @@ -111,6 +141,15 @@ export const runRedisTrack = async ({ const { updates, fullCus, rolloverUpdates } = result; + const aiCreditCost = buildAiCreditCostProperty({ + featureDeductions, + updates, + fullCustomer, + }); + if (aiCreditCost) { + body.properties = { ...(body.properties ?? {}), credit_cost: aiCreditCost }; + } + // Queue sync and event queueSyncItem({ ctx, diff --git a/server/src/internal/balances/track/v3/runRedisTrackV3.ts b/server/src/internal/balances/track/v3/runRedisTrackV3.ts index 494a393b6..c571e575d 100644 --- a/server/src/internal/balances/track/v3/runRedisTrackV3.ts +++ b/server/src/internal/balances/track/v3/runRedisTrackV3.ts @@ -22,6 +22,27 @@ import type { FeatureDeduction } from "../../utils/types/featureDeduction.js"; import type { RolloverUpdate } from "../../utils/types/rolloverUpdate.js"; import { handleRedisTrackErrorV3 } from "./handleRedisTrackErrorV3.js"; +const buildAiCreditCostProperty = ({ + featureDeductions, + deductions, +}: { + featureDeductions: FeatureDeduction[]; + deductions: TrackDeduction[]; +}): Record | undefined => { + const aiDeduction = featureDeductions.find((d) => d.tokenUsage); + if (!aiDeduction) return; + + const creditCost: Record = {}; + for (const deduction of deductions) { + if (deduction.feature_id === aiDeduction.feature.id) continue; + if (!deduction.value) continue; + creditCost[deduction.feature_id] = + (creditCost[deduction.feature_id] ?? 0) + deduction.value; + } + + return Object.keys(creditCost).length > 0 ? creditCost : undefined; +}; + const queueSyncItem = ({ ctx, body, @@ -148,6 +169,14 @@ export const runRedisTrackV3 = async ({ mutationLogs, }); + const aiCreditCost = buildAiCreditCostProperty({ + featureDeductions, + deductions, + }); + if (aiCreditCost) { + body.properties = { ...(body.properties ?? {}), credit_cost: aiCreditCost }; + } + queueEvent({ ctx, body, fullSubject, deductions, internalProductId }); const { balance, balances } = await deductionToTrackResponseV2({ diff --git a/server/src/internal/balances/utils/deduction/computeCreditCosts.ts b/server/src/internal/balances/utils/deduction/computeCreditCosts.ts index 05889e041..0ffe954a9 100644 --- a/server/src/internal/balances/utils/deduction/computeCreditCosts.ts +++ b/server/src/internal/balances/utils/deduction/computeCreditCosts.ts @@ -20,20 +20,31 @@ export const computeCreditCosts = async ({ }): Promise => { const costMap = new Map(); - if (deduction.precomputedCreditCost != null) { - const cost = deduction.precomputedCreditCost; - return () => cost; - } - const tokens = deduction.tokenUsage - ? { input: deduction.tokenUsage.inputTokens, output: deduction.tokenUsage.outputTokens } + ? { + input: deduction.tokenUsage.inputTokens, + output: deduction.tokenUsage.outputTokens, + } : undefined; await Promise.all( cusEnts.map(async (ce) => { + // Precomputed cost (from /track/tokens) is in the AI credit feature's + // native unit (USD). It applies 1:1 to that feature's own entitlement, + // but parent credit systems still need their schema ratio applied — + // fall through to getCreditCost with amount = precomputed cost. + if ( + deduction.precomputedCreditCost != null && + ce.entitlement.feature.id === deduction.feature.id + ) { + costMap.set(ce.id, deduction.precomputedCreditCost); + return; + } + const creditCost = await getCreditCost({ featureId: deduction.feature.id, creditSystem: ce.entitlement.feature, + amount: deduction.precomputedCreditCost, modelName: deduction.tokenUsage?.modelName, tokens, }); diff --git a/server/tests/integration/balances/track/basic/track-tokens-orbs.test.ts b/server/tests/integration/balances/track/basic/track-tokens-orbs.test.ts new file mode 100644 index 000000000..36dcde8b8 --- /dev/null +++ b/server/tests/integration/balances/track/basic/track-tokens-orbs.test.ts @@ -0,0 +1,80 @@ +import { expect, test } from "bun:test"; + +import type { ApiCustomerV3, TrackResponseV2 } from "@autumn/shared"; +import { TestFeature } from "@tests/setup/v2Features.js"; +import { items } from "@tests/utils/fixtures/items.js"; +import { products } from "@tests/utils/fixtures/products.js"; +import { initScenario, s } from "@tests/utils/testInitUtils/initScenario.js"; +import chalk from "chalk"; +import { Decimal } from "decimal.js"; + +// ═══════════════════════════════════════════════════════════════════ +// TRACK-TOKENS-ORBS: AI credit system nested inside a parent credit system +// Verifies that a single /track/tokens call deducts USD from the AI credit +// feature AND deducts the ratio-mapped amount from any parent credit +// system whose schema references it. +// ═══════════════════════════════════════════════════════════════════ + +test.concurrent( + `${chalk.yellowBright("track-tokens-orbs: AI credit system inside parent credit system deducts both balances")}`, + async () => { + const aiCreditsItem = items.free({ + featureId: TestFeature.AiCredits, + includedUsage: 100, // $100 of AI usage + }); + const orbsItem = items.free({ + featureId: TestFeature.Orbs, + includedUsage: 50_000, // 50,000 orbs + }); + const freeProd = products.base({ + id: "free", + items: [aiCreditsItem, orbsItem], + }); + + const { customerId, autumnV1, autumnV2 } = await initScenario({ + customerId: "track-tokens-orbs", + setup: [ + s.customer({ testClock: false }), + s.products({ list: [freeProd] }), + ], + actions: [s.attach({ productId: freeProd.id })], + }); + + // custom/internal-model: input_cost=5 $/M, output_cost=15 $/M, markup=0% + const inputTokens = 10_000; + const outputTokens = 5_000; + const expectedUsdCost = new Decimal(5) + .mul(inputTokens) + .add(new Decimal(15).mul(outputTokens)) + .div(1_000_000) + .toNumber(); // 0.125 + + // Orbs schema: 1000 orbs per $1 of AI usage + const expectedOrbsCost = new Decimal(expectedUsdCost).mul(1000).toNumber(); // 125 + + const trackRes: TrackResponseV2 = await autumnV2.post("/track_tokens", { + customer_id: customerId, + feature_id: TestFeature.AiCredits, + model_id: "custom/internal-model", + input_tokens: inputTokens, + output_tokens: outputTokens, + }); + + expect(trackRes.customer_id).toBe(customerId); + expect(trackRes.value).toBeCloseTo(expectedUsdCost, 10); + + const customer = await autumnV1.customers.get(customerId); + + // AI credit feature balance dropped by USD cost (1:1) + expect(customer.features[TestFeature.AiCredits]).toMatchObject({ + balance: new Decimal(100).minus(expectedUsdCost).toNumber(), + usage: expectedUsdCost, + }); + + // Parent orbs balance dropped by USD cost × 1000 + expect(customer.features[TestFeature.Orbs]).toMatchObject({ + balance: new Decimal(50_000).minus(expectedOrbsCost).toNumber(), + usage: expectedOrbsCost, + }); + }, +); diff --git a/server/tests/setup/v2Features.ts b/server/tests/setup/v2Features.ts index 2a4150b83..dd66e05aa 100644 --- a/server/tests/setup/v2Features.ts +++ b/server/tests/setup/v2Features.ts @@ -29,6 +29,8 @@ export enum TestFeature { AiCredits = "ai_credits", // AI credit system (models.dev pricing) AiCredits2 = "ai_credits_2", // second AI credit system (for disambiguation tests) + + Orbs = "orbs", // credit system that wraps an AI credit system (1000 orbs per $1) } export const getFeatures = ({ orgId }: { orgId: string }) => ({ @@ -158,4 +160,15 @@ export const getFeatures = ({ orgId }: { orgId: string }) => ({ }, }, }), + [TestFeature.Orbs]: constructCreditSystem({ + featureId: TestFeature.Orbs, + orgId, + env: AppEnv.Sandbox, + schema: [ + { + metered_feature_id: TestFeature.AiCredits, + credit_cost: 1000, // 1000 orbs per $1 of AI usage + }, + ], + }), }); From 7833dc29f9e9b9254c1504216ffd62dd11166e2a Mon Sep 17 00:00:00 2001 From: Ridhwan Hussain <73362400+TheUntraceable@users.noreply.github.com.> Date: Mon, 1 Jun 2026 14:45:12 +0100 Subject: [PATCH 17/46] chore: dedupe buildAiCreditCostProperty --- .../track/utils/buildAiCreditCostProperty.ts | 28 +++++++++++++++++++ .../balances/track/utils/runRedisTrack.ts | 23 ++++++--------- .../balances/track/v3/runRedisTrackV3.ts | 27 ++++-------------- 3 files changed, 41 insertions(+), 37 deletions(-) create mode 100644 server/src/internal/balances/track/utils/buildAiCreditCostProperty.ts diff --git a/server/src/internal/balances/track/utils/buildAiCreditCostProperty.ts b/server/src/internal/balances/track/utils/buildAiCreditCostProperty.ts new file mode 100644 index 000000000..66156e522 --- /dev/null +++ b/server/src/internal/balances/track/utils/buildAiCreditCostProperty.ts @@ -0,0 +1,28 @@ +import type { FeatureDeduction } from "../../utils/types/featureDeduction.js"; + +/** + * Aggregates the amounts deducted from each metered feature into a `credit_cost` + * map for an AI credit-system track, so the event records what each feature was + * charged. Returns undefined when the track isn't an AI credit deduction or + * nothing chargeable was deducted. The AI feature's own deduction is excluded so + * the map only contains the downstream metered features it consumed. + */ +export const buildAiCreditCostProperty = ({ + featureDeductions, + entries, +}: { + featureDeductions: FeatureDeduction[]; + entries: Array<{ featureId: string; amount: number }>; +}): Record | undefined => { + const aiDeduction = featureDeductions.find((d) => d.tokenUsage); + if (!aiDeduction) return; + + const creditCost: Record = {}; + for (const { featureId, amount } of entries) { + if (featureId === aiDeduction.feature.id) continue; + if (!amount) continue; + creditCost[featureId] = (creditCost[featureId] ?? 0) + amount; + } + + return Object.keys(creditCost).length > 0 ? creditCost : undefined; +}; diff --git a/server/src/internal/balances/track/utils/runRedisTrack.ts b/server/src/internal/balances/track/utils/runRedisTrack.ts index 01455f1cd..9f5763c81 100644 --- a/server/src/internal/balances/track/utils/runRedisTrack.ts +++ b/server/src/internal/balances/track/utils/runRedisTrack.ts @@ -15,20 +15,16 @@ import { globalSyncBatchingManagerV2 } from "../../utils/sync/SyncBatchingManage import type { DeductionUpdate } from "../../utils/types/deductionUpdate.js"; import type { FeatureDeduction } from "../../utils/types/featureDeduction.js"; import type { RolloverUpdate } from "../../utils/types/rolloverUpdate.js"; +import { buildAiCreditCostProperty } from "./buildAiCreditCostProperty.js"; import { handleRedisTrackError } from "./handleRedisTrackError.js"; -const buildAiCreditCostProperty = ({ - featureDeductions, +const aiCreditCostEntries = ({ updates, fullCustomer, }: { - featureDeductions: FeatureDeduction[]; updates: Record; fullCustomer: FullCustomer; -}): Record | undefined => { - const aiDeduction = featureDeductions.find((d) => d.tokenUsage); - if (!aiDeduction) return; - +}): Array<{ featureId: string; amount: number }> => { const cusEntIdToFeatureId = new Map(); for (const cp of fullCustomer.customer_products) { for (const ce of cp.customer_entitlements ?? []) { @@ -36,15 +32,13 @@ const buildAiCreditCostProperty = ({ } } - const creditCost: Record = {}; + const entries: Array<{ featureId: string; amount: number }> = []; for (const [cusEntId, update] of Object.entries(updates)) { const featureId = cusEntIdToFeatureId.get(cusEntId); - if (!featureId || featureId === aiDeduction.feature.id) continue; - if (update.deducted === 0) continue; - creditCost[featureId] = (creditCost[featureId] ?? 0) + update.deducted; + if (!featureId) continue; + entries.push({ featureId, amount: update.deducted }); } - - return Object.keys(creditCost).length > 0 ? creditCost : undefined; + return entries; }; const queueSyncItem = ({ @@ -143,8 +137,7 @@ export const runRedisTrack = async ({ const aiCreditCost = buildAiCreditCostProperty({ featureDeductions, - updates, - fullCustomer, + entries: aiCreditCostEntries({ updates, fullCustomer }), }); if (aiCreditCost) { body.properties = { ...(body.properties ?? {}), credit_cost: aiCreditCost }; diff --git a/server/src/internal/balances/track/v3/runRedisTrackV3.ts b/server/src/internal/balances/track/v3/runRedisTrackV3.ts index c571e575d..fa7d99ef8 100644 --- a/server/src/internal/balances/track/v3/runRedisTrackV3.ts +++ b/server/src/internal/balances/track/v3/runRedisTrackV3.ts @@ -20,29 +20,9 @@ import { import { globalSyncBatchingManagerV3 } from "@/internal/balances/utils/sync/SyncBatchingManagerV3.js"; import type { FeatureDeduction } from "../../utils/types/featureDeduction.js"; import type { RolloverUpdate } from "../../utils/types/rolloverUpdate.js"; +import { buildAiCreditCostProperty } from "../utils/buildAiCreditCostProperty.js"; import { handleRedisTrackErrorV3 } from "./handleRedisTrackErrorV3.js"; -const buildAiCreditCostProperty = ({ - featureDeductions, - deductions, -}: { - featureDeductions: FeatureDeduction[]; - deductions: TrackDeduction[]; -}): Record | undefined => { - const aiDeduction = featureDeductions.find((d) => d.tokenUsage); - if (!aiDeduction) return; - - const creditCost: Record = {}; - for (const deduction of deductions) { - if (deduction.feature_id === aiDeduction.feature.id) continue; - if (!deduction.value) continue; - creditCost[deduction.feature_id] = - (creditCost[deduction.feature_id] ?? 0) + deduction.value; - } - - return Object.keys(creditCost).length > 0 ? creditCost : undefined; -}; - const queueSyncItem = ({ ctx, body, @@ -171,7 +151,10 @@ export const runRedisTrackV3 = async ({ const aiCreditCost = buildAiCreditCostProperty({ featureDeductions, - deductions, + entries: deductions.map((d) => ({ + featureId: d.feature_id, + amount: d.value ?? 0, + })), }); if (aiCreditCost) { body.properties = { ...(body.properties ?? {}), credit_cost: aiCreditCost }; From bd41bde109ae12d57f8ff43d37a3b9434b4e4370 Mon Sep 17 00:00:00 2001 From: Ridhwan Hussain <73362400+TheUntraceable@users.noreply.github.com.> Date: Mon, 1 Jun 2026 14:46:06 +0100 Subject: [PATCH 18/46] fix(balances): update operationId and spec identifiers for track_tokens route --- packages/openapi/v2.3/contracts/balancesContract.ts | 4 ++-- packages/openapi/v2.3/jsDocs/balancesJsDocs.ts | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/packages/openapi/v2.3/contracts/balancesContract.ts b/packages/openapi/v2.3/contracts/balancesContract.ts index a18a81560..558b03dd1 100644 --- a/packages/openapi/v2.3/contracts/balancesContract.ts +++ b/packages/openapi/v2.3/contracts/balancesContract.ts @@ -163,12 +163,12 @@ export const balancesTrackTokensContract = oc .route({ method: "POST", path: "/v1/balances.track_tokens", - operationId: "trackTokens", + operationId: "track_tokens", description: balancesTrackTokensJsDoc, spec: (spec) => withAcceptedResponse( spec, - "trackTokens", + "track_tokens", "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.", ), }) diff --git a/packages/openapi/v2.3/jsDocs/balancesJsDocs.ts b/packages/openapi/v2.3/jsDocs/balancesJsDocs.ts index 12ec2c81a..559d7db3c 100644 --- a/packages/openapi/v2.3/jsDocs/balancesJsDocs.ts +++ b/packages/openapi/v2.3/jsDocs/balancesJsDocs.ts @@ -81,7 +81,7 @@ export const balancesTrackTokensJsDoc = createJSDocDescription({ }, }), ], - methodName: "trackTokens", + methodName: "track_tokens", 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.", }); From 8c3d0107aa2bcfec06f36953749270d82669843f Mon Sep 17 00:00:00 2001 From: Ridhwan Hussain <73362400+TheUntraceable@users.noreply.github.com.> Date: Mon, 1 Jun 2026 14:47:13 +0100 Subject: [PATCH 19/46] chore(tests): don't use non-null assertion --- .../track/basic/track-deductions.test.ts | 6 ++-- .../balances/track/basic/track-tokens.test.ts | 32 +++++++++++++------ 2 files changed, 26 insertions(+), 12 deletions(-) diff --git a/server/tests/integration/balances/track/basic/track-deductions.test.ts b/server/tests/integration/balances/track/basic/track-deductions.test.ts index 47dc727f2..7a4b6e6c8 100644 --- a/server/tests/integration/balances/track/basic/track-deductions.test.ts +++ b/server/tests/integration/balances/track/basic/track-deductions.test.ts @@ -288,7 +288,9 @@ test.concurrent( const creditFeature = ctx.features.find( (f) => f.id === TestFeature.Credits, ); - expect(creditFeature).toBeDefined(); + if (!creditFeature) { + throw new Error(`${TestFeature.Credits} feature not found`); + } const customerBefore = await autumnV1.customers.get(customerId); @@ -306,7 +308,7 @@ test.concurrent( const overflowAmount = 50; const expectedCreditCost = await getCreditCost({ featureId: TestFeature.Action1, - creditSystem: creditFeature!, + creditSystem: creditFeature, amount: overflowAmount, }); diff --git a/server/tests/integration/balances/track/basic/track-tokens.test.ts b/server/tests/integration/balances/track/basic/track-tokens.test.ts index b9d15dedb..a97065fb7 100644 --- a/server/tests/integration/balances/track/basic/track-tokens.test.ts +++ b/server/tests/integration/balances/track/basic/track-tokens.test.ts @@ -32,6 +32,9 @@ test.concurrent(`${chalk.yellowBright("track-tokens-1: basic trackTokens with mo const aiCreditFeature = ctx.features.find( (f) => f.id === TestFeature.AiCredits, ); + if (!aiCreditFeature) { + throw new Error(`${TestFeature.AiCredits} feature not found`); + } const customerBefore = await autumnV1.customers.get(customerId); @@ -42,8 +45,8 @@ test.concurrent(`${chalk.yellowBright("track-tokens-1: basic trackTokens with mo const modelId = "anthropic/claude-sonnet-4-20250514"; const expectedCost = await getCreditCost({ - featureId: aiCreditFeature!.id, - creditSystem: aiCreditFeature!, + featureId: aiCreditFeature.id, + creditSystem: aiCreditFeature, modelName: modelId, tokens: { input: inputTokens, output: outputTokens }, }); @@ -109,14 +112,17 @@ test.concurrent(`${chalk.yellowBright("track-tokens-2: disambiguation error and const aiCreditFeature = ctx.features.find( (f) => f.id === TestFeature.AiCredits, ); + if (!aiCreditFeature) { + throw new Error(`${TestFeature.AiCredits} feature not found`); + } const inputTokens = 2000; const outputTokens = 1000; const modelId = "anthropic/claude-sonnet-4-20250514"; const expectedCost = await getCreditCost({ - featureId: aiCreditFeature!.id, - creditSystem: aiCreditFeature!, + featureId: aiCreditFeature.id, + creditSystem: aiCreditFeature, modelName: modelId, tokens: { input: inputTokens, output: outputTokens }, }); @@ -237,6 +243,9 @@ test.concurrent(`${chalk.yellowBright("track-tokens-4: models.dev markup and non const aiCreditFeature = ctx.features.find( (f) => f.id === TestFeature.AiCredits, ); + if (!aiCreditFeature) { + throw new Error(`${TestFeature.AiCredits} feature not found`); + } // anthropic/claude-haiku-3.5 has 20% markup in test config const inputTokens = 50000; @@ -244,8 +253,8 @@ test.concurrent(`${chalk.yellowBright("track-tokens-4: models.dev markup and non const modelId = "anthropic/claude-3-5-haiku-20241022"; const expectedCost = await getCreditCost({ - featureId: aiCreditFeature!.id, - creditSystem: aiCreditFeature!, + featureId: aiCreditFeature.id, + creditSystem: aiCreditFeature, modelName: modelId, tokens: { input: inputTokens, output: outputTokens }, }); @@ -306,11 +315,14 @@ test.concurrent(`${chalk.yellowBright("track-tokens-5: multiple tracks accumulat const aiCreditFeature = ctx.features.find( (f) => f.id === TestFeature.AiCredits, ); + if (!aiCreditFeature) { + throw new Error(`${TestFeature.AiCredits} feature not found`); + } // First track: custom/internal-model (input_cost=5, output_cost=15, markup=0%) const cost1 = await getCreditCost({ - featureId: aiCreditFeature!.id, - creditSystem: aiCreditFeature!, + featureId: aiCreditFeature.id, + creditSystem: aiCreditFeature, modelName: "custom/internal-model", tokens: { input: 5000, output: 2000 }, }); @@ -325,8 +337,8 @@ test.concurrent(`${chalk.yellowBright("track-tokens-5: multiple tracks accumulat // Second track: custom/marked-up-model (input_cost=10, output_cost=30, markup=50%) const cost2 = await getCreditCost({ - featureId: aiCreditFeature!.id, - creditSystem: aiCreditFeature!, + featureId: aiCreditFeature.id, + creditSystem: aiCreditFeature, modelName: "custom/marked-up-model", tokens: { input: 3000, output: 1000 }, }); From c3dd74c3eb7eaca2cd1e818df6253d93ba3b66c8 Mon Sep 17 00:00:00 2001 From: Ridhwan Hussain <73362400+TheUntraceable@users.noreply.github.com.> Date: Mon, 1 Jun 2026 14:48:16 +0100 Subject: [PATCH 20/46] chore: dedupe buildUnlimitedPlanMutationLog --- .../buildUnlimitedPlanMutationLog.ts | 37 +++++++++++++++++++ .../deductionV2/executePostgresDeductionV2.ts | 27 +++++--------- .../deductionV2/executeRedisDeductionV2.ts | 27 +++++--------- 3 files changed, 55 insertions(+), 36 deletions(-) create mode 100644 server/src/internal/balances/utils/deductionV2/buildUnlimitedPlanMutationLog.ts diff --git a/server/src/internal/balances/utils/deductionV2/buildUnlimitedPlanMutationLog.ts b/server/src/internal/balances/utils/deductionV2/buildUnlimitedPlanMutationLog.ts new file mode 100644 index 000000000..36c3cb53b --- /dev/null +++ b/server/src/internal/balances/utils/deductionV2/buildUnlimitedPlanMutationLog.ts @@ -0,0 +1,37 @@ +import type { FullCusEntWithFullCusProduct } from "@autumn/shared"; +import type { MutationLogItem } from "../types/mutationLogItem.js"; + +/** + * Attribute a track event to an unlimited plan even though we skip the actual + * deduction. Without this, resolveInternalProductIdForEvent gets an empty + * mutation log and the event lands in "No plan". Returns null when there is no + * unlimited entitlement to attribute to, or the resolved delta is zero. + */ +export const buildUnlimitedPlanMutationLog = ({ + unlimitedCusEnt, + toDeduct, + fallbackDeduction, + entityId, +}: { + unlimitedCusEnt: FullCusEntWithFullCusProduct | undefined; + toDeduct: number | null | undefined; + fallbackDeduction: number | null | undefined; + entityId?: string | null; +}): MutationLogItem | null => { + if (!unlimitedCusEnt) return null; + + const syntheticDelta = -(toDeduct ?? fallbackDeduction ?? 1); + if (syntheticDelta === 0) return null; + + return { + target_type: "customer_entitlement", + customer_entitlement_id: unlimitedCusEnt.id, + rollover_id: null, + entity_id: entityId ?? null, + credit_cost: 1, + balance_delta: syntheticDelta, + adjustment_delta: 0, + usage_delta: 0, + value_delta: 0, + }; +}; diff --git a/server/src/internal/balances/utils/deductionV2/executePostgresDeductionV2.ts b/server/src/internal/balances/utils/deductionV2/executePostgresDeductionV2.ts index 1b6d5ca84..894076631 100644 --- a/server/src/internal/balances/utils/deductionV2/executePostgresDeductionV2.ts +++ b/server/src/internal/balances/utils/deductionV2/executePostgresDeductionV2.ts @@ -17,6 +17,7 @@ import type { DeductionUpdate } from "../types/deductionUpdate.js"; import type { FeatureDeduction } from "../types/featureDeduction.js"; import type { MutationLogItem } from "../types/mutationLogItem.js"; import { applyDeductionUpdateToFullSubject } from "./applyDeductionUpdateToFullSubject.js"; +import { buildUnlimitedPlanMutationLog } from "./buildUnlimitedPlanMutationLog.js"; import { applyRolloverUpdatesToFullSubject } from "./applyRolloverUpdatesToFullSubject.js"; import { logDeductionUpdatesV2 } from "./logDeductionUpdatesV2.js"; import { mutationLogsToFeaturesV2 } from "./mutationLogsToFeaturesV2.js"; @@ -130,24 +131,14 @@ export const executePostgresDeductionV2 = async ({ redisInstance: ctx.redisV2, }); } - // Attribute the event to the unlimited plan even though we skip - // the actual deduction. Without this, resolveInternalProductIdForEvent - // gets an empty mutation log and the event lands in "No plan". - if (unlimitedCusEnt) { - const syntheticDelta = -(toDeduct ?? deduction.deduction ?? 1); - if (syntheticDelta !== 0) { - allMutationLogs.push({ - target_type: "customer_entitlement", - customer_entitlement_id: unlimitedCusEnt.id, - rollover_id: null, - entity_id: entityId ?? null, - credit_cost: 1, - balance_delta: syntheticDelta, - adjustment_delta: 0, - usage_delta: 0, - value_delta: 0, - }); - } + const unlimitedPlanLog = buildUnlimitedPlanMutationLog({ + unlimitedCusEnt, + toDeduct, + fallbackDeduction: deduction.deduction, + entityId, + }); + if (unlimitedPlanLog) { + allMutationLogs.push(unlimitedPlanLog); } continue; } diff --git a/server/src/internal/balances/utils/deductionV2/executeRedisDeductionV2.ts b/server/src/internal/balances/utils/deductionV2/executeRedisDeductionV2.ts index bfd9cdcb2..d40a508dc 100644 --- a/server/src/internal/balances/utils/deductionV2/executeRedisDeductionV2.ts +++ b/server/src/internal/balances/utils/deductionV2/executeRedisDeductionV2.ts @@ -29,6 +29,7 @@ import type { LuaDeductionResult } from "../types/redisDeductionResult.js"; import type { RolloverUpdate } from "../types/rolloverUpdate.js"; import { applyDeductionUpdateToFullSubject } from "./applyDeductionUpdateToFullSubject.js"; import { applyRolloverUpdatesToFullSubject } from "./applyRolloverUpdatesToFullSubject.js"; +import { buildUnlimitedPlanMutationLog } from "./buildUnlimitedPlanMutationLog.js"; import { logDeductionUpdatesV2 } from "./logDeductionUpdatesV2.js"; import { mutationLogsToFeaturesV2 } from "./mutationLogsToFeaturesV2.js"; import { normalizeDeductionSyncStateV2 } from "./normalizeDeductionSyncStateV2.js"; @@ -141,24 +142,14 @@ export const executeRedisDeductionV2 = async ({ redisInstance: redisInstance ?? ctx.redisV2, }); } - // Attribute the event to the unlimited plan even though we skip - // the actual deduction. Without this, resolveInternalProductIdForEvent - // gets an empty mutation log and the event lands in "No plan". - if (unlimitedCusEnt) { - const syntheticDelta = -(toDeduct ?? deduction.deduction ?? 1); - if (syntheticDelta !== 0) { - allMutationLogs.push({ - target_type: "customer_entitlement", - customer_entitlement_id: unlimitedCusEnt.id, - rollover_id: null, - entity_id: entityId ?? null, - credit_cost: 1, - balance_delta: syntheticDelta, - adjustment_delta: 0, - usage_delta: 0, - value_delta: 0, - }); - } + const unlimitedPlanLog = buildUnlimitedPlanMutationLog({ + unlimitedCusEnt, + toDeduct, + fallbackDeduction: deduction.deduction, + entityId, + }); + if (unlimitedPlanLog) { + allMutationLogs.push(unlimitedPlanLog); } continue; } From 0d85809ad72e9d7c90b0fb14f4e3fc1c7a4e72b6 Mon Sep 17 00:00:00 2001 From: Ridhwan Hussain <73362400+TheUntraceable@users.noreply.github.com.> Date: Mon, 1 Jun 2026 15:02:24 +0100 Subject: [PATCH 21/46] revert(balances): keep trackTokens operationId camelCase The track_tokens route path is already snake_case (correct). operationId/SDK method names are camelCase across all ~40 routes (batchTrack, getOrCreateCustomer, redeemReferralCode...); Speakeasy emits the override verbatim, so snake_case here produced a lone track_tokens SDK method. Reverts bd41bde10. --- packages/openapi/v2.3/contracts/balancesContract.ts | 4 ++-- packages/openapi/v2.3/jsDocs/balancesJsDocs.ts | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/packages/openapi/v2.3/contracts/balancesContract.ts b/packages/openapi/v2.3/contracts/balancesContract.ts index 558b03dd1..a18a81560 100644 --- a/packages/openapi/v2.3/contracts/balancesContract.ts +++ b/packages/openapi/v2.3/contracts/balancesContract.ts @@ -163,12 +163,12 @@ export const balancesTrackTokensContract = oc .route({ method: "POST", path: "/v1/balances.track_tokens", - operationId: "track_tokens", + operationId: "trackTokens", description: balancesTrackTokensJsDoc, spec: (spec) => withAcceptedResponse( spec, - "track_tokens", + "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.", ), }) diff --git a/packages/openapi/v2.3/jsDocs/balancesJsDocs.ts b/packages/openapi/v2.3/jsDocs/balancesJsDocs.ts index 559d7db3c..12ec2c81a 100644 --- a/packages/openapi/v2.3/jsDocs/balancesJsDocs.ts +++ b/packages/openapi/v2.3/jsDocs/balancesJsDocs.ts @@ -81,7 +81,7 @@ export const balancesTrackTokensJsDoc = createJSDocDescription({ }, }), ], - methodName: "track_tokens", + 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.", }); From d4cea174c6ce83e4c92f91c485efb6a738a485a7 Mon Sep 17 00:00:00 2001 From: Ridhwan Hussain <73362400+TheUntraceable@users.noreply.github.com.> Date: Mon, 1 Jun 2026 15:28:02 +0100 Subject: [PATCH 22/46] feat: enhance AI credit feature resolution and update trackTokensParams description --- .../track/utils/getTokenTrackParams.ts | 119 +++++++++++++----- .../api/balances/track/trackTokensParams.ts | 2 +- 2 files changed, 87 insertions(+), 34 deletions(-) diff --git a/server/src/internal/balances/track/utils/getTokenTrackParams.ts b/server/src/internal/balances/track/utils/getTokenTrackParams.ts index 4e9aef820..11fbf0c90 100644 --- a/server/src/internal/balances/track/utils/getTokenTrackParams.ts +++ b/server/src/internal/balances/track/utils/getTokenTrackParams.ts @@ -6,53 +6,100 @@ import { type TrackParams, type TrackTokensParams, } from "@autumn/shared"; +import { fullCustomerToCustomerEntitlements } from "@autumn/shared"; +import { fullSubjectToFullCustomer } from "@autumn/shared"; import type { AutumnContext } from "@/honoUtils/HonoEnv.js"; +import { getOrSetCachedFullSubject } from "@/internal/customers/cache/fullSubject/actions/getOrSetCachedFullSubject.js"; +import { getOrSetCachedFullCustomer } from "@/internal/customers/cusUtils/fullCustomerCacheUtils/getOrSetCachedFullCustomer.js"; import { getCreditCost } from "@/internal/features/creditSystemUtils.js"; +import { isFullSubjectRolloutEnabled } from "@/internal/misc/rollouts/fullSubjectRolloutUtils.js"; import type { FeatureDeduction } from "../../utils/types/featureDeduction.js"; -const resolveAiCreditFeature = ({ +const resolveAiCreditFeatureById = ({ features, featureId, }: { features: Feature[]; - featureId?: string; + featureId: string; }): Feature => { - if (featureId) { - const candidate = features.find((f) => f.id === featureId); - if (!candidate) { - throw new RecaseError({ - message: `Feature ${featureId} not found`, - code: ErrCode.FeatureNotFound, - statusCode: 404, - }); - } - if (candidate.type !== FeatureType.AiCreditSystem) { - throw new RecaseError({ - message: `Feature ${featureId} is not an AI credit system feature`, - code: ErrCode.InvalidRequest, - statusCode: 400, - }); - } - return candidate; - } - - const matches = features.filter((f) => f.type === FeatureType.AiCreditSystem); - if (matches.length === 0) { + const candidate = features.find((f) => f.id === featureId); + if (!candidate) { throw new RecaseError({ - message: "No AI credit system feature found for this organization", + message: `Feature ${featureId} not found`, code: ErrCode.FeatureNotFound, statusCode: 404, }); } - if (matches.length > 1) { + if (candidate.type !== FeatureType.AiCreditSystem) { throw new RecaseError({ - message: - "Multiple AI credit system features found. Please specify a feature_id to disambiguate.", + message: `Feature ${featureId} is not an AI credit system feature`, code: ErrCode.InvalidRequest, statusCode: 400, }); } - return matches[0]; + return candidate; +}; + +const resolveAiCreditFeatureFromEntitlements = async ({ + ctx, + customerId, + entityId, +}: { + ctx: AutumnContext; + customerId: string; + entityId?: string; +}): Promise => { + const fullCustomer = isFullSubjectRolloutEnabled({ ctx }) + ? fullSubjectToFullCustomer({ + fullSubject: await getOrSetCachedFullSubject({ + ctx, + customerId, + entityId, + source: "resolveAiCreditFeature", + }), + }) + : await getOrSetCachedFullCustomer({ + ctx, + customerId, + entityId, + source: "resolveAiCreditFeature", + }); + + const entity = entityId + ? fullCustomer.entities?.find((e) => e.id === entityId) + : undefined; + + const cusEnts = fullCustomerToCustomerEntitlements({ + fullCustomer, + entity, + }); + + const aiCreditFeatures = [ + ...new Map( + cusEnts + .filter( + (ce) => ce.entitlement.feature.type === FeatureType.AiCreditSystem, + ) + .map((ce) => [ce.entitlement.feature.id, ce.entitlement.feature]), + ).values(), + ]; + + if (aiCreditFeatures.length === 0) { + throw new RecaseError({ + message: "No AI credit system feature found for this customer", + code: ErrCode.FeatureNotFound, + statusCode: 404, + }); + } + if (aiCreditFeatures.length > 1) { + throw new RecaseError({ + message: + "Multiple AI credit system features found for this customer. Please specify a feature_id to disambiguate.", + code: ErrCode.InvalidRequest, + statusCode: 400, + }); + } + return aiCreditFeatures[0]; }; export const getTokenTrackParams = async ({ @@ -62,10 +109,16 @@ export const getTokenTrackParams = async ({ ctx: AutumnContext; input: TrackTokensParams; }): Promise<{ body: TrackParams; featureDeductions: FeatureDeduction[] }> => { - const aiCreditFeature = resolveAiCreditFeature({ - features: ctx.features, - featureId: input.feature_id, - }); + const aiCreditFeature = input.feature_id + ? resolveAiCreditFeatureById({ + features: ctx.features, + featureId: input.feature_id, + }) + : await resolveAiCreditFeatureFromEntitlements({ + ctx, + customerId: input.customer_id, + entityId: input.entity_id, + }); const cost = await getCreditCost({ featureId: aiCreditFeature.id, @@ -110,4 +163,4 @@ export const getTokenTrackParams = async ({ }; return { body, featureDeductions }; -}; +}; \ No newline at end of file diff --git a/shared/api/balances/track/trackTokensParams.ts b/shared/api/balances/track/trackTokensParams.ts index 3d79830a5..818ed980a 100644 --- a/shared/api/balances/track/trackTokensParams.ts +++ b/shared/api/balances/track/trackTokensParams.ts @@ -11,7 +11,7 @@ export const TrackTokensParamsSchema = z.object({ }), feature_id: z.string().optional().meta({ description: - "The ID of the AI credit system feature. Auto-detected if omitted.", + "The ID of the AI credit system feature. Auto-detected from the customer's entitlements if omitted — only required when a customer has multiple AI credit systems.", }), model_id: z.string().meta({ description: "The AI model name with provider prefix (e.g., 'anthropic/claude-opus-4-6').", From 9e4909a9fa9ccd06dbd64598975b30bbf3d1397b Mon Sep 17 00:00:00 2001 From: Charlie Lamb Date: Mon, 1 Jun 2026 17:28:40 +0100 Subject: [PATCH 23/46] chore: dry up credit system checks --- server/src/internal/features/creditSystemUtils.ts | 3 ++- .../internal/features/featureActions/createFeature.ts | 6 ++---- .../internal/features/featureActions/updateFeature.ts | 5 ++--- shared/utils/featureUtils/apiFeatureToDbFeature.ts | 4 ++-- .../featureUtils/classifyFeature/isAnyCreditSystem.ts | 4 ++++ shared/utils/featureUtils/index.ts | 4 ++++ .../features/components/CreateFeatureSheet.tsx | 6 ++---- .../features/feature-list/FeatureListTable.tsx | 10 +++------- .../components/new-feature/NewFeatureBehaviour.tsx | 6 ++---- 9 files changed, 23 insertions(+), 25 deletions(-) create mode 100644 shared/utils/featureUtils/classifyFeature/isAnyCreditSystem.ts diff --git a/server/src/internal/features/creditSystemUtils.ts b/server/src/internal/features/creditSystemUtils.ts index 3b47a3009..08b6340d6 100644 --- a/server/src/internal/features/creditSystemUtils.ts +++ b/server/src/internal/features/creditSystemUtils.ts @@ -4,6 +4,7 @@ import { type Feature, FeatureType, InternalError, + isAnyCreditSystem, RecaseError, } from "@autumn/shared"; import { Decimal } from "decimal.js"; @@ -170,7 +171,7 @@ export const getCreditCost = async ({ modelName?: string; tokens?: TokenInput; }) => { - if (creditSystem.type !== FeatureType.CreditSystem && creditSystem.type !== FeatureType.AiCreditSystem) { + if (!isAnyCreditSystem(creditSystem.type)) { return amount; } if (creditSystem.type === FeatureType.AiCreditSystem) { diff --git a/server/src/internal/features/featureActions/createFeature.ts b/server/src/internal/features/featureActions/createFeature.ts index 5e073b6a5..6fd71732e 100644 --- a/server/src/internal/features/featureActions/createFeature.ts +++ b/server/src/internal/features/featureActions/createFeature.ts @@ -2,6 +2,7 @@ import { CreateFeatureSchema, type Feature, FeatureType, + isAnyCreditSystem, type ModelMarkups, } from "@autumn/shared"; import type { AutumnContext } from "@/honoUtils/HonoEnv.js"; @@ -20,10 +21,7 @@ const validateFeature = (data: any, allFeatures: Feature[]) => { let config = data.config; if (featureType === FeatureType.Metered) { config = validateMeteredConfig(config); - } else if ( - featureType === FeatureType.CreditSystem || - featureType === FeatureType.AiCreditSystem - ) { + } else if (isAnyCreditSystem(featureType)) { config = validateCreditSystem(config, featureType); if (featureType === FeatureType.CreditSystem) { validateCreditSystemSchemaReferences({ diff --git a/server/src/internal/features/featureActions/updateFeature.ts b/server/src/internal/features/featureActions/updateFeature.ts index c597bbf4b..38d4b9472 100644 --- a/server/src/internal/features/featureActions/updateFeature.ts +++ b/server/src/internal/features/featureActions/updateFeature.ts @@ -3,6 +3,7 @@ import { ErrCode, type Feature, FeatureType, + isAnyCreditSystem, type ModelMarkups, notNullish, } from "@autumn/shared"; @@ -230,9 +231,7 @@ export const updateFeature = async ({ } // Queue cache clear for credit system if schema or model markups changed - const isCreditSystem = - feature.type === FeatureType.CreditSystem || - feature.type === FeatureType.AiCreditSystem; + const isCreditSystem = isAnyCreditSystem(feature.type); if (isCreditSystem && updatedFeature) { const schemaChanged = updates.config != null && diff --git a/shared/utils/featureUtils/apiFeatureToDbFeature.ts b/shared/utils/featureUtils/apiFeatureToDbFeature.ts index 4c1d0cb6e..c621bf054 100644 --- a/shared/utils/featureUtils/apiFeatureToDbFeature.ts +++ b/shared/utils/featureUtils/apiFeatureToDbFeature.ts @@ -8,6 +8,7 @@ import { FeatureUsageType, } from "@models/featureModels/featureEnums.js"; import type { Feature } from "@models/featureModels/featureModels.js"; +import { isAnyCreditSystem } from "./classifyFeature/isAnyCreditSystem.js"; import { AppEnv } from "@models/genModels/genEnums.js"; import type { ApiFeatureV1 } from "../../api/features/apiFeatureV1.js"; import type { @@ -208,8 +209,7 @@ export const dbToApiFeatureV1 = ({ name: dbFeature.name, type: dbFeature.type, consumable: - dbFeature.type === FeatureType.CreditSystem || - dbFeature.type === FeatureType.AiCreditSystem || + isAnyCreditSystem(dbFeature.type) || dbFeature.config?.usage_type === FeatureUsageType.Single, credit_schema: Array.isArray(dbFeature.config?.schema) diff --git a/shared/utils/featureUtils/classifyFeature/isAnyCreditSystem.ts b/shared/utils/featureUtils/classifyFeature/isAnyCreditSystem.ts new file mode 100644 index 000000000..6f4bac496 --- /dev/null +++ b/shared/utils/featureUtils/classifyFeature/isAnyCreditSystem.ts @@ -0,0 +1,4 @@ +import { FeatureType } from "@models/featureModels/featureEnums"; + +export const isAnyCreditSystem = (type: FeatureType): boolean => + type === FeatureType.CreditSystem || type === FeatureType.AiCreditSystem; diff --git a/shared/utils/featureUtils/index.ts b/shared/utils/featureUtils/index.ts index 4d47d8777..184129007 100644 --- a/shared/utils/featureUtils/index.ts +++ b/shared/utils/featureUtils/index.ts @@ -1,4 +1,5 @@ import { isAllocatedFeature } from "@utils/featureUtils/classifyFeature/isAllocatedFeature"; +import { isAnyCreditSystem } from "@utils/featureUtils/classifyFeature/isAnyCreditSystem"; import { isConsumableFeature } from "@utils/featureUtils/classifyFeature/isConsumableFeature"; import { findFeatureById } from "@utils/featureUtils/findFeatureUtils"; @@ -8,9 +9,12 @@ export * from "./convertFeatureUtils"; export * from "./creditSystemUtils"; export * from "./findFeatureUtils"; +export { isAnyCreditSystem } from "@utils/featureUtils/classifyFeature/isAnyCreditSystem"; + export const featureUtils = { isConsumable: isConsumableFeature, isAllocated: isAllocatedFeature, + isAnyCreditSystem, find: { byId: findFeatureById, diff --git a/vite/src/views/products/features/components/CreateFeatureSheet.tsx b/vite/src/views/products/features/components/CreateFeatureSheet.tsx index ab61c4be4..e5003ff12 100644 --- a/vite/src/views/products/features/components/CreateFeatureSheet.tsx +++ b/vite/src/views/products/features/components/CreateFeatureSheet.tsx @@ -3,6 +3,7 @@ import { type CreditSchemaItem, FeatureType, FeatureUsageType, + isAnyCreditSystem, } from "@autumn/shared"; import type { AxiosError } from "axios"; import { useEffect, useState } from "react"; @@ -52,10 +53,7 @@ function CreateFeatureSheet({ const handleCreateFeature = async () => { // Validate credit system specific fields first - if ( - feature.type === FeatureType.CreditSystem || - feature.type === FeatureType.AiCreditSystem - ) { + if (isAnyCreditSystem(feature.type)) { const validationError = validateCreditSystem(feature); if (validationError) { toast.error(validationError); diff --git a/vite/src/views/products/features/feature-list/FeatureListTable.tsx b/vite/src/views/products/features/feature-list/FeatureListTable.tsx index a34d65026..8d65826a6 100644 --- a/vite/src/views/products/features/feature-list/FeatureListTable.tsx +++ b/vite/src/views/products/features/feature-list/FeatureListTable.tsx @@ -1,4 +1,4 @@ -import { AppEnv, type Feature, FeatureType } from "@autumn/shared"; +import { AppEnv, type Feature, isAnyCreditSystem } from "@autumn/shared"; import { ArrowSquareOutIcon, CoinsIcon, LegoIcon } from "@phosphor-icons/react"; import { useMemo, useState } from "react"; import { Table } from "@/components/general/table"; @@ -29,19 +29,15 @@ export function FeatureListTable() { // Filter features and credit systems based on archived state const { regularFeatures, creditSystems, hasEventNames } = useMemo(() => { - const isCreditType = (type: string) => - type === FeatureType.CreditSystem || - type === FeatureType.AiCreditSystem; - const regularFeatures = features?.filter((feature) => { - if (isCreditType(feature.type)) return false; + if (isAnyCreditSystem(feature.type)) return false; return queryStates.showArchivedFeatures ? feature.archived : !feature.archived; }); const creditSystems = features?.filter((feature) => { - if (!isCreditType(feature.type)) return false; + if (!isAnyCreditSystem(feature.type)) return false; return queryStates.showArchivedFeatures ? feature.archived : !feature.archived; diff --git a/vite/src/views/products/plan/components/new-feature/NewFeatureBehaviour.tsx b/vite/src/views/products/plan/components/new-feature/NewFeatureBehaviour.tsx index 928f19176..1dcf82a4e 100644 --- a/vite/src/views/products/plan/components/new-feature/NewFeatureBehaviour.tsx +++ b/vite/src/views/products/plan/components/new-feature/NewFeatureBehaviour.tsx @@ -2,6 +2,7 @@ import { type CreateFeature, FeatureType, FeatureUsageType, + isAnyCreditSystem, } from "@autumn/shared"; import { AreaRadioGroupItem } from "@/components/v2/radio-groups/AreaRadioGroupItem"; import { RadioGroup } from "@/components/v2/radio-groups/RadioGroup"; @@ -62,10 +63,7 @@ export function NewFeatureBehaviour({ feature: CreateFeature; setFeature: (feature: CreateFeature) => void; }) { - if ( - feature.type === FeatureType.CreditSystem || - feature.type === FeatureType.AiCreditSystem - ) { + if (isAnyCreditSystem(feature.type)) { return ; } From 7abe65753067896b3d2e714dce1794f7a44f904d Mon Sep 17 00:00:00 2001 From: Charlie Lamb Date: Mon, 1 Jun 2026 17:43:32 +0100 Subject: [PATCH 24/46] chore: separate ai credit system utils --- .../internal/features/aiCreditSystemUtils.ts | 91 +++++++++++++++++++ .../internal/features/creditSystemUtils.ts | 88 +----------------- 2 files changed, 95 insertions(+), 84 deletions(-) create mode 100644 server/src/internal/features/aiCreditSystemUtils.ts diff --git a/server/src/internal/features/aiCreditSystemUtils.ts b/server/src/internal/features/aiCreditSystemUtils.ts new file mode 100644 index 000000000..234d82643 --- /dev/null +++ b/server/src/internal/features/aiCreditSystemUtils.ts @@ -0,0 +1,91 @@ +import { + ErrCode, + type Feature, + InternalError, + RecaseError, +} from "@autumn/shared"; +import { Decimal } from "decimal.js"; +import { getModelsDevPricing } from "@/internal/features/utils/getModelPricing.js"; + +export type TokenInput = { input: number; output: number }; + +// Costs are in $/M tokens; markup is a percentage (e.g. 20 = +20%). +const computeMarkedUpCost = ({ + inputCostPerMillion, + outputCostPerMillion, + input, + output, + markup, +}: { + inputCostPerMillion: Decimal.Value; + outputCostPerMillion: Decimal.Value; + input: number; + output: number; + markup: number; +}) => + new Decimal(inputCostPerMillion) + .mul(input) + .add(new Decimal(outputCostPerMillion).mul(output)) + .div(1_000_000) + .mul(new Decimal(1).add(new Decimal(markup).div(100))) + .toNumber(); + +export const getModelCreditCost = async ({ + modelName, + creditSystem, + input, + output, +}: { + modelName: string; + creditSystem: Feature; +} & TokenInput) => { + const markups = creditSystem.model_markups || {}; + const markupEntry = markups[modelName]; + const { markup } = markupEntry ?? { markup: 0 }; + + if (modelName.startsWith("custom/")) { + if (markupEntry?.input_cost == null || markupEntry?.output_cost == null) { + throw new RecaseError({ + message: `Custom model ${modelName} is missing input_cost or output_cost in model_markups`, + code: ErrCode.InvalidRequest, + data: { modelName }, + }); + } + return computeMarkedUpCost({ + inputCostPerMillion: markupEntry.input_cost, + outputCostPerMillion: markupEntry.output_cost, + input, + output, + markup, + }); + } + + const pricingData = await getModelsDevPricing(); + if (!pricingData) { + throw new InternalError({ + message: "Failed to fetch models.dev pricing data", + code: ErrCode.InternalError, + }); + } + + const [providerKey, ...modelParts] = modelName.split("/"); + const modelKey = modelParts.join("/"); + const model = pricingData[providerKey]?.models[modelKey]; + + if (!model) { + throw new RecaseError({ + message: `Model ${modelName} not found in models.dev pricing data ${providerKey} provider config.`, + code: ErrCode.InvalidRequest, + statusCode: 400, + data: { modelName }, + }); + } + + return computeMarkedUpCost({ + inputCostPerMillion: model.cost.input, + outputCostPerMillion: model.cost.output, + input, + output, + markup, + }); +}; diff --git a/server/src/internal/features/creditSystemUtils.ts b/server/src/internal/features/creditSystemUtils.ts index 08b6340d6..ace6f4504 100644 --- a/server/src/internal/features/creditSystemUtils.ts +++ b/server/src/internal/features/creditSystemUtils.ts @@ -3,14 +3,14 @@ import { ErrCode, type Feature, FeatureType, - InternalError, isAnyCreditSystem, RecaseError, } from "@autumn/shared"; import { Decimal } from "decimal.js"; -import { getModelsDevPricing } from "@/internal/features/utils/getModelPricing.js"; - -type TokenInput = { input: number; output: number }; +import { + getModelCreditCost, + type TokenInput, +} from "@/internal/features/aiCreditSystemUtils.js"; const creditSystemContainsFeature = ({ creditSystem, @@ -77,86 +77,6 @@ export const featureToCreditSystem = ({ return amount; }; -// Costs are in $/M tokens; markup is a percentage (e.g. 20 = +20%). -const computeMarkedUpCost = ({ - inputCostPerMillion, - outputCostPerMillion, - input, - output, - markup, -}: { - inputCostPerMillion: Decimal.Value; - outputCostPerMillion: Decimal.Value; - input: number; - output: number; - markup: number; -}) => - new Decimal(inputCostPerMillion) - .mul(input) - .add(new Decimal(outputCostPerMillion).mul(output)) - .div(1_000_000) - .mul(new Decimal(1).add(new Decimal(markup).div(100))) - .toNumber(); - -const getModelCreditCost = async ({ - modelName, - creditSystem, - input, - output, -}: { - modelName: string; - creditSystem: Feature; -} & TokenInput) => { - const markups = creditSystem.model_markups || {}; - const markupEntry = markups[modelName]; - const { markup } = markupEntry ?? { markup: 0 }; - - if (modelName.startsWith("custom/")) { - if (markupEntry?.input_cost == null || markupEntry?.output_cost == null) { - throw new RecaseError({ - message: `Custom model ${modelName} is missing input_cost or output_cost in model_markups`, - code: ErrCode.InvalidRequest, - data: { modelName }, - }); - } - return computeMarkedUpCost({ - inputCostPerMillion: markupEntry.input_cost, - outputCostPerMillion: markupEntry.output_cost, - input, - output, - markup, - }); - } - - const pricingData = await getModelsDevPricing(); - if (!pricingData) { - throw new InternalError({ - message: "Failed to fetch models.dev pricing data", - code: ErrCode.InternalError, - }); - } - - const [providerKey, ...modelParts] = modelName.split("/"); - const modelKey = modelParts.join("/"); - const model = pricingData[providerKey]?.models[modelKey]; - - if (!model) { - throw new RecaseError({ - message: `Model ${modelName} not found in models.dev pricing data ${providerKey} provider config.`, - code: ErrCode.InvalidRequest, - statusCode: 400, - data: { modelName }, - }); - } - - return computeMarkedUpCost({ - inputCostPerMillion: model.cost.input, - outputCostPerMillion: model.cost.output, - input, - output, - markup, - }); -}; export const getCreditCost = async ({ featureId, From 34778145e1849afa32771fc2851d9d554671aab8 Mon Sep 17 00:00:00 2001 From: Charlie Lamb Date: Tue, 2 Jun 2026 19:54:47 +0100 Subject: [PATCH 25/46] chore: agent comments --- .../hooks/useCreditSystemForm.ts | 23 +++++++------------ 1 file changed, 8 insertions(+), 15 deletions(-) diff --git a/vite/src/views/products/features/credit-systems/hooks/useCreditSystemForm.ts b/vite/src/views/products/features/credit-systems/hooks/useCreditSystemForm.ts index 55bb89b7c..9f643f08a 100644 --- a/vite/src/views/products/features/credit-systems/hooks/useCreditSystemForm.ts +++ b/vite/src/views/products/features/credit-systems/hooks/useCreditSystemForm.ts @@ -1,7 +1,5 @@ import type { Feature, ModelMarkups } from "@autumn/shared"; import { FeatureType } from "@autumn/shared"; -import { useStore } from "@tanstack/react-form"; -import { useEffect, useRef } from "react"; import { useAppForm } from "@/hooks/form/form"; export interface CreditSystemFormValues { @@ -28,25 +26,20 @@ export function useCreditSystemForm({ name: feature?.name ?? "", id: feature?.id ?? "", type: feature?.type ?? FeatureType.CreditSystem, - config: feature?.config ?? { schema: [{ metered_feature_id: "", feature_amount: 1, credit_amount: 0 }] }, + config: feature?.config ?? { + schema: [ + { metered_feature_id: "", feature_amount: 1, credit_amount: 0 }, + ], + }, event_names: feature?.event_names ?? [], - model_markups: (feature?.model_markups as CreditSystemFormValues["model_markups"]) ?? {}, + model_markups: + (feature?.model_markups as CreditSystemFormValues["model_markups"]) ?? + {}, defaultMarkup: 0, } satisfies CreditSystemFormValues, onSubmit: onSubmit ? ({ value }) => onSubmit(value) : undefined, }); - // Form-level `listeners.onChange` only fires when a FieldApi instance is - // registered for the changed field (see form-core FormApi.setFieldValue). - // None of these fields are mounted via , so we subscribe to the - // store directly and push value changes out to the caller. - const values = useStore(form.store, (s) => s.values); - const onChangeRef = useRef(onChange); - onChangeRef.current = onChange; - useEffect(() => { - onChangeRef.current?.(values); - }, [values]); - return form; } From fc1b4d455b4ab44262833f473f59d60d1f6caed4 Mon Sep 17 00:00:00 2001 From: Ridhwan Hussain <73362400+TheUntraceable@users.noreply.github.com.> Date: Wed, 3 Jun 2026 14:12:35 +0100 Subject: [PATCH 26/46] fix: make model markups use provider markup if available for placeholder --- .../components/AiCreditSchemaTable.tsx | 124 +++++++++++++++--- 1 file changed, 104 insertions(+), 20 deletions(-) diff --git a/vite/src/views/products/features/credit-systems/components/AiCreditSchemaTable.tsx b/vite/src/views/products/features/credit-systems/components/AiCreditSchemaTable.tsx index c5fef3d33..698ed2202 100644 --- a/vite/src/views/products/features/credit-systems/components/AiCreditSchemaTable.tsx +++ b/vite/src/views/products/features/credit-systems/components/AiCreditSchemaTable.tsx @@ -1,12 +1,12 @@ import type { ModelsDevProvider } from "@autumn/shared"; -import type { ColumnDef, Row } from "@tanstack/react-table"; import { useStore } from "@tanstack/react-form"; +import type { ColumnDef, Row } from "@tanstack/react-table"; import { InfoIcon, PlusIcon, X } from "lucide-react"; import { useMemo } from "react"; import { Table } from "@/components/general/table"; import { IconButton } from "@/components/v2/buttons/IconButton"; +import { Input } from "@/components/v2/inputs/Input"; import { - Tooltip, TooltipContent, TooltipTrigger, @@ -22,6 +22,34 @@ interface ModelRow { modelKey: string; } +function MarkupCell({ + form, + fullId, + providerKey, +}: { + form: CreditSystemFormInstance; + fullId: string; + providerKey: string; +}) { + const defaultMarkup = useStore(form.store, (s) => s.values.defaultMarkup); + const providerMarkup = useStore( + form.store, + (s) => s.values.provider_markups[providerKey]?.markup, + ); + const inheritedMarkup = providerMarkup ?? defaultMarkup; + + return ( + + ); +} + interface AiCreditSchemaTableProps { form: CreditSystemFormInstance; providerKey: string; @@ -30,6 +58,8 @@ interface AiCreditSchemaTableProps { provider: ModelsDevProvider; isLoading: boolean; removeKeys: (keys: string[]) => void; + removeProvider: (providerKey: string) => void; + setProviderMarkup: (providerKey: string, value: number | undefined) => void; renameKey: (oldKey: string, newKey: string) => void; } @@ -46,10 +76,18 @@ export function AiCreditSchemaTable({ provider, isLoading, removeKeys, + removeProvider, + setProviderMarkup, renameKey, }: AiCreditSchemaTableProps) { const isCustom = providerKey === "custom"; + const defaultMarkup = useStore(form.store, (s) => s.values.defaultMarkup); + const providerMarkup = useStore( + form.store, + (s) => s.values.provider_markups[providerKey]?.markup, + ); + const data: ModelRow[] = useMemo( () => modelFullIds.map((fullId) => { @@ -72,7 +110,10 @@ export function AiCreditSchemaTable({ - renameKey(`${providerKey}/${modelKey}`, `${providerKey}/${newKey}`) + renameKey( + `${providerKey}/${modelKey}`, + `${providerKey}/${newKey}`, + ) } /> ); @@ -81,7 +122,10 @@ export function AiCreditSchemaTable({ - renameKey(`${providerKey}/${modelKey}`, `${providerKey}/${newKey}`) + renameKey( + `${providerKey}/${modelKey}`, + `${providerKey}/${newKey}`, + ) } provider={provider} isLoading={isLoading} @@ -140,12 +184,10 @@ export function AiCreditSchemaTable({ id: "markup", size: 80, cell: ({ row }: { row: Row }) => ( - ), }, @@ -155,7 +197,10 @@ export function AiCreditSchemaTable({ size: 40, enableSorting: false, cell: ({ row }: { row: Row }) => ( -
e.stopPropagation()}> +
e.stopPropagation()} + > - Use format custom/modelId in API tracking + Use format{" "} + + custom/modelId + {" "} + in API tracking )} - } - onClick={() => removeKeys(modelFullIds)} - className="!text-subtle hover:!text-foreground" - /> +
+ {!isCustom && ( +
+ Markup % + { + const raw = e.target.value; + if (raw === "" || /^-?\d*\.?\d*$/.test(raw)) { + if (raw === "") { + setProviderMarkup(providerKey, undefined); + } else { + const parsed = Number(raw); + if (!Number.isNaN(parsed)) { + setProviderMarkup(providerKey, parsed); + } + } + } + }} + placeholder={String(defaultMarkup)} + className="w-16 text-sm text-right" + /> +
+ )} + } + onClick={() => removeProvider(providerKey)} + className="!text-subtle hover:!text-foreground" + /> +
@@ -236,17 +313,24 @@ export function AiCreditSchemaTable({ onClick={() => form.setFieldValue("model_markups", (prev) => { if (isCustom) { - const existing = Object.keys(prev).filter((k) => k.startsWith("custom/")); + const existing = Object.keys(prev).filter((k) => + k.startsWith("custom/"), + ); let i = 1; while (existing.includes(`custom/model-${i}`)) i++; - return { ...prev, [`custom/model-${i}`]: { input_cost: 0, output_cost: 0 } }; + return { + ...prev, + [`custom/model-${i}`]: { input_cost: 0, output_cost: 0 }, + }; } const usedKeys = new Set( Object.keys(prev) .filter((k) => k.startsWith(`${providerKey}/`)) .map((k) => k.slice(`${providerKey}/`.length)), ); - const nextKey = Object.keys(provider.models).find((k) => !usedKeys.has(k)); + const nextKey = Object.keys(provider.models).find( + (k) => !usedKeys.has(k), + ); if (!nextKey) return prev; return { ...prev, [`${providerKey}/${nextKey}`]: {} }; }) From fddb4d4a99eee81b71f3205f9455cdf3da15159a Mon Sep 17 00:00:00 2001 From: Ridhwan Hussain <73362400+TheUntraceable@users.noreply.github.com.> Date: Wed, 3 Jun 2026 15:58:57 +0100 Subject: [PATCH 27/46] feat: implement tiered AI markup resolution in credit system --- server/src/init.ts | 3 + .../internal/features/creditSystemUtils.ts | 34 ++++- .../features/featureActions/updateFeature.ts | 66 +++++++++- server/src/internal/features/featureUtils.ts | 28 +++++ .../features/utils/constructFeatureUtils.ts | 7 ++ .../ai-markup-resolution.test.ts | 116 ++++++++++++++++++ .../track/basic/track-tokens-tiered.test.ts | 76 ++++++++++++ server/tests/setup/v2Features.ts | 23 ++++ shared/api/features/apiFeatureV1.ts | 16 ++- .../crud/common/baseFeatureParamsV1.ts | 17 ++- .../featureConfig/creditConfig.ts | 18 ++- shared/utils/agentTypes.ts | 23 +++- .../featureUtils/apiFeatureToDbFeature.ts | 19 ++- .../components/CreateFeatureSheet.tsx | 6 + .../components/UpdateFeatureSheet.tsx | 6 + .../components/AiCreditSchema.tsx | 18 ++- .../components/EditableNumberCell.tsx | 13 +- .../components/UpdateCreditSystemSheet.tsx | 27 ++-- .../credit-systems/hooks/useAiProviders.ts | 62 +++++++++- .../hooks/useCreditSystemForm.ts | 20 ++- .../utils/validateCreditSystem.ts | 12 +- .../new-feature/NewFeatureBehaviour.tsx | 20 ++- 22 files changed, 566 insertions(+), 64 deletions(-) create mode 100644 server/tests/advanced/creditSystems/ai-markup-resolution.test.ts create mode 100644 server/tests/integration/balances/track/basic/track-tokens-tiered.test.ts diff --git a/server/src/init.ts b/server/src/init.ts index 6d9ab9d43..707bd94f9 100644 --- a/server/src/init.ts +++ b/server/src/init.ts @@ -45,6 +45,7 @@ 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"; @@ -66,6 +67,8 @@ 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()]); startRedisMonitor(); diff --git a/server/src/internal/features/creditSystemUtils.ts b/server/src/internal/features/creditSystemUtils.ts index 3b47a3009..2c5742e3e 100644 --- a/server/src/internal/features/creditSystemUtils.ts +++ b/server/src/internal/features/creditSystemUtils.ts @@ -97,6 +97,29 @@ const computeMarkedUpCost = ({ .mul(new Decimal(1).add(new Decimal(markup).div(100))) .toNumber(); +const resolveAiMarkup = ({ + modelName, + creditSystem, + modelMarkup, +}: { + modelName: string; + creditSystem: Feature; + modelMarkup?: { markup?: number | null } | null; +}) => { + if (modelMarkup?.markup != null) { + return modelMarkup.markup; + } + + const [providerKey] = modelName.split("/"); + const providerMarkup = + creditSystem.config?.provider_markups?.[providerKey]?.markup; + if (providerMarkup != null) { + return providerMarkup; + } + + return creditSystem.config?.default_markup ?? 0; +}; + const getModelCreditCost = async ({ modelName, creditSystem, @@ -108,7 +131,11 @@ const getModelCreditCost = async ({ } & TokenInput) => { const markups = creditSystem.model_markups || {}; const markupEntry = markups[modelName]; - const { markup } = markupEntry ?? { markup: 0 }; + const markup = resolveAiMarkup({ + modelName, + creditSystem, + modelMarkup: markupEntry, + }); if (modelName.startsWith("custom/")) { if (markupEntry?.input_cost == null || markupEntry?.output_cost == null) { @@ -170,7 +197,10 @@ export const getCreditCost = async ({ modelName?: string; tokens?: TokenInput; }) => { - if (creditSystem.type !== FeatureType.CreditSystem && creditSystem.type !== FeatureType.AiCreditSystem) { + if ( + creditSystem.type !== FeatureType.CreditSystem && + creditSystem.type !== FeatureType.AiCreditSystem + ) { return amount; } if (creditSystem.type === FeatureType.AiCreditSystem) { diff --git a/server/src/internal/features/featureActions/updateFeature.ts b/server/src/internal/features/featureActions/updateFeature.ts index c597bbf4b..143daad0f 100644 --- a/server/src/internal/features/featureActions/updateFeature.ts +++ b/server/src/internal/features/featureActions/updateFeature.ts @@ -1,5 +1,6 @@ import { type CreditSchemaItem, + type CreditSystemConfig, ErrCode, type Feature, FeatureType, @@ -29,10 +30,6 @@ interface UpdateFeatureParams { updates: Partial; } -/** - * Checks if the credit schema has changed between old and new config. - * Returns true if schema changed (different items or different credit amounts). - */ const areModelMarkupsEqual = ({ a, b, @@ -61,6 +58,36 @@ const areModelMarkupsEqual = ({ return true; }; +const areProviderMarkupsEqual = ({ + a, + b, +}: { + a: CreditSystemConfig["provider_markups"]; + b: CreditSystemConfig["provider_markups"]; +}): boolean => { + const aIsAbsent = a == null; + const bIsAbsent = b == null; + if (aIsAbsent && bIsAbsent) return true; + if (aIsAbsent || bIsAbsent) return false; + + const aKeys = Object.keys(a); + const bKeys = Object.keys(b); + if (aKeys.length !== bKeys.length) return false; + + for (const key of aKeys) { + const aEntry = a[key]; + const bEntry = b[key]; + if (!bEntry) return false; + if (aEntry.markup !== bEntry.markup) return false; + } + + return true; +}; + +/** + * Checks if the credit schema has changed between old and new config. + * Returns true if schema changed (different items or different credit amounts). + */ const hasCreditSchemaChanged = ({ oldSchema, newSchema, @@ -88,6 +115,26 @@ const hasCreditSchemaChanged = ({ return false; }; +const hasAiMarkupConfigChanged = ({ + oldConfig, + newConfig, +}: { + oldConfig: CreditSystemConfig | undefined; + newConfig: CreditSystemConfig | undefined; +}): boolean => { + if ( + (oldConfig?.default_markup ?? undefined) !== + (newConfig?.default_markup ?? undefined) + ) { + return true; + } + + return !areProviderMarkupsEqual({ + a: oldConfig?.provider_markups, + b: newConfig?.provider_markups, + }); +}; + /** * Updates an existing feature with full validation logic */ @@ -229,7 +276,6 @@ export const updateFeature = async ({ }); } - // Queue cache clear for credit system if schema or model markups changed const isCreditSystem = feature.type === FeatureType.CreditSystem || feature.type === FeatureType.AiCreditSystem; @@ -248,7 +294,15 @@ export const updateFeature = async ({ b: feature.model_markups, }); - if (schemaChanged || markupsChanged) { + const aiMarkupConfigChanged = + feature.type === FeatureType.AiCreditSystem && + updates.config != null && + hasAiMarkupConfigChanged({ + oldConfig: feature.config, + newConfig, + }); + + if (schemaChanged || markupsChanged || aiMarkupConfigChanged) { await addTaskToQueue({ jobName: JobName.ClearCreditSystemCustomerCache, payload: { diff --git a/server/src/internal/features/featureUtils.ts b/server/src/internal/features/featureUtils.ts index 49e737982..6a3877ec6 100644 --- a/server/src/internal/features/featureUtils.ts +++ b/server/src/internal/features/featureUtils.ts @@ -77,6 +77,34 @@ export const validateCreditSystem = ( } const newConfig = { ...config, schema, usage_type: FeatureUsageType.Single }; + const defaultMarkup = newConfig.default_markup; + if (defaultMarkup != null) { + const parsedDefaultMarkup = Number(defaultMarkup); + if (Number.isNaN(parsedDefaultMarkup) || parsedDefaultMarkup < 0) { + throw new RecaseError({ + message: "Default markup should be a non-negative number", + code: ErrCode.InvalidFeature, + statusCode: 400, + }); + } + newConfig.default_markup = parsedDefaultMarkup; + } + + const providerMarkups = newConfig.provider_markups; + if (providerMarkups != null) { + for (const [provider, entry] of Object.entries(providerMarkups)) { + const markup = Number(entry?.markup); + if (!provider || Number.isNaN(markup) || markup < 0) { + throw new RecaseError({ + message: "Provider markups must be non-negative numbers", + code: ErrCode.InvalidFeature, + statusCode: 400, + }); + } + entry.markup = markup; + } + } + for (let i = 0; i < newConfig.schema.length; i++) { const creditAmount = parseFloat( newConfig.schema[i].credit_amount.toString(), diff --git a/server/src/internal/features/utils/constructFeatureUtils.ts b/server/src/internal/features/utils/constructFeatureUtils.ts index 29c712eed..3fb8925ed 100644 --- a/server/src/internal/features/utils/constructFeatureUtils.ts +++ b/server/src/internal/features/utils/constructFeatureUtils.ts @@ -5,6 +5,7 @@ import { FeatureType, FeatureUsageType, type ModelMarkups, + type ProviderMarkups, } from "@autumn/shared"; import { generateId, keyToTitle } from "@server/utils/genUtils"; @@ -167,16 +168,22 @@ export const constructAiCreditSystem = ({ orgId, env, modelMarkups, + defaultMarkup, + providerMarkups, }: { featureId: string; name?: string; orgId: string; env: AppEnv; modelMarkups: ModelMarkups; + defaultMarkup?: number; + providerMarkups?: ProviderMarkups; }) => { const config = { schema: [], usage_type: FeatureUsageType.Single, + default_markup: defaultMarkup, + provider_markups: providerMarkups, }; const newFeature: Feature = { diff --git a/server/tests/advanced/creditSystems/ai-markup-resolution.test.ts b/server/tests/advanced/creditSystems/ai-markup-resolution.test.ts new file mode 100644 index 000000000..390b3307b --- /dev/null +++ b/server/tests/advanced/creditSystems/ai-markup-resolution.test.ts @@ -0,0 +1,116 @@ +import { describe, expect, test } from "bun:test"; +import { + type Feature, + FeatureType, + FeatureUsageType, + type ModelMarkups, + type ProviderMarkups, +} from "@autumn/shared"; +import { getCreditCost } from "@/internal/features/creditSystemUtils.js"; + +// Custom models carry their own input/output costs, so getCreditCost resolves +// them without hitting the models.dev pricing fetch — ideal for unit-testing +// the tiered markup resolution (model > provider > global > none). +const CUSTOM_MODEL = "custom/foo"; +const TOKENS = { input: 1000, output: 500 }; +// base cost = (1000 * 1000 + 500 * 2000) / 1_000_000 = 2.0 +const BASE_COST = 2.0; +const INPUT_COST = 1000; +const OUTPUT_COST = 2000; + +const makeAiCredit = ({ + model_markups, + default_markup, + provider_markups, +}: { + model_markups: ModelMarkups; + default_markup?: number; + provider_markups?: ProviderMarkups; +}): Feature => ({ + internal_id: "fe_ai_credits", + org_id: "org_test", + created_at: Date.now(), + env: "sandbox" as Feature["env"], + id: "ai_credits", + name: "AI Credits", + type: FeatureType.AiCreditSystem, + config: { + schema: [], + usage_type: FeatureUsageType.Single, + default_markup, + provider_markups, + }, + archived: false, + event_names: [], + model_markups, +}); + +const cost = (creditSystem: Feature) => + getCreditCost({ + featureId: "ai_credits", + creditSystem, + tokens: TOKENS, + modelName: CUSTOM_MODEL, + }); + +describe("getCreditCost — tiered AI markup resolution", () => { + test("per-model markup wins over provider and global", async () => { + const creditSystem = makeAiCredit({ + model_markups: { + [CUSTOM_MODEL]: { + input_cost: INPUT_COST, + output_cost: OUTPUT_COST, + markup: 50, + }, + }, + provider_markups: { custom: { markup: 20 } }, + default_markup: 10, + }); + expect(await cost(creditSystem)).toBeCloseTo(BASE_COST * 1.5); + }); + + test("falls back to provider markup when model markup is omitted", async () => { + const creditSystem = makeAiCredit({ + model_markups: { + [CUSTOM_MODEL]: { input_cost: INPUT_COST, output_cost: OUTPUT_COST }, + }, + provider_markups: { custom: { markup: 20 } }, + default_markup: 10, + }); + expect(await cost(creditSystem)).toBeCloseTo(BASE_COST * 1.2); + }); + + test("falls back to global default markup when model and provider are omitted", async () => { + const creditSystem = makeAiCredit({ + model_markups: { + [CUSTOM_MODEL]: { input_cost: INPUT_COST, output_cost: OUTPUT_COST }, + }, + default_markup: 10, + }); + expect(await cost(creditSystem)).toBeCloseTo(BASE_COST * 1.1); + }); + + test("bills at base cost (1:1) when no markup is configured anywhere", async () => { + const creditSystem = makeAiCredit({ + model_markups: { + [CUSTOM_MODEL]: { input_cost: INPUT_COST, output_cost: OUTPUT_COST }, + }, + }); + expect(await cost(creditSystem)).toBeCloseTo(BASE_COST); + }); + + test("explicit per-model markup of 0 overrides provider and global", async () => { + const creditSystem = makeAiCredit({ + model_markups: { + [CUSTOM_MODEL]: { + input_cost: INPUT_COST, + output_cost: OUTPUT_COST, + markup: 0, + }, + }, + provider_markups: { custom: { markup: 20 } }, + default_markup: 10, + }); + expect(await cost(creditSystem)).toBeCloseTo(BASE_COST); + }); +}); diff --git a/server/tests/integration/balances/track/basic/track-tokens-tiered.test.ts b/server/tests/integration/balances/track/basic/track-tokens-tiered.test.ts new file mode 100644 index 000000000..8dec5f359 --- /dev/null +++ b/server/tests/integration/balances/track/basic/track-tokens-tiered.test.ts @@ -0,0 +1,76 @@ +import { expect, test } from "bun:test"; + +import type { ApiCustomerV3, TrackResponseV2 } from "@autumn/shared"; +import { TestFeature } from "@tests/setup/v2Features.js"; +import { items } from "@tests/utils/fixtures/items.js"; +import { products } from "@tests/utils/fixtures/products.js"; +import { initScenario, s } from "@tests/utils/testInitUtils/initScenario.js"; +import chalk from "chalk"; +import { Decimal } from "decimal.js"; + +// ═══════════════════════════════════════════════════════════════════ +// TRACK-TOKENS-TIERED: per-model override vs provider markup fallback +// +// Uses custom/* models so pricing is deterministic (no models.dev fetch). +// AiCreditsTiered config: defaultMarkup=10, providerMarkups.custom=30. +// ═══════════════════════════════════════════════════════════════════ + +test.concurrent( + `${chalk.yellowBright("track-tokens-tiered: per-model override wins over provider markup fallback")}`, + async () => { + const aiCreditsItem = items.free({ + featureId: TestFeature.AiCreditsTiered, + includedUsage: 1000, + }); + const freeProd = products.base({ + id: "free", + items: [aiCreditsItem], + }); + + const { customerId, autumnV1, autumnV2 } = await initScenario({ + customerId: "track-tokens-tiered", + setup: [ + s.customer({ testClock: false }), + s.products({ list: [freeProd] }), + ], + actions: [s.attach({ productId: freeProd.id })], + }); + + const inputTokens = 10000; + const outputTokens = 5000; + // base = (10000 * 10 + 5000 * 20) / 1_000_000 = 0.2 + const baseCost = new Decimal(10) + .mul(inputTokens) + .add(new Decimal(20).mul(outputTokens)) + .div(1_000_000); + + // Per-model override of 5% wins over provider (30%) and global (10%). + const overrideCost = baseCost.mul(1.05).toNumber(); // 0.21 + const overrideRes: TrackResponseV2 = await autumnV2.post("/track_tokens", { + customer_id: customerId, + feature_id: TestFeature.AiCreditsTiered, + model_id: "custom/override-model", + input_tokens: inputTokens, + output_tokens: outputTokens, + }); + expect(overrideRes.value).toBeCloseTo(overrideCost, 10); + + // No per-model markup -> inherits the "custom" provider markup of 30%. + const providerCost = baseCost.mul(1.3).toNumber(); // 0.26 + const providerRes: TrackResponseV2 = await autumnV2.post("/track_tokens", { + customer_id: customerId, + feature_id: TestFeature.AiCreditsTiered, + model_id: "custom/provider-fallback-model", + input_tokens: inputTokens, + output_tokens: outputTokens, + }); + expect(providerRes.value).toBeCloseTo(providerCost, 10); + + const totalCost = new Decimal(overrideCost).plus(providerCost).toNumber(); + const customer = await autumnV1.customers.get(customerId); + expect(customer.features[TestFeature.AiCreditsTiered]).toMatchObject({ + balance: new Decimal(1000).minus(totalCost).toNumber(), + usage: totalCost, + }); + }, +); diff --git a/server/tests/setup/v2Features.ts b/server/tests/setup/v2Features.ts index dd66e05aa..56b117301 100644 --- a/server/tests/setup/v2Features.ts +++ b/server/tests/setup/v2Features.ts @@ -29,6 +29,7 @@ export enum TestFeature { AiCredits = "ai_credits", // AI credit system (models.dev pricing) AiCredits2 = "ai_credits_2", // second AI credit system (for disambiguation tests) + AiCreditsTiered = "ai_credits_tiered", // AI credit system with global + provider markup tiers Orbs = "orbs", // credit system that wraps an AI credit system (1000 orbs per $1) } @@ -160,6 +161,28 @@ export const getFeatures = ({ orgId }: { orgId: string }) => ({ }, }, }), + [TestFeature.AiCreditsTiered]: constructAiCreditSystem({ + featureId: TestFeature.AiCreditsTiered, + orgId, + env: AppEnv.Sandbox, + defaultMarkup: 10, + providerMarkups: { + custom: { markup: 30 }, + }, + modelMarkups: { + // Per-model override wins over provider/global. + "custom/override-model": { + markup: 5, + input_cost: 10, + output_cost: 20, + }, + // No markup -> inherits the "custom" provider markup (30%). + "custom/provider-fallback-model": { + input_cost: 10, + output_cost: 20, + }, + }, + }), [TestFeature.Orbs]: constructCreditSystem({ featureId: TestFeature.Orbs, orgId, diff --git a/shared/api/features/apiFeatureV1.ts b/shared/api/features/apiFeatureV1.ts index bc3b92cfc..9ff81d6f4 100644 --- a/shared/api/features/apiFeatureV1.ts +++ b/shared/api/features/apiFeatureV1.ts @@ -1,5 +1,8 @@ import { z } from "zod/v4"; -import { ModelMarkupsSchema } from "../../models/featureModels/featureConfig/creditConfig"; +import { + ModelMarkupsSchema, + ProviderMarkupsSchema, +} from "../../models/featureModels/featureConfig/creditConfig"; import { FeatureType } from "../../models/featureModels/featureEnums"; export const ApiFeatureV1Schema = z.object({ @@ -44,7 +47,16 @@ export const ApiFeatureV1Schema = z.object({ }), model_markups: ModelMarkupsSchema.optional().meta({ - description: "Per-model markup percentages for AI credit systems.", + description: "Per-model markup overrides for AI credit systems.", + }), + + default_markup: z.number().min(0).optional().meta({ + description: "Default percentage markup for AI credit systems.", + }), + + provider_markups: ProviderMarkupsSchema.optional().meta({ + description: + "Per-provider default markup percentages for AI credit systems.", }), display: z diff --git a/shared/api/features/crud/common/baseFeatureParamsV1.ts b/shared/api/features/crud/common/baseFeatureParamsV1.ts index 69fee0f97..e9a103ed7 100644 --- a/shared/api/features/crud/common/baseFeatureParamsV1.ts +++ b/shared/api/features/crud/common/baseFeatureParamsV1.ts @@ -1,5 +1,8 @@ import { z } from "zod/v4"; -import { ModelMarkupsSchema } from "../../../../models/featureModels/featureConfig/creditConfig"; +import { + ModelMarkupsSchema, + ProviderMarkupsSchema, +} from "../../../../models/featureModels/featureConfig/creditConfig"; import { FeatureType } from "../../../../models/featureModels/featureEnums"; import { idRegex } from "../../../../utils/utils"; @@ -49,7 +52,17 @@ export const BaseFeatureV1ParamsSchema = z.object({ model_markups: ModelMarkupsSchema.optional().meta({ description: - "Per-model markup percentages for AI credit systems. Maps model IDs to their markup configuration.", + "Per-model markup overrides for AI credit systems. Maps model IDs to their markup configuration.", + }), + + default_markup: z.number().min(0).optional().meta({ + description: + "Default percentage markup for this AI credit system. Used when no model or provider markup applies.", + }), + + provider_markups: ProviderMarkupsSchema.optional().meta({ + description: + "Per-provider default markup percentages for AI credit systems. Provider keys match the first segment of model_id.", }), event_names: z.array(z.string()).optional(), diff --git a/shared/models/featureModels/featureConfig/creditConfig.ts b/shared/models/featureModels/featureConfig/creditConfig.ts index 5bc5b9b51..69f51666a 100644 --- a/shared/models/featureModels/featureConfig/creditConfig.ts +++ b/shared/models/featureModels/featureConfig/creditConfig.ts @@ -7,6 +7,17 @@ export const CreditSchemaItemSchema = z.object({ credit_amount: z.number(), }); +const MarkupEntrySchema = z.object({ + markup: z.number().min(0), // percentage markup, e.g. 20 for 20% +}); + +export const ProviderMarkupsSchema = z + .record( + z.string(), // Provider key from the model name, e.g. "openrouter" in "openrouter/anthropic/claude" + MarkupEntrySchema, + ) + .nullish(); + export const CreditSystemConfigSchema = z.object({ schema: z.array( z.object({ @@ -15,13 +26,15 @@ export const CreditSystemConfigSchema = z.object({ }), ), usage_type: z.nativeEnum(FeatureUsageType), + default_markup: z.number().min(0).optional(), + provider_markups: ProviderMarkupsSchema, }); export const ModelMarkupsSchema = z .record( z.string(), // Represents the model name in "provider/model" format, e.g. "anthropic/claude-2" - z.object({ - markup: z.number().min(0), // percentage markup, e.g. 20 for 20% + MarkupEntrySchema.extend({ + markup: z.number().min(0).optional(), // Omit to inherit provider/global markup input_cost: z.number().min(0).optional(), // $/M tokens, required for custom/ models output_cost: z.number().min(0).optional(), // $/M tokens, required for custom/ models }), @@ -31,3 +44,4 @@ export const ModelMarkupsSchema = z export type CreditSystemConfig = z.infer; export type CreditSchemaItem = z.infer; export type ModelMarkups = z.infer; +export type ProviderMarkups = z.infer; diff --git a/shared/utils/agentTypes.ts b/shared/utils/agentTypes.ts index 5eea84691..dd338b179 100644 --- a/shared/utils/agentTypes.ts +++ b/shared/utils/agentTypes.ts @@ -10,7 +10,10 @@ * - Converters: AgentFeature ↔ Feature, AgentProduct ↔ ProductV2 */ -import type { ModelMarkups } from "../models/featureModels/featureConfig/creditConfig.js"; +import type { + ModelMarkups, + ProviderMarkups, +} from "../models/featureModels/featureConfig/creditConfig.js"; import { FeatureType, FeatureUsageType, @@ -44,6 +47,8 @@ export interface AgentFeature { credit_cost: number; }> | null; model_markups?: ModelMarkups; + default_markup?: number | null; + provider_markups?: ProviderMarkups; } export interface AgentProductItem { @@ -121,6 +126,16 @@ export function agentFeatureToFeature(agentFeature: AgentFeature): Feature { credit_amount: s.credit_cost, })); } + if (agentFeature.type === "ai_credit_system") { + config.schema = []; + config.usage_type = FeatureUsageType.Single; + if (agentFeature.default_markup != null) { + config.default_markup = agentFeature.default_markup; + } + if (agentFeature.provider_markups != null) { + config.provider_markups = agentFeature.provider_markups; + } + } return { internal_id: agentFeature.id, @@ -178,6 +193,10 @@ export function agentProductToProductV2(product: AgentProduct): ProductV2 { // ============ SHARED → AGENT CONVERTERS ============ function mapFeatureTypeToAgentType(feature: Feature): AgentFeatureType { + if (feature.type === FeatureType.AiCreditSystem) { + return "ai_credit_system"; + } + if (feature.type === FeatureType.CreditSystem) { return "credit_system"; } @@ -225,6 +244,8 @@ export function featureToAgentFeature(feature: Feature): AgentFeature { } if (feature.type === FeatureType.AiCreditSystem) { agentFeature.model_markups = feature.model_markups; + agentFeature.default_markup = feature.config?.default_markup; + agentFeature.provider_markups = feature.config?.provider_markups; } return agentFeature; diff --git a/shared/utils/featureUtils/apiFeatureToDbFeature.ts b/shared/utils/featureUtils/apiFeatureToDbFeature.ts index 4c1d0cb6e..864b4372b 100644 --- a/shared/utils/featureUtils/apiFeatureToDbFeature.ts +++ b/shared/utils/featureUtils/apiFeatureToDbFeature.ts @@ -84,11 +84,24 @@ export const featureV1ToDbFeatureConfig = ({ originalFeature: Feature; }) => { const type = apiFeature.type || originalFeature.type; + const hasProviderMarkups = "provider_markups" in apiFeature; + const hasDefaultMarkup = "default_markup" in apiFeature; - if (apiFeature.type === FeatureType.AiCreditSystem) { + if ( + type === FeatureType.AiCreditSystem && + (apiFeature.type === FeatureType.AiCreditSystem || + hasDefaultMarkup || + hasProviderMarkups) + ) { return { schema: [], usage_type: FeatureUsageType.Single, + default_markup: hasDefaultMarkup + ? apiFeature.default_markup + : originalFeature.config?.default_markup, + provider_markups: hasProviderMarkups + ? apiFeature.provider_markups + : originalFeature.config?.provider_markups, }; } @@ -151,6 +164,8 @@ export const featureV1ToDbFeature = ({ if (apiFeature.type === FeatureType.AiCreditSystem) { newConfig.usage_type = FeatureUsageType.Single; newConfig.schema = []; + newConfig.default_markup = apiFeature.default_markup; + newConfig.provider_markups = apiFeature.provider_markups; } if (apiFeature.credit_schema) { @@ -219,6 +234,8 @@ export const dbToApiFeatureV1 = ({ })) : undefined, model_markups: dbFeature.model_markups ?? undefined, + default_markup: dbFeature.config?.default_markup ?? undefined, + provider_markups: dbFeature.config?.provider_markups ?? undefined, event_names: Array.isArray(dbFeature.event_names) ? dbFeature.event_names : [], diff --git a/vite/src/views/products/features/components/CreateFeatureSheet.tsx b/vite/src/views/products/features/components/CreateFeatureSheet.tsx index ab61c4be4..6569a0f4e 100644 --- a/vite/src/views/products/features/components/CreateFeatureSheet.tsx +++ b/vite/src/views/products/features/components/CreateFeatureSheet.tsx @@ -83,6 +83,12 @@ function CreateFeatureSheet({ type: feature.type, consumable: feature.config?.usage_type === FeatureUsageType.Single, model_markups: feature.model_markups ?? undefined, + default_markup: isAiCreditSystem + ? feature.config?.default_markup + : undefined, + provider_markups: isAiCreditSystem + ? feature.config?.provider_markups + : undefined, credit_schema: isAiCreditSystem ? undefined : feature.config?.schema?.map((x: CreditSchemaItem) => ({ diff --git a/vite/src/views/products/features/components/UpdateFeatureSheet.tsx b/vite/src/views/products/features/components/UpdateFeatureSheet.tsx index 60126f471..b1097af6e 100644 --- a/vite/src/views/products/features/components/UpdateFeatureSheet.tsx +++ b/vite/src/views/products/features/components/UpdateFeatureSheet.tsx @@ -67,6 +67,12 @@ function UpdateFeatureSheet({ type: feature.type, consumable: feature.config?.usage_type === FeatureUsageType.Single, model_markups: feature.model_markups ?? undefined, + default_markup: isAiCreditSystem + ? feature.config?.default_markup + : undefined, + provider_markups: isAiCreditSystem + ? feature.config?.provider_markups + : undefined, event_names: feature.event_names, display: undefined, credit_schema: isAiCreditSystem diff --git a/vite/src/views/products/features/credit-systems/components/AiCreditSchema.tsx b/vite/src/views/products/features/credit-systems/components/AiCreditSchema.tsx index f45798d27..edaf869d8 100644 --- a/vite/src/views/products/features/credit-systems/components/AiCreditSchema.tsx +++ b/vite/src/views/products/features/credit-systems/components/AiCreditSchema.tsx @@ -2,8 +2,8 @@ import { PlusIcon } 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 type { CreditSystemFormInstance } from "../hooks/useCreditSystemForm"; import { useAiProviders } from "../hooks/useAiProviders"; +import type { CreditSystemFormInstance } from "../hooks/useCreditSystemForm"; import { AiCreditSchemaTable } from "./AiCreditSchemaTable"; interface AiCreditSchemaProps { @@ -20,6 +20,8 @@ export function AiCreditSchema({ form }: AiCreditSchemaProps) { availableProviders, addProvider, removeKeys, + removeProvider, + setProviderMarkup, renameKey, } = useAiProviders(form); @@ -46,7 +48,8 @@ export function AiCreditSchema({ form }: AiCreditSchemaProps) { const provider = providers[providerKey]; const modelFullIds = providerGroups[providerKey] ?? []; const providerName = - provider?.name ?? providerKey.charAt(0).toUpperCase() + providerKey.slice(1); + provider?.name ?? + providerKey.charAt(0).toUpperCase() + providerKey.slice(1); return ( ); })}
-
e.stopPropagation()}> +
e.stopPropagation()} + > Add Provider s.values.model_markups[fullId]?.[field], ); - const placeholder = useStore(form.store, (s) => - useDefaultAsPlaceholder ? String(s.values.defaultMarkup) : "0", - ); + const placeholder = useDefaultAsPlaceholder + ? String(inheritedPlaceholder) + : "0"; const [local, setLocal] = useState(""); const [focused, setFocused] = useState(false); - const hasValue = allowUndefined ? currentValue != null && currentValue !== 0 : currentValue != null; + const hasValue = allowUndefined + ? currentValue != null && currentValue !== 0 + : currentValue != null; const displayed = focused ? local : hasValue ? String(currentValue) : ""; return ( diff --git a/vite/src/views/products/features/credit-systems/components/UpdateCreditSystemSheet.tsx b/vite/src/views/products/features/credit-systems/components/UpdateCreditSystemSheet.tsx index 504c827f9..b817cb45b 100644 --- a/vite/src/views/products/features/credit-systems/components/UpdateCreditSystemSheet.tsx +++ b/vite/src/views/products/features/credit-systems/components/UpdateCreditSystemSheet.tsx @@ -56,15 +56,6 @@ function UpdateCreditSystemSheet({ const isAiCreditSystem = values.type === FeatureType.AiCreditSystem; - const finalMarkups = { ...values.model_markups }; - if (isAiCreditSystem) { - for (const [key, entry] of Object.entries(finalMarkups)) { - if (entry.markup == null) { - finalMarkups[key] = { ...entry, markup: values.defaultMarkup }; - } - } - } - await FeatureService.updateFeature( axiosInstance, selectedCreditSystem.id, @@ -72,12 +63,17 @@ function UpdateCreditSystemSheet({ id: values.id, name: values.name, type: values.type, - model_markups: isAiCreditSystem ? finalMarkups : undefined, + model_markups: isAiCreditSystem ? values.model_markups : undefined, + default_markup: isAiCreditSystem ? values.defaultMarkup : undefined, + provider_markups: isAiCreditSystem + ? values.provider_markups + : undefined, credit_schema: isAiCreditSystem ? undefined : values.config?.schema?.map((x: CreditSchemaItem) => ({ metered_feature_id: x.metered_feature_id, - credit_cost: x.credit_amount != null ? Number(x.credit_amount) : 0, + credit_cost: + x.credit_amount != null ? Number(x.credit_amount) : 0, })), event_names: values.event_names, display: undefined, @@ -86,7 +82,10 @@ function UpdateCreditSystemSheet({ await refetch(); toast.success("Credit system updated successfully"); - onSuccess?.(selectedCreditSystem.id, values.id || selectedCreditSystem.id); + onSuccess?.( + selectedCreditSystem.id, + values.id || selectedCreditSystem.id, + ); setOpen(false); }, }); @@ -122,7 +121,9 @@ function UpdateCreditSystemSheet({ className="w-full" onClick={() => form.handleSubmit().catch((err: AxiosError) => { - toast.error(getBackendErr(err, "Failed to update credit system")); + toast.error( + getBackendErr(err, "Failed to update credit system"), + ); }) } metaShortcut="enter" diff --git a/vite/src/views/products/features/credit-systems/hooks/useAiProviders.ts b/vite/src/views/products/features/credit-systems/hooks/useAiProviders.ts index 30aaafd46..08a8db63c 100644 --- a/vite/src/views/products/features/credit-systems/hooks/useAiProviders.ts +++ b/vite/src/views/products/features/credit-systems/hooks/useAiProviders.ts @@ -17,16 +17,37 @@ export function useAiProviders(form: CreditSystemFormInstance) { const { providers, isLoading } = useModelsDevPricing(); const modelMarkups = useStore(form.store, (s) => s.values.model_markups); const defaultMarkup = useStore(form.store, (s) => s.values.defaultMarkup); + const providerMarkups = useStore( + form.store, + (s) => s.values.provider_markups, + ); - const providerGroups = useMemo(() => groupByProvider(modelMarkups), [modelMarkups]); - const activeProviderKeys = Object.keys(providerGroups); + const providerGroups = useMemo( + () => groupByProvider(modelMarkups), + [modelMarkups], + ); + // A provider is "active" if it has model overrides OR a provider-level markup. + const activeProviderKeys = useMemo( + () => + Array.from( + new Set([ + ...Object.keys(providerGroups), + ...Object.keys(providerMarkups), + ]), + ), + [providerGroups, providerMarkups], + ); const availableProviders = useMemo(() => { const filtered = Object.values(providers).filter( (p) => !activeProviderKeys.includes(p.id), ); if (!activeProviderKeys.includes("custom")) { - filtered.push({ id: "custom", name: "Custom", models: {} } as ModelsDevProvider); + filtered.push({ + id: "custom", + name: "Custom", + models: {}, + } as ModelsDevProvider); } return filtered; }, [providers, activeProviderKeys]); @@ -34,10 +55,15 @@ export function useAiProviders(form: CreditSystemFormInstance) { const addProvider = (providerKey: string) => { form.setFieldValue("model_markups", (prev) => { if (providerKey === "custom") { - const existing = Object.keys(prev).filter((k) => k.startsWith("custom/")); + const existing = Object.keys(prev).filter((k) => + k.startsWith("custom/"), + ); let i = 1; while (existing.includes(`custom/model-${i}`)) i++; - return { ...prev, [`custom/model-${i}`]: { input_cost: 0, output_cost: 0 } }; + return { + ...prev, + [`custom/model-${i}`]: { input_cost: 0, output_cost: 0 }, + }; } const provider = providers[providerKey]; if (!provider) return prev; @@ -54,6 +80,29 @@ export function useAiProviders(form: CreditSystemFormInstance) { return updated; }); + const setProviderMarkup = (providerKey: string, value: number | undefined) => + form.setFieldValue("provider_markups", (prev) => { + const updated = { ...prev }; + if (value == null) { + delete updated[providerKey]; + } else { + updated[providerKey] = { markup: value }; + } + return updated; + }); + + // Removes the whole provider section: all its model overrides and its markup. + const removeProvider = (providerKey: string) => { + form.setFieldValue("model_markups", (prev) => { + const updated = { ...prev }; + for (const k of Object.keys(updated)) { + if (k.split("/")[0] === providerKey) delete updated[k]; + } + return updated; + }); + setProviderMarkup(providerKey, undefined); + }; + const renameKey = (oldKey: string, newKey: string) => form.setFieldValue("model_markups", (prev) => { if (newKey in prev) return prev; @@ -68,11 +117,14 @@ export function useAiProviders(form: CreditSystemFormInstance) { providers, isLoading, defaultMarkup, + providerMarkups, providerGroups, activeProviderKeys, availableProviders, addProvider, removeKeys, + removeProvider, + setProviderMarkup, renameKey, }; } diff --git a/vite/src/views/products/features/credit-systems/hooks/useCreditSystemForm.ts b/vite/src/views/products/features/credit-systems/hooks/useCreditSystemForm.ts index 55bb89b7c..404ab9b67 100644 --- a/vite/src/views/products/features/credit-systems/hooks/useCreditSystemForm.ts +++ b/vite/src/views/products/features/credit-systems/hooks/useCreditSystemForm.ts @@ -11,7 +11,10 @@ export interface CreditSystemFormValues { config: Record; event_names: string[]; model_markups: NonNullable; + /** Global default markup for the AI credit system (persisted to config.default_markup). */ defaultMarkup: number; + /** Per-provider default markups (persisted to config.provider_markups). */ + provider_markups: Record; } export function useCreditSystemForm({ @@ -28,10 +31,21 @@ export function useCreditSystemForm({ name: feature?.name ?? "", id: feature?.id ?? "", type: feature?.type ?? FeatureType.CreditSystem, - config: feature?.config ?? { schema: [{ metered_feature_id: "", feature_amount: 1, credit_amount: 0 }] }, + config: feature?.config ?? { + schema: [ + { metered_feature_id: "", feature_amount: 1, credit_amount: 0 }, + ], + }, event_names: feature?.event_names ?? [], - model_markups: (feature?.model_markups as CreditSystemFormValues["model_markups"]) ?? {}, - defaultMarkup: 0, + model_markups: + (feature?.model_markups as CreditSystemFormValues["model_markups"]) ?? + {}, + defaultMarkup: + (feature?.config?.default_markup as number | undefined) ?? 0, + provider_markups: + (feature?.config + ?.provider_markups as CreditSystemFormValues["provider_markups"]) ?? + {}, } satisfies CreditSystemFormValues, onSubmit: onSubmit ? ({ value }) => onSubmit(value) : undefined, }); diff --git a/vite/src/views/products/features/credit-systems/utils/validateCreditSystem.ts b/vite/src/views/products/features/credit-systems/utils/validateCreditSystem.ts index 32c27e890..518a92b4a 100644 --- a/vite/src/views/products/features/credit-systems/utils/validateCreditSystem.ts +++ b/vite/src/views/products/features/credit-systems/utils/validateCreditSystem.ts @@ -10,13 +10,11 @@ export const validateCreditSystem = ( const isAiCreditSystem = creditSystem.type === FeatureType.AiCreditSystem; if (isAiCreditSystem) { - if ( - !creditSystem.model_markups || - Object.keys(creditSystem.model_markups).length === 0 - ) - return "Add at least one model markup"; - - for (const [modelId, entry] of Object.entries(creditSystem.model_markups)) { + // No per-model rows is valid: such systems bill at the base cost, + // adjusted by any provider-level or global default markup. + for (const [modelId, entry] of Object.entries( + creditSystem.model_markups ?? {}, + )) { if (!modelId) return "Select a model for each row"; if (modelId.startsWith("custom/")) { const customModelKey = modelId.slice("custom/".length); diff --git a/vite/src/views/products/plan/components/new-feature/NewFeatureBehaviour.tsx b/vite/src/views/products/plan/components/new-feature/NewFeatureBehaviour.tsx index 928f19176..508c0353e 100644 --- a/vite/src/views/products/plan/components/new-feature/NewFeatureBehaviour.tsx +++ b/vite/src/views/products/plan/components/new-feature/NewFeatureBehaviour.tsx @@ -32,22 +32,18 @@ function NewFeatureCreditSchema({ }, onChange: (values) => { const isAi = values.type === FeatureType.AiCreditSystem; - const materializedMarkups = isAi - ? Object.fromEntries( - Object.entries(values.model_markups ?? {}).map(([key, entry]) => [ - key, - entry?.markup == null - ? { ...entry, markup: values.defaultMarkup } - : entry, - ]), - ) - : values.model_markups; setFeature({ ...feature, type: values.type, - config: values.config, - model_markups: materializedMarkups, + config: isAi + ? { + ...values.config, + default_markup: values.defaultMarkup, + provider_markups: values.provider_markups, + } + : values.config, + model_markups: values.model_markups, }); }, }); From e4201b741eb9dc35e9544c1097547808e7daaa18 Mon Sep 17 00:00:00 2001 From: Ridhwan Hussain <73362400+TheUntraceable@users.noreply.github.com.> Date: Thu, 4 Jun 2026 23:42:47 +0100 Subject: [PATCH 28/46] feat: add multi-pool token tracking and granular AI cost calculation --- packages/ai-sdk/src/index.ts | 142 +++- .../atmn/src/compose/models/featureModels.ts | 148 +++-- .../track/utils/getTokenTrackParams.ts | 16 +- .../internal/features/aiCreditSystemUtils.ts | 211 ++++-- .../creditSystems/ai-model-resolution.test.ts | 250 +++++++ .../balances/track/basic/track-tokens.test.ts | 617 +++++++++--------- .../api/balances/track/trackTokensParams.ts | 24 +- shared/models/aiModels/modelsDevTypes.ts | 64 +- .../components/AiCreditSchemaTable.tsx | 39 +- .../components/CreditSystemSchema.tsx | 12 +- .../credit-systems/hooks/useAiProviders.ts | 28 +- .../credit-systems/utils/modelMarkupUtils.ts | 24 + .../utils/validateCreditSystem.ts | 13 +- .../feature-list/CreditListColumns.tsx | 17 +- 14 files changed, 1102 insertions(+), 503 deletions(-) create mode 100644 server/tests/advanced/creditSystems/ai-model-resolution.test.ts create mode 100644 vite/src/views/products/features/credit-systems/utils/modelMarkupUtils.ts diff --git a/packages/ai-sdk/src/index.ts b/packages/ai-sdk/src/index.ts index 2f32eb08f..ce2f5a9ff 100644 --- a/packages/ai-sdk/src/index.ts +++ b/packages/ai-sdk/src/index.ts @@ -6,20 +6,48 @@ import { } from "ai"; import type { Autumn } from "autumn-js"; -type TokenCount = - | number - | { - total?: number | null; - } - | null - | undefined; +// Standalone published package: must not import from the internal @autumn/shared workspace. +const PROVIDER_SEPARATOR = "/"; -type TokenUsage = (LanguageModelV3Usage | LanguageModelUsage) & { - promptTokens?: TokenCount; - completionTokens?: TokenCount; +type NestedCount = { total?: number | null } | null; + +/** + * Lenient view over the AI SDK usage shapes we accept: the nested + * `LanguageModelV3Usage`, the flat `ai` `LanguageModelUsage` (with token details), and + * legacy `promptTokens`/`completionTokens` objects. + */ +type AnyUsage = (LanguageModelV3Usage | LanguageModelUsage) & { + promptTokens?: number | NestedCount; + completionTokens?: number | NestedCount; + inputTokenDetails?: { + noCacheTokens?: number | null; + cacheReadTokens?: number | null; + cacheWriteTokens?: number | null; + } | null; + outputTokenDetails?: { + textTokens?: number | null; + reasoningTokens?: number | null; + } | null; + cachedInputTokens?: number | null; + reasoningTokens?: number | null; }; -export const withTokenTracking = ({ +type ExclusivePools = { + inputTokens: number; + outputTokens: number; + cacheReadTokens: number; + cacheWriteTokens: number; + reasoningTokens: number; +}; + +const flatCount = ( + value: number | NestedCount | undefined, +): number | undefined => { + if (typeof value === "number") return value; + return value?.total ?? undefined; +}; + +export const withAutumn = ({ autumn, model, customerId, @@ -44,31 +72,86 @@ export const withTokenTracking = ({ properties?: Record; }) => { const provider = providerId ?? model.provider; - const modelName = `${provider}/${model.modelId}`; + const modelName = `${provider}${PROVIDER_SEPARATOR}${model.modelId}`; - const resolveTokens = (tokens: TokenCount, label: string): number => { - const value = typeof tokens === "number" ? tokens : tokens?.total; - if (value == null) + const required = (value: number | undefined, label: string): number => { + if (value == null) { throw new Error( `[Autumn] ${label} token usage was not returned by the model provider (${modelName}). This provider may not support usage tracking.`, ); + } return value; }; - const trackUsage = async (usage: TokenUsage) => { + const normalizeUsage = (usage: AnyUsage): ExclusivePools => { + const input = usage.inputTokens; + const output = usage.outputTokens; + + if (input != null && typeof input === "object") { + const cacheReadTokens = input.cacheRead ?? 0; + const cacheWriteTokens = input.cacheWrite ?? 0; + const textInput = + input.noCache ?? + (input.total != null + ? input.total - cacheReadTokens - cacheWriteTokens + : undefined); + const out = typeof output === "object" ? output : null; + const reasoningTokens = out?.reasoning ?? 0; + const textOutput = + out?.text ?? + (out?.total != null ? out.total - reasoningTokens : undefined); + return { + inputTokens: required(textInput, "Input"), + outputTokens: required(textOutput, "Output"), + cacheReadTokens: Math.max(0, cacheReadTokens), + cacheWriteTokens: Math.max(0, cacheWriteTokens), + reasoningTokens: Math.max(0, reasoningTokens), + }; + } + + const inputDetails = usage.inputTokenDetails; + const outputDetails = usage.outputTokenDetails; + const cacheReadTokens = + inputDetails?.cacheReadTokens ?? usage.cachedInputTokens ?? 0; + const cacheWriteTokens = inputDetails?.cacheWriteTokens ?? 0; + const reasoningTokens = + outputDetails?.reasoningTokens ?? usage.reasoningTokens ?? 0; + + const rawInput = + typeof input === "number" ? input : flatCount(usage.promptTokens); + const textInput = + inputDetails?.noCacheTokens ?? + (rawInput != null + ? rawInput - cacheReadTokens - cacheWriteTokens + : undefined); + + const rawOutput = + typeof output === "number" ? output : flatCount(usage.completionTokens); + const textOutput = + outputDetails?.textTokens ?? + (rawOutput != null ? rawOutput - reasoningTokens : undefined); + + return { + inputTokens: Math.max(0, required(textInput, "Input")), + outputTokens: Math.max(0, required(textOutput, "Output")), + cacheReadTokens: Math.max(0, cacheReadTokens), + cacheWriteTokens: Math.max(0, cacheWriteTokens), + reasoningTokens: Math.max(0, reasoningTokens), + }; + }; + + const trackUsage = async (usage: AnyUsage) => { try { - // @ts-ignore trackTokens is generated from OpenAPI; local autumn-js types may not include it yet. + const pools = normalizeUsage(usage); + // @ts-expect-error trackTokens is generated from OpenAPI; local autumn-js types may not include it yet. await autumn.balances.trackTokens({ customerId, modelId: modelName, - inputTokens: resolveTokens( - usage.inputTokens ?? usage.promptTokens, - "Input", - ), - outputTokens: resolveTokens( - usage.outputTokens ?? usage.completionTokens, - "Output", - ), + inputTokens: pools.inputTokens, + outputTokens: pools.outputTokens, + cacheReadTokens: pools.cacheReadTokens, + cacheWriteTokens: pools.cacheWriteTokens, + reasoningTokens: pools.reasoningTokens, featureId, entityId, properties, @@ -82,7 +165,7 @@ export const withTokenTracking = ({ specificationVersion: "v3", wrapGenerate: async ({ doGenerate }) => { const result = await doGenerate(); - await trackUsage(result.usage); + await trackUsage(result.usage as AnyUsage); return result; }, wrapStream: async ({ doStream }) => { @@ -90,13 +173,14 @@ export const withTokenTracking = ({ let trackingPromise: Promise | undefined; - type StreamChunk = - typeof stream extends ReadableStream ? T : never; + type StreamChunk = typeof stream extends ReadableStream + ? T + : never; const transformStream = new TransformStream({ transform(chunk, controller) { if (chunk.type === "finish" && chunk.usage) { - trackingPromise = trackUsage(chunk.usage); + trackingPromise = trackUsage(chunk.usage as AnyUsage); } controller.enqueue(chunk); }, diff --git a/packages/atmn/src/compose/models/featureModels.ts b/packages/atmn/src/compose/models/featureModels.ts index fac022ce7..826059b0d 100644 --- a/packages/atmn/src/compose/models/featureModels.ts +++ b/packages/atmn/src/compose/models/featureModels.ts @@ -4,98 +4,108 @@ import { z } from "zod/v4"; - export const FeatureSchema = z.object({ - id: z.string().meta({ - description: - "The unique identifier for this feature, used in /check and /track calls.", - }), - name: z.string().meta({ - description: - "Human-readable name displayed in the dashboard and billing UI.", - }), - eventNames: z.array(z.string()).optional().meta({ - description: - "Event names that trigger this feature's balance. Allows multiple features to respond to a single event.", - }), - creditSchema: z - .array( - z.object({ - metered_feature_id: z.string().meta({ - description: - "ID of the metered feature that draws from this credit system.", - }), - credit_cost: z.number().meta({ - description: "Credits consumed per unit of the metered feature.", - }), - }), - ) - .optional() - .meta({ - description: - "For credit_system features: maps metered features to their credit costs.", - }), - archived: z.boolean().meta({ - description: - "Whether the feature is archived and hidden from the dashboard.", - }) + id: z.string().meta({ + description: + "The unique identifier for this feature, used in /check and /track calls.", + }), + name: z.string().meta({ + description: + "Human-readable name displayed in the dashboard and billing UI.", + }), + eventNames: z.array(z.string()).optional().meta({ + description: + "Event names that trigger this feature's balance. Allows multiple features to respond to a single event.", + }), + creditSchema: z + .array( + z.object({ + metered_feature_id: z.string().meta({ + description: + "ID of the metered feature that draws from this credit system.", + }), + credit_cost: z.number().meta({ + description: "Credits consumed per unit of the metered feature.", + }), + }), + ) + .optional() + .meta({ + description: + "For credit_system features: maps metered features to their credit costs.", + }), + archived: z.boolean().meta({ + description: + "Whether the feature is archived and hidden from the dashboard.", + }), }); - - // Base fields shared by all feature types type FeatureBase = { - /** Unique identifier for the feature */ - id: string; - /** Display name for the feature */ - name: string; - /** Whether the feature is archived */ - archived?: boolean; - /** Event names that trigger this feature */ - eventNames?: string[]; - /** Credit schema for credit_system features */ - creditSchema?: Array<{ - meteredFeatureId: string; - creditCost: number; - }>; + /** Unique identifier for the feature */ + id: string; + /** Display name for the feature */ + name: string; + /** Whether the feature is archived */ + archived?: boolean; + /** Event names that trigger this feature */ + eventNames?: string[]; + /** Credit schema for credit_system features */ + creditSchema?: Array<{ + meteredFeatureId: string; + creditCost: number; + }>; }; /** Boolean feature - no consumable field allowed */ export type BooleanFeature = FeatureBase & { - type: "boolean"; - consumable?: never; + type: "boolean"; + consumable?: never; }; /** Metered feature - requires consumable field */ export type MeteredFeature = FeatureBase & { - type: "metered"; - /** Whether usage is consumed (true) or accumulated (false) */ - consumable: boolean; + type: "metered"; + /** Whether usage is consumed (true) or accumulated (false) */ + consumable: boolean; }; /** Credit system feature - always consumable */ export type CreditSystemFeature = FeatureBase & { - type: "credit_system"; - /** Credit systems are always consumable */ - consumable?: true; - /** Required: defines how credits map to metered features */ - creditSchema: Array<{ - meteredFeatureId: string; - creditCost: number; - }>; + type: "credit_system"; + /** Credit systems are always consumable */ + consumable?: true; + /** Required: defines how credits map to metered features */ + creditSchema: Array<{ + meteredFeatureId: string; + creditCost: number; + }>; }; export type ModelMarkupEntry = { - markup: number; - inputCost?: number; - outputCost?: number; + /** Per-model markup override. Omit to inherit provider/global markup. */ + markup?: number; + inputCost?: number; + outputCost?: number; +}; + +export type ProviderMarkupEntry = { + markup: number; }; /** AI credit system feature - uses model-based pricing */ export type AiCreditSystemFeature = FeatureBase & { - type: "ai_credit_system"; - modelMarkups?: Record; + type: "ai_credit_system"; + /** Per-model markup overrides (highest priority). */ + modelMarkups?: Record; + /** Default markup applied when no model or provider markup matches. */ + defaultMarkup?: number; + /** Per-provider default markups, keyed by the first segment of the model id. */ + providerMarkups?: Record; }; -export type Feature = BooleanFeature | MeteredFeature | CreditSystemFeature | AiCreditSystemFeature; - +export type Feature = + | BooleanFeature + | MeteredFeature + | CreditSystemFeature + | AiCreditSystemFeature; diff --git a/server/src/internal/balances/track/utils/getTokenTrackParams.ts b/server/src/internal/balances/track/utils/getTokenTrackParams.ts index 11fbf0c90..59ddd9bfe 100644 --- a/server/src/internal/balances/track/utils/getTokenTrackParams.ts +++ b/server/src/internal/balances/track/utils/getTokenTrackParams.ts @@ -2,12 +2,12 @@ import { ErrCode, type Feature, FeatureType, + fullCustomerToCustomerEntitlements, + fullSubjectToFullCustomer, RecaseError, type TrackParams, type TrackTokensParams, } from "@autumn/shared"; -import { fullCustomerToCustomerEntitlements } from "@autumn/shared"; -import { fullSubjectToFullCustomer } from "@autumn/shared"; import type { AutumnContext } from "@/honoUtils/HonoEnv.js"; import { getOrSetCachedFullSubject } from "@/internal/customers/cache/fullSubject/actions/getOrSetCachedFullSubject.js"; import { getOrSetCachedFullCustomer } from "@/internal/customers/cusUtils/fullCustomerCacheUtils/getOrSetCachedFullCustomer.js"; @@ -127,6 +127,11 @@ export const getTokenTrackParams = async ({ tokens: { input: input.input_tokens, output: input.output_tokens, + cacheRead: input.cache_read_tokens, + cacheWrite: input.cache_write_tokens, + audioInput: input.audio_input_tokens, + audioOutput: input.audio_output_tokens, + reasoning: input.reasoning_tokens, }, }); @@ -153,6 +158,11 @@ export const getTokenTrackParams = async ({ model: input.model_id, input_tokens: input.input_tokens, output_tokens: input.output_tokens, + cache_read_tokens: input.cache_read_tokens, + cache_write_tokens: input.cache_write_tokens, + audio_input_tokens: input.audio_input_tokens, + audio_output_tokens: input.audio_output_tokens, + reasoning_tokens: input.reasoning_tokens, cost, }, idempotency_key: input.idempotency_key, @@ -163,4 +173,4 @@ export const getTokenTrackParams = async ({ }; return { body, featureDeductions }; -}; \ No newline at end of file +}; diff --git a/server/src/internal/features/aiCreditSystemUtils.ts b/server/src/internal/features/aiCreditSystemUtils.ts index cd055fa9a..1ba507c68 100644 --- a/server/src/internal/features/aiCreditSystemUtils.ts +++ b/server/src/internal/features/aiCreditSystemUtils.ts @@ -1,34 +1,155 @@ import { ErrCode, type Feature, - InternalError, + isCustomModel, + type ModelsDevCost, + type ModelsDevCostTier, + type ModelsDevModel, + type ModelsDevProvider, RecaseError, + splitModelId, } from "@autumn/shared"; import { Decimal } from "decimal.js"; import { getModelsDevPricing } from "@/internal/features/utils/getModelPricing.js"; -export type TokenInput = { input: number; output: number }; - -// Costs are in $/M tokens; markup is a percentage (e.g. 20 = +20%). -const computeMarkedUpCost = ({ - inputCostPerMillion, - outputCostPerMillion, - input, - output, - markup, -}: { - inputCostPerMillion: Decimal.Value; - outputCostPerMillion: Decimal.Value; +export type TokenInput = { input: number; output: number; + cacheRead?: number; + cacheWrite?: number; + audioInput?: number; + audioOutput?: number; + reasoning?: number; +}; + +type ModelPricingData = Record; + +const LARGE_CONTEXT_THRESHOLD = 200_000; + +type ResolvedModel = + | { custom: true } + | { + custom: false; + providerKey: string; + modelKey: string; + model: ModelsDevModel; + }; + +const modelNotFoundError = (modelName: string) => + new RecaseError({ + message: `Model ${modelName} not found in models.dev pricing data`, + code: ErrCode.InvalidRequest, + statusCode: 400, + data: { modelName }, + }); + +/** + * Resolve a `model_id` to a models.dev entry by exact `/` lookup. The id is + * split on the first `/` (so openrouter slugs like `openrouter/openai/gpt-4o` keep their inner + * `/`); the model key must match a models.dev entry exactly. `custom/` models skip resolution. + */ +const resolveModel = ({ + modelName, + pricingData, +}: { + modelName: string; + pricingData: ModelPricingData; +}): ResolvedModel => { + if (isCustomModel(modelName)) { + return { custom: true }; + } + + const { provider, modelKey } = splitModelId(modelName); + const model = provider ? pricingData[provider]?.models[modelKey] : undefined; + if (!(provider && model)) { + throw modelNotFoundError(modelName); + } + + return { custom: false, providerKey: provider, modelKey, model }; +}; + +/** + * Resolve the effective per-token rates for a request, overlaying the active long-context + * tier (or `context_over_200k`) onto the base rates. Tier-level `cache_read`/`cache_write` + * override the base cache rates when present, so cache tokens above the threshold are billed + * at the tier rate too — not just input/output. + */ +const getEffectiveCost = ( + cost: ModelsDevCost, + totalInputTokens: number, +): ModelsDevCost => { + if (cost.tiers?.length) { + let chosen: ModelsDevCostTier | undefined; + for (const tier of cost.tiers) { + if ( + totalInputTokens > tier.tier.size && + (!chosen || tier.tier.size > chosen.tier.size) + ) { + chosen = tier; + } + } + if (chosen) { + return { + ...cost, + input: chosen.input, + output: chosen.output, + cache_read: chosen.cache_read ?? cost.cache_read, + cache_write: chosen.cache_write ?? cost.cache_write, + }; + } + return cost; + } + if (cost.context_over_200k && totalInputTokens > LARGE_CONTEXT_THRESHOLD) { + return { + ...cost, + input: cost.context_over_200k.input, + output: cost.context_over_200k.output, + cache_read: cost.context_over_200k.cache_read ?? cost.cache_read, + cache_write: cost.context_over_200k.cache_write ?? cost.cache_write, + }; + } + return cost; +}; + +const computeCost = ({ + cost, + tokens, + markup, +}: { + cost: ModelsDevCost; + tokens: TokenInput; markup: number; -}) => - new Decimal(inputCostPerMillion) - .mul(input) - .add(new Decimal(outputCostPerMillion).mul(output)) +}): number => { + const cacheRead = tokens.cacheRead ?? 0; + const cacheWrite = tokens.cacheWrite ?? 0; + const audioInput = tokens.audioInput ?? 0; + const audioOutput = tokens.audioOutput ?? 0; + const reasoning = tokens.reasoning ?? 0; + + const totalInput = tokens.input + cacheRead + cacheWrite; + const effective = getEffectiveCost(cost, totalInput); + const inputRate = effective.input; + const outputRate = effective.output; + + // Pools without a published rate fall back to the base text rate. + const cacheReadRate = effective.cache_read ?? inputRate; + const cacheWriteRate = effective.cache_write ?? inputRate; + const audioInputRate = effective.input_audio ?? inputRate; + const audioOutputRate = effective.output_audio ?? outputRate; + const reasoningRate = effective.reasoning ?? outputRate; + + return new Decimal(inputRate) + .mul(tokens.input) + .add(new Decimal(outputRate).mul(tokens.output)) + .add(new Decimal(cacheReadRate).mul(cacheRead)) + .add(new Decimal(cacheWriteRate).mul(cacheWrite)) + .add(new Decimal(audioInputRate).mul(audioInput)) + .add(new Decimal(audioOutputRate).mul(audioOutput)) + .add(new Decimal(reasoningRate).mul(reasoning)) .div(1_000_000) .mul(new Decimal(1).add(new Decimal(markup).div(100))) .toNumber(); +}; const resolveAiMarkup = ({ modelName, @@ -43,9 +164,10 @@ const resolveAiMarkup = ({ return modelMarkup.markup; } - const [providerKey] = modelName.split("/"); - const providerMarkup = - creditSystem.config?.provider_markups?.[providerKey]?.markup; + const { provider } = splitModelId(modelName); + const providerMarkup = provider + ? creditSystem.config?.provider_markups?.[provider]?.markup + : undefined; if (providerMarkup != null) { return providerMarkup; } @@ -56,13 +178,15 @@ const resolveAiMarkup = ({ export const getModelCreditCost = async ({ modelName, creditSystem, - input, - output, + ...tokens }: { modelName: string; creditSystem: Feature; -} & TokenInput) => { +} & TokenInput): Promise => { const markups = creditSystem.model_markups || {}; + const pricingData = await getModelsDevPricing(); + const resolved = resolveModel({ modelName, pricingData }); + const markupEntry = markups[modelName]; const markup = resolveAiMarkup({ modelName, @@ -70,7 +194,9 @@ export const getModelCreditCost = async ({ modelMarkup: markupEntry, }); - if (modelName.startsWith("custom/")) { + // Custom models carry no models.dev rates; they bill input/output at the user-supplied + // costs only (cache/audio/reasoning pools are not priced for custom models). + if (resolved.custom) { if (markupEntry?.input_cost == null || markupEntry?.output_cost == null) { throw new RecaseError({ message: `Custom model ${modelName} is missing input_cost or output_cost in model_markups`, @@ -78,41 +204,16 @@ export const getModelCreditCost = async ({ data: { modelName }, }); } - return computeMarkedUpCost({ - inputCostPerMillion: markupEntry.input_cost, - outputCostPerMillion: markupEntry.output_cost, - input, - output, + return computeCost({ + cost: { input: markupEntry.input_cost, output: markupEntry.output_cost }, + tokens: { input: tokens.input, output: tokens.output }, markup, }); } - const pricingData = await getModelsDevPricing(); - if (!pricingData) { - throw new InternalError({ - message: "Failed to fetch models.dev pricing data", - code: ErrCode.InternalError, - }); - } - - const [providerKey, ...modelParts] = modelName.split("/"); - const modelKey = modelParts.join("/"); - const model = pricingData[providerKey]?.models[modelKey]; - - if (!model) { - throw new RecaseError({ - message: `Model ${modelName} not found in models.dev pricing data ${providerKey} provider config.`, - code: ErrCode.InvalidRequest, - statusCode: 400, - data: { modelName }, - }); - } - - return computeMarkedUpCost({ - inputCostPerMillion: model.cost.input, - outputCostPerMillion: model.cost.output, - input, - output, + return computeCost({ + cost: resolved.model.cost, + tokens, markup, }); }; diff --git a/server/tests/advanced/creditSystems/ai-model-resolution.test.ts b/server/tests/advanced/creditSystems/ai-model-resolution.test.ts new file mode 100644 index 000000000..d63689846 --- /dev/null +++ b/server/tests/advanced/creditSystems/ai-model-resolution.test.ts @@ -0,0 +1,250 @@ +import { describe, expect, mock, test } from "bun:test"; +import { + type Feature, + FeatureType, + FeatureUsageType, + type ModelMarkups, + type ModelsDevProvider, +} from "@autumn/shared"; + +// Stub the models.dev fetch so resolution + pricing are deterministic and offline. +const pricingData: Record = { + anthropic: { + id: "anthropic", + name: "Anthropic", + models: { + "claude-opus-4-5": { + id: "claude-opus-4-5", + name: "Claude Opus 4.5", + cost: { input: 5, output: 25, cache_read: 0.5, cache_write: 6.25 }, + }, + "claude-3-5-haiku-20241022": { + id: "claude-3-5-haiku-20241022", + name: "Claude 3.5 Haiku", + cost: { input: 1, output: 5 }, + }, + }, + }, + openai: { + id: "openai", + name: "OpenAI", + models: { + "gpt-4o-2024-08-06": { + id: "gpt-4o-2024-08-06", + name: "GPT-4o", + cost: { input: 2.5, output: 10, cache_read: 1.25 }, + }, + "gpt-5": { + id: "gpt-5", + name: "GPT-5", + cost: { + input: 1, + output: 2, + cache_read: 0.5, + tiers: [ + { + input: 2, + output: 4, + cache_read: 1, + tier: { type: "context", size: 200_000 }, + }, + ], + context_over_200k: { input: 2, output: 4, cache_read: 1 }, + }, + }, + "omni-audio": { + id: "omni-audio", + name: "Omni Audio", + cost: { + input: 2, + output: 4, + input_audio: 3.5, + output_audio: 7, + reasoning: 10, + }, + }, + "no-cache-model": { + id: "no-cache-model", + name: "No Cache", + cost: { input: 10, output: 20 }, + }, + }, + }, + openrouter: { + id: "openrouter", + name: "OpenRouter", + models: { + // Openrouter keys contain a subprovider slash; the id keeps it after the + // first `/` split (openrouter/openai/gpt-4o-2024-08-06). + "openai/gpt-4o-2024-08-06": { + id: "openai/gpt-4o-2024-08-06", + name: "GPT-4o (OpenRouter)", + cost: { input: 3, output: 12 }, + }, + }, + }, +}; + +mock.module("@/internal/features/utils/getModelPricing.js", () => ({ + getModelsDevPricing: async () => pricingData, +})); + +const { getModelCreditCost } = await import( + "@/internal/features/aiCreditSystemUtils.js" +); + +const makeFeature = (model_markups: ModelMarkups = {}): Feature => ({ + internal_id: "fe_ai", + org_id: "org_test", + created_at: Date.now(), + env: "sandbox" as Feature["env"], + id: "ai_credits", + name: "AI Credits", + type: FeatureType.AiCreditSystem, + config: { schema: [], usage_type: FeatureUsageType.Single }, + archived: false, + event_names: [], + model_markups, +}); + +const PER_MILLION = 1_000_000; + +describe("resolveModel — exact match", () => { + test("exact provider-scoped match", async () => { + const cost = await getModelCreditCost({ + modelName: "anthropic/claude-opus-4-5", + creditSystem: makeFeature(), + input: 1000, + output: 500, + }); + expect(cost).toBeCloseTo((5 * 1000 + 25 * 500) / PER_MILLION, 10); + }); + + test("openrouter slug keeps its inner slash (split on the first '/')", async () => { + const cost = await getModelCreditCost({ + modelName: "openrouter/openai/gpt-4o-2024-08-06", + creditSystem: makeFeature(), + input: 1000, + output: 0, + }); + // openrouter's own rate (3), not openai's direct rate (2.5). + expect(cost).toBeCloseTo((3 * 1000) / PER_MILLION, 10); + }); + + test("throws when the model key does not match exactly", async () => { + expect( + getModelCreditCost({ + modelName: "anthropic/claude-3-5-haiku", + creditSystem: makeFeature(), + input: 1, + output: 1, + }), + ).rejects.toThrow(/not found/); + }); + + test("throws for a bare name with no provider", async () => { + expect( + getModelCreditCost({ + modelName: "gpt-4o", + creditSystem: makeFeature(), + input: 1, + output: 1, + }), + ).rejects.toThrow(/not found/); + }); +}); + +describe("computeCost — token pools", () => { + test("bills cache read/write at their own rates", async () => { + const cost = await getModelCreditCost({ + modelName: "anthropic/claude-opus-4-5", + creditSystem: makeFeature(), + input: 1000, + output: 500, + cacheRead: 2000, + cacheWrite: 100, + }); + const expected = + (5 * 1000 + 25 * 500 + 0.5 * 2000 + 6.25 * 100) / PER_MILLION; + expect(cost).toBeCloseTo(expected, 10); + }); + + test("uses the long-context tier once the input exceeds the threshold", async () => { + const large = await getModelCreditCost({ + modelName: "openai/gpt-5", + creditSystem: makeFeature(), + input: 300_000, + output: 1000, + }); + expect(large).toBeCloseTo((2 * 300_000 + 4 * 1000) / PER_MILLION, 10); + + const small = await getModelCreditCost({ + modelName: "openai/gpt-5", + creditSystem: makeFeature(), + input: 100_000, + output: 1000, + }); + expect(small).toBeCloseTo((1 * 100_000 + 2 * 1000) / PER_MILLION, 10); + }); + + test("bills cache at the tier rate above the context threshold", async () => { + const above = await getModelCreditCost({ + modelName: "openai/gpt-5", + creditSystem: makeFeature(), + input: 300_000, + output: 0, + cacheRead: 1000, + }); + // totalInput 301k > 200k -> tier input (2) and tier cache_read (1), not base 0.5. + expect(above).toBeCloseTo((2 * 300_000 + 1 * 1000) / PER_MILLION, 10); + + const below = await getModelCreditCost({ + modelName: "openai/gpt-5", + creditSystem: makeFeature(), + input: 1000, + output: 0, + cacheRead: 1000, + }); + // totalInput 2k < 200k -> base input (1) and base cache_read (0.5). + expect(below).toBeCloseTo((1 * 1000 + 0.5 * 1000) / PER_MILLION, 10); + }); + + test("bills audio and reasoning pools at their modality rates", async () => { + const cost = await getModelCreditCost({ + modelName: "openai/omni-audio", + creditSystem: makeFeature(), + input: 1000, + output: 500, + audioInput: 200, + audioOutput: 100, + reasoning: 50, + }); + const expected = + (2 * 1000 + 4 * 500 + 3.5 * 200 + 7 * 100 + 10 * 50) / PER_MILLION; + expect(cost).toBeCloseTo(expected, 10); + }); + + test("falls back to the base input rate for missing cache rates", async () => { + const cost = await getModelCreditCost({ + modelName: "openai/no-cache-model", + creditSystem: makeFeature(), + input: 0, + output: 0, + cacheRead: 1000, + }); + // no published cache_read rate -> bill at the base input rate (10) + expect(cost).toBeCloseTo((10 * 1000) / PER_MILLION, 10); + }); + + test("applies markup to the full total", async () => { + const cost = await getModelCreditCost({ + modelName: "anthropic/claude-opus-4-5", + creditSystem: makeFeature({ + "anthropic/claude-opus-4-5": { markup: 50 }, + }), + input: 1000, + output: 500, + }); + expect(cost).toBeCloseTo(((5 * 1000 + 25 * 500) / PER_MILLION) * 1.5, 10); + }); +}); diff --git a/server/tests/integration/balances/track/basic/track-tokens.test.ts b/server/tests/integration/balances/track/basic/track-tokens.test.ts index a97065fb7..d500f4e8a 100644 --- a/server/tests/integration/balances/track/basic/track-tokens.test.ts +++ b/server/tests/integration/balances/track/basic/track-tokens.test.ts @@ -13,349 +13,380 @@ import { getCreditCost } from "@/internal/features/creditSystemUtils.js"; // TRACK-TOKENS-1: Basic trackTokens with models.dev pricing // ═══════════════════════════════════════════════════════════════════ -test.concurrent(`${chalk.yellowBright("track-tokens-1: basic trackTokens with models.dev pricing")}`, async () => { - const aiCreditsItem = items.free({ - featureId: TestFeature.AiCredits, - includedUsage: 1000, - }); - const freeProd = products.base({ - id: "free", - items: [aiCreditsItem], - }); +test.concurrent( + `${chalk.yellowBright("track-tokens-1: basic trackTokens with models.dev pricing")}`, + async () => { + const aiCreditsItem = items.free({ + featureId: TestFeature.AiCredits, + includedUsage: 1000, + }); + const freeProd = products.base({ + id: "free", + items: [aiCreditsItem], + }); - const { customerId, autumnV1, autumnV2, ctx } = await initScenario({ - customerId: "track-tokens-1", - setup: [s.customer({ testClock: false }), s.products({ list: [freeProd] })], - actions: [s.attach({ productId: freeProd.id })], - }); + const { customerId, autumnV1, autumnV2, ctx } = await initScenario({ + customerId: "track-tokens-1", + setup: [ + s.customer({ testClock: false }), + s.products({ list: [freeProd] }), + ], + actions: [s.attach({ productId: freeProd.id })], + }); - const aiCreditFeature = ctx.features.find( - (f) => f.id === TestFeature.AiCredits, - ); - if (!aiCreditFeature) { - throw new Error(`${TestFeature.AiCredits} feature not found`); - } + const aiCreditFeature = ctx.features.find( + (f) => f.id === TestFeature.AiCredits, + ); + if (!aiCreditFeature) { + throw new Error(`${TestFeature.AiCredits} feature not found`); + } - const customerBefore = - await autumnV1.customers.get(customerId); - expect(customerBefore.features[TestFeature.AiCredits].balance).toBe(1000); + const customerBefore = + await autumnV1.customers.get(customerId); + expect(customerBefore.features[TestFeature.AiCredits].balance).toBe(1000); - const inputTokens = 1000; - const outputTokens = 500; - const modelId = "anthropic/claude-sonnet-4-20250514"; + const inputTokens = 1000; + const outputTokens = 500; + const modelId = "anthropic/claude-sonnet-4-20250514"; - const expectedCost = await getCreditCost({ - featureId: aiCreditFeature.id, - creditSystem: aiCreditFeature, - modelName: modelId, - tokens: { input: inputTokens, output: outputTokens }, - }); + const expectedCost = await getCreditCost({ + featureId: aiCreditFeature.id, + creditSystem: aiCreditFeature, + modelName: modelId, + tokens: { input: inputTokens, output: outputTokens }, + }); - const trackRes: TrackResponseV2 = await autumnV2.post("/track_tokens", { - customer_id: customerId, - feature_id: TestFeature.AiCredits, - model_id: modelId, - input_tokens: inputTokens, - output_tokens: outputTokens, - }); + const trackRes: TrackResponseV2 = await autumnV2.post("/track_tokens", { + customer_id: customerId, + feature_id: TestFeature.AiCredits, + model_id: modelId, + input_tokens: inputTokens, + output_tokens: outputTokens, + }); - expect(trackRes.customer_id).toBe(customerId); - expect(trackRes.value).toBeCloseTo(expectedCost, 10); + expect(trackRes.customer_id).toBe(customerId); + expect(trackRes.value).toBeCloseTo(expectedCost, 10); - const customerAfter = await autumnV1.customers.get(customerId); - expect(customerAfter.features[TestFeature.AiCredits]).toMatchObject({ - balance: new Decimal(1000).minus(expectedCost).toNumber(), - usage: expectedCost, - }); -}); + const customerAfter = + await autumnV1.customers.get(customerId); + expect(customerAfter.features[TestFeature.AiCredits]).toMatchObject({ + balance: new Decimal(1000).minus(expectedCost).toNumber(), + usage: expectedCost, + }); + }, +); // ═══════════════════════════════════════════════════════════════════ // TRACK-TOKENS-2: Disambiguation error + explicit feature_id resolution // ═══════════════════════════════════════════════════════════════════ -test.concurrent(`${chalk.yellowBright("track-tokens-2: disambiguation error and explicit feature_id resolution")}`, async () => { - const aiCreditsItem = items.free({ - featureId: TestFeature.AiCredits, - includedUsage: 500, - }); - const aiCredits2Item = items.free({ - featureId: TestFeature.AiCredits2, - includedUsage: 500, - }); - const freeProd = products.base({ - id: "free", - items: [aiCreditsItem, aiCredits2Item], - }); - - const { customerId, autumnV1, autumnV2, ctx } = await initScenario({ - customerId: "track-tokens-2", - setup: [s.customer({ testClock: false }), s.products({ list: [freeProd] })], - actions: [s.attach({ productId: freeProd.id })], - }); - - // Without feature_id, should fail with disambiguation error - let error: any; - try { - await autumnV2.post("/track_tokens", { - customer_id: customerId, - model_id: "anthropic/claude-sonnet-4-20250514", - input_tokens: 100, - output_tokens: 50, +test.concurrent( + `${chalk.yellowBright("track-tokens-2: disambiguation error and explicit feature_id resolution")}`, + async () => { + const aiCreditsItem = items.free({ + featureId: TestFeature.AiCredits, + includedUsage: 500, + }); + const aiCredits2Item = items.free({ + featureId: TestFeature.AiCredits2, + includedUsage: 500, + }); + const freeProd = products.base({ + id: "free", + items: [aiCreditsItem, aiCredits2Item], }); - } catch (e) { - error = e; - } - expect(error).toBeDefined(); - expect(error.message).toContain("Multiple AI credit system features"); - // With explicit feature_id, should succeed and only deduct from AiCredits - const aiCreditFeature = ctx.features.find( - (f) => f.id === TestFeature.AiCredits, - ); - if (!aiCreditFeature) { - throw new Error(`${TestFeature.AiCredits} feature not found`); - } + const { customerId, autumnV1, autumnV2, ctx } = await initScenario({ + customerId: "track-tokens-2", + setup: [ + s.customer({ testClock: false }), + s.products({ list: [freeProd] }), + ], + actions: [s.attach({ productId: freeProd.id })], + }); - const inputTokens = 2000; - const outputTokens = 1000; - const modelId = "anthropic/claude-sonnet-4-20250514"; + // Without feature_id, should fail with disambiguation error + let error: any; + try { + await autumnV2.post("/track_tokens", { + customer_id: customerId, + model_id: "anthropic/claude-sonnet-4-20250514", + input_tokens: 100, + output_tokens: 50, + }); + } catch (e) { + error = e; + } + expect(error).toBeDefined(); + expect(error.message).toContain("Multiple AI credit system features"); - const expectedCost = await getCreditCost({ - featureId: aiCreditFeature.id, - creditSystem: aiCreditFeature, - modelName: modelId, - tokens: { input: inputTokens, output: outputTokens }, - }); + // With explicit feature_id, should succeed and only deduct from AiCredits + const aiCreditFeature = ctx.features.find( + (f) => f.id === TestFeature.AiCredits, + ); + if (!aiCreditFeature) { + throw new Error(`${TestFeature.AiCredits} feature not found`); + } - const trackRes: TrackResponseV2 = await autumnV2.post("/track_tokens", { - customer_id: customerId, - feature_id: TestFeature.AiCredits, - model_id: modelId, - input_tokens: inputTokens, - output_tokens: outputTokens, - }); + const inputTokens = 2000; + const outputTokens = 1000; + const modelId = "anthropic/claude-sonnet-4-20250514"; - expect(trackRes.customer_id).toBe(customerId); - expect(trackRes.value).toBeCloseTo(expectedCost, 10); + const expectedCost = await getCreditCost({ + featureId: aiCreditFeature.id, + creditSystem: aiCreditFeature, + modelName: modelId, + tokens: { input: inputTokens, output: outputTokens }, + }); - const customer = await autumnV1.customers.get(customerId); - expect(customer.features[TestFeature.AiCredits]).toMatchObject({ - balance: new Decimal(500).minus(expectedCost).toNumber(), - usage: expectedCost, - }); - expect(customer.features[TestFeature.AiCredits2]).toMatchObject({ - balance: 500, - usage: 0, - }); -}); + const trackRes: TrackResponseV2 = await autumnV2.post("/track_tokens", { + customer_id: customerId, + feature_id: TestFeature.AiCredits, + model_id: modelId, + input_tokens: inputTokens, + output_tokens: outputTokens, + }); + + expect(trackRes.customer_id).toBe(customerId); + expect(trackRes.value).toBeCloseTo(expectedCost, 10); + + const customer = await autumnV1.customers.get(customerId); + expect(customer.features[TestFeature.AiCredits]).toMatchObject({ + balance: new Decimal(500).minus(expectedCost).toNumber(), + usage: expectedCost, + }); + expect(customer.features[TestFeature.AiCredits2]).toMatchObject({ + balance: 500, + usage: 0, + }); + }, +); // ═══════════════════════════════════════════════════════════════════ // TRACK-TOKENS-3: custom/* model pricing (with and without markup) // ═══════════════════════════════════════════════════════════════════ -test.concurrent(`${chalk.yellowBright("track-tokens-3: custom model pricing with and without markup")}`, async () => { - const aiCreditsItem = items.free({ - featureId: TestFeature.AiCredits, - includedUsage: 1000, - }); - const freeProd = products.base({ - id: "free", - items: [aiCreditsItem], - }); +test.concurrent( + `${chalk.yellowBright("track-tokens-3: custom model pricing with and without markup")}`, + async () => { + const aiCreditsItem = items.free({ + featureId: TestFeature.AiCredits, + includedUsage: 1000, + }); + const freeProd = products.base({ + id: "free", + items: [aiCreditsItem], + }); - const { customerId, autumnV1, autumnV2 } = await initScenario({ - customerId: "track-tokens-3", - setup: [s.customer({ testClock: false }), s.products({ list: [freeProd] })], - actions: [s.attach({ productId: freeProd.id })], - }); + const { customerId, autumnV1, autumnV2 } = await initScenario({ + customerId: "track-tokens-3", + setup: [ + s.customer({ testClock: false }), + s.products({ list: [freeProd] }), + ], + actions: [s.attach({ productId: freeProd.id })], + }); - // custom/internal-model: input_cost=5 $/M, output_cost=15 $/M, markup=0% - const expectedCostNoMarkup = new Decimal(5) - .mul(10000) - .add(new Decimal(15).mul(5000)) - .div(1_000_000) - .toNumber(); // 0.125 + // custom/internal-model: input_cost=5 $/M, output_cost=15 $/M, markup=0% + const expectedCostNoMarkup = new Decimal(5) + .mul(10000) + .add(new Decimal(15).mul(5000)) + .div(1_000_000) + .toNumber(); // 0.125 - const trackRes1: TrackResponseV2 = await autumnV2.post("/track_tokens", { - customer_id: customerId, - feature_id: TestFeature.AiCredits, - model_id: "custom/internal-model", - input_tokens: 10000, - output_tokens: 5000, - }); + const trackRes1: TrackResponseV2 = await autumnV2.post("/track_tokens", { + customer_id: customerId, + feature_id: TestFeature.AiCredits, + model_id: "custom/internal-model", + input_tokens: 10000, + output_tokens: 5000, + }); - expect(trackRes1.value).toBeCloseTo(expectedCostNoMarkup, 10); + expect(trackRes1.value).toBeCloseTo(expectedCostNoMarkup, 10); - // custom/marked-up-model: input_cost=10 $/M, output_cost=30 $/M, markup=50% - const baseCost = new Decimal(10) - .mul(8000) - .add(new Decimal(30).mul(2000)) - .div(1_000_000); - const expectedCostWithMarkup = baseCost - .mul(new Decimal(1).add(new Decimal(50).div(100))) - .toNumber(); // 0.21 + // custom/marked-up-model: input_cost=10 $/M, output_cost=30 $/M, markup=50% + const baseCost = new Decimal(10) + .mul(8000) + .add(new Decimal(30).mul(2000)) + .div(1_000_000); + const expectedCostWithMarkup = baseCost + .mul(new Decimal(1).add(new Decimal(50).div(100))) + .toNumber(); // 0.21 - const trackRes2: TrackResponseV2 = await autumnV2.post("/track_tokens", { - customer_id: customerId, - feature_id: TestFeature.AiCredits, - model_id: "custom/marked-up-model", - input_tokens: 8000, - output_tokens: 2000, - }); + const trackRes2: TrackResponseV2 = await autumnV2.post("/track_tokens", { + customer_id: customerId, + feature_id: TestFeature.AiCredits, + model_id: "custom/marked-up-model", + input_tokens: 8000, + output_tokens: 2000, + }); - expect(trackRes2.value).toBeCloseTo(expectedCostWithMarkup, 10); + expect(trackRes2.value).toBeCloseTo(expectedCostWithMarkup, 10); - const totalCost = new Decimal(expectedCostNoMarkup) - .plus(expectedCostWithMarkup) - .toNumber(); + const totalCost = new Decimal(expectedCostNoMarkup) + .plus(expectedCostWithMarkup) + .toNumber(); - const customer = await autumnV1.customers.get(customerId); - expect(customer.features[TestFeature.AiCredits]).toMatchObject({ - balance: new Decimal(1000).minus(totalCost).toNumber(), - usage: totalCost, - }); -}); + const customer = await autumnV1.customers.get(customerId); + expect(customer.features[TestFeature.AiCredits]).toMatchObject({ + balance: new Decimal(1000).minus(totalCost).toNumber(), + usage: totalCost, + }); + }, +); // ═══════════════════════════════════════════════════════════════════ // TRACK-TOKENS-4: models.dev pricing with markup + error for non-AI feature // ═══════════════════════════════════════════════════════════════════ -test.concurrent(`${chalk.yellowBright("track-tokens-4: models.dev markup and non-AI feature_id error")}`, async () => { - const aiCreditsItem = items.free({ - featureId: TestFeature.AiCredits, - includedUsage: 1000, - }); - const creditsItem = items.free({ - featureId: TestFeature.Credits, - includedUsage: 100, - }); - const freeProd = products.base({ - id: "free", - items: [aiCreditsItem, creditsItem], - }); - - const { customerId, autumnV1, autumnV2, ctx } = await initScenario({ - customerId: "track-tokens-4", - setup: [s.customer({ testClock: false }), s.products({ list: [freeProd] })], - actions: [s.attach({ productId: freeProd.id })], - }); - - const aiCreditFeature = ctx.features.find( - (f) => f.id === TestFeature.AiCredits, - ); - if (!aiCreditFeature) { - throw new Error(`${TestFeature.AiCredits} feature not found`); - } - - // anthropic/claude-haiku-3.5 has 20% markup in test config - const inputTokens = 50000; - const outputTokens = 10000; - const modelId = "anthropic/claude-3-5-haiku-20241022"; - - const expectedCost = await getCreditCost({ - featureId: aiCreditFeature.id, - creditSystem: aiCreditFeature, - modelName: modelId, - tokens: { input: inputTokens, output: outputTokens }, - }); - - const trackRes: TrackResponseV2 = await autumnV2.post("/track_tokens", { - customer_id: customerId, - feature_id: TestFeature.AiCredits, - model_id: modelId, - input_tokens: inputTokens, - output_tokens: outputTokens, - }); - - expect(trackRes.value).toBeCloseTo(expectedCost, 10); - - const customer = await autumnV1.customers.get(customerId); - expect(customer.features[TestFeature.AiCredits]).toMatchObject({ - balance: new Decimal(1000).minus(expectedCost).toNumber(), - usage: expectedCost, - }); - - // Pointing at a regular credit system should fail - let error: any; - try { - await autumnV2.post("/track_tokens", { - customer_id: customerId, - feature_id: TestFeature.Credits, - model_id: "anthropic/claude-sonnet-4-20250514", - input_tokens: 100, - output_tokens: 50, +test.concurrent( + `${chalk.yellowBright("track-tokens-4: models.dev markup and non-AI feature_id error")}`, + async () => { + const aiCreditsItem = items.free({ + featureId: TestFeature.AiCredits, + includedUsage: 1000, }); - } catch (e) { - error = e; - } - expect(error).toBeDefined(); - expect(error.message).toContain("not an AI credit system"); -}); + const creditsItem = items.free({ + featureId: TestFeature.Credits, + includedUsage: 100, + }); + const freeProd = products.base({ + id: "free", + items: [aiCreditsItem, creditsItem], + }); + + const { customerId, autumnV1, autumnV2, ctx } = await initScenario({ + customerId: "track-tokens-4", + setup: [ + s.customer({ testClock: false }), + s.products({ list: [freeProd] }), + ], + actions: [s.attach({ productId: freeProd.id })], + }); + + const aiCreditFeature = ctx.features.find( + (f) => f.id === TestFeature.AiCredits, + ); + if (!aiCreditFeature) { + throw new Error(`${TestFeature.AiCredits} feature not found`); + } + + // anthropic/claude-haiku-3.5 has 20% markup in test config + const inputTokens = 50000; + const outputTokens = 10000; + const modelId = "anthropic/claude-3-5-haiku-20241022"; + + const expectedCost = await getCreditCost({ + featureId: aiCreditFeature.id, + creditSystem: aiCreditFeature, + modelName: modelId, + tokens: { input: inputTokens, output: outputTokens }, + }); + + const trackRes: TrackResponseV2 = await autumnV2.post("/track_tokens", { + customer_id: customerId, + feature_id: TestFeature.AiCredits, + model_id: modelId, + input_tokens: inputTokens, + output_tokens: outputTokens, + }); + + expect(trackRes.value).toBeCloseTo(expectedCost, 10); + + const customer = await autumnV1.customers.get(customerId); + expect(customer.features[TestFeature.AiCredits]).toMatchObject({ + balance: new Decimal(1000).minus(expectedCost).toNumber(), + usage: expectedCost, + }); + + // Pointing at a regular credit system should fail + let error: any; + try { + await autumnV2.post("/track_tokens", { + customer_id: customerId, + feature_id: TestFeature.Credits, + model_id: "anthropic/claude-sonnet-4-20250514", + input_tokens: 100, + output_tokens: 50, + }); + } catch (e) { + error = e; + } + expect(error).toBeDefined(); + expect(error.message).toContain("not an AI credit system"); + }, +); // ═══════════════════════════════════════════════════════════════════ // TRACK-TOKENS-5: Multiple tracks accumulate correctly // ═══════════════════════════════════════════════════════════════════ -test.concurrent(`${chalk.yellowBright("track-tokens-5: multiple tracks accumulate balance deductions")}`, async () => { - const aiCreditsItem = items.free({ - featureId: TestFeature.AiCredits, - includedUsage: 1000, - }); - const freeProd = products.base({ - id: "free", - items: [aiCreditsItem], - }); +test.concurrent( + `${chalk.yellowBright("track-tokens-5: multiple tracks accumulate balance deductions")}`, + async () => { + const aiCreditsItem = items.free({ + featureId: TestFeature.AiCredits, + includedUsage: 1000, + }); + const freeProd = products.base({ + id: "free", + items: [aiCreditsItem], + }); - const { customerId, autumnV1, autumnV2, ctx } = await initScenario({ - customerId: "track-tokens-5", - setup: [s.customer({ testClock: false }), s.products({ list: [freeProd] })], - actions: [s.attach({ productId: freeProd.id })], - }); + const { customerId, autumnV1, autumnV2, ctx } = await initScenario({ + customerId: "track-tokens-5", + setup: [ + s.customer({ testClock: false }), + s.products({ list: [freeProd] }), + ], + actions: [s.attach({ productId: freeProd.id })], + }); - const aiCreditFeature = ctx.features.find( - (f) => f.id === TestFeature.AiCredits, - ); - if (!aiCreditFeature) { - throw new Error(`${TestFeature.AiCredits} feature not found`); - } + const aiCreditFeature = ctx.features.find( + (f) => f.id === TestFeature.AiCredits, + ); + if (!aiCreditFeature) { + throw new Error(`${TestFeature.AiCredits} feature not found`); + } - // First track: custom/internal-model (input_cost=5, output_cost=15, markup=0%) - const cost1 = await getCreditCost({ - featureId: aiCreditFeature.id, - creditSystem: aiCreditFeature, - modelName: "custom/internal-model", - tokens: { input: 5000, output: 2000 }, - }); + // First track: custom/internal-model (input_cost=5, output_cost=15, markup=0%) + const cost1 = await getCreditCost({ + featureId: aiCreditFeature.id, + creditSystem: aiCreditFeature, + modelName: "custom/internal-model", + tokens: { input: 5000, output: 2000 }, + }); - await autumnV2.post("/track_tokens", { - customer_id: customerId, - feature_id: TestFeature.AiCredits, - model_id: "custom/internal-model", - input_tokens: 5000, - output_tokens: 2000, - }); + await autumnV2.post("/track_tokens", { + customer_id: customerId, + feature_id: TestFeature.AiCredits, + model_id: "custom/internal-model", + input_tokens: 5000, + output_tokens: 2000, + }); - // Second track: custom/marked-up-model (input_cost=10, output_cost=30, markup=50%) - const cost2 = await getCreditCost({ - featureId: aiCreditFeature.id, - creditSystem: aiCreditFeature, - modelName: "custom/marked-up-model", - tokens: { input: 3000, output: 1000 }, - }); + // Second track: custom/marked-up-model (input_cost=10, output_cost=30, markup=50%) + const cost2 = await getCreditCost({ + featureId: aiCreditFeature.id, + creditSystem: aiCreditFeature, + modelName: "custom/marked-up-model", + tokens: { input: 3000, output: 1000 }, + }); - await autumnV2.post("/track_tokens", { - customer_id: customerId, - feature_id: TestFeature.AiCredits, - model_id: "custom/marked-up-model", - input_tokens: 3000, - output_tokens: 1000, - }); + await autumnV2.post("/track_tokens", { + customer_id: customerId, + feature_id: TestFeature.AiCredits, + model_id: "custom/marked-up-model", + input_tokens: 3000, + output_tokens: 1000, + }); - const totalCost = new Decimal(cost1).plus(cost2).toNumber(); + const totalCost = new Decimal(cost1).plus(cost2).toNumber(); - const customer = await autumnV1.customers.get(customerId); - expect(customer.features[TestFeature.AiCredits]).toMatchObject({ - balance: new Decimal(1000).minus(totalCost).toNumber(), - usage: totalCost, - }); -}); + const customer = await autumnV1.customers.get(customerId); + expect(customer.features[TestFeature.AiCredits]).toMatchObject({ + balance: new Decimal(1000).minus(totalCost).toNumber(), + usage: totalCost, + }); + }, +); diff --git a/shared/api/balances/track/trackTokensParams.ts b/shared/api/balances/track/trackTokensParams.ts index 818ed980a..a8ebef2c5 100644 --- a/shared/api/balances/track/trackTokensParams.ts +++ b/shared/api/balances/track/trackTokensParams.ts @@ -14,13 +14,31 @@ export const TrackTokensParamsSchema = z.object({ "The ID of the AI credit system feature. Auto-detected from the customer's entitlements if omitted — only required when a customer has multiple AI credit systems.", }), model_id: z.string().meta({ - description: "The AI model name with provider prefix (e.g., 'anthropic/claude-opus-4-6').", + description: + "The AI model as '/' (e.g. 'anthropic/claude-opus-4-8', 'openrouter/openai/gpt-4o'). The provider is the first path segment and must match a provider + model key in models.dev.", }), input_tokens: z.number().int().nonnegative().meta({ - description: "Number of input tokens consumed.", + description: + "Number of non-cached text input tokens consumed. Exclusive of cache and audio token pools.", }), output_tokens: z.number().int().nonnegative().meta({ - description: "Number of output tokens consumed.", + description: + "Number of text output tokens consumed. Exclusive of the reasoning and audio output pools.", + }), + cache_read_tokens: z.number().int().nonnegative().optional().meta({ + description: "Number of cached input tokens read.", + }), + cache_write_tokens: z.number().int().nonnegative().optional().meta({ + description: "Number of input tokens written to the cache.", + }), + audio_input_tokens: z.number().int().nonnegative().optional().meta({ + description: "Number of audio input tokens consumed.", + }), + audio_output_tokens: z.number().int().nonnegative().optional().meta({ + description: "Number of audio output tokens generated.", + }), + reasoning_tokens: z.number().int().nonnegative().optional().meta({ + description: "Number of reasoning tokens generated.", }), properties: z.record(z.string(), z.any()).optional().meta({ description: "Additional properties to attach to this usage event.", diff --git a/shared/models/aiModels/modelsDevTypes.ts b/shared/models/aiModels/modelsDevTypes.ts index 35f7bb30e..db0d72eb5 100644 --- a/shared/models/aiModels/modelsDevTypes.ts +++ b/shared/models/aiModels/modelsDevTypes.ts @@ -1,12 +1,68 @@ +/** Separator between the provider key and the model key in an AI credit `model_id`. */ +export const PROVIDER_SEPARATOR = "/"; + +/** Prefix marking a user-defined model whose price comes from `model_markups`. */ +export const CUSTOM_PROVIDER = "custom"; + +/** + * Split a `model_id` on the first {@link PROVIDER_SEPARATOR}. The model key keeps any + * remaining separators (e.g. openrouter slugs like `openrouter/openai/gpt-4o`). When no + * separator is present the id is a bare model name and `provider` is `undefined`. + */ +export const splitModelId = ( + id: string, +): { provider: string | undefined; modelKey: string } => { + const index = id.indexOf(PROVIDER_SEPARATOR); + if (index === -1) { + return { provider: undefined, modelKey: id }; + } + return { + provider: id.slice(0, index), + modelKey: id.slice(index + PROVIDER_SEPARATOR.length), + }; +}; + +/** Build a canonical `model_id` from a provider key and model key. */ +export const joinModelId = (provider: string, modelKey: string): string => + `${provider}${PROVIDER_SEPARATOR}${modelKey}`; + +/** Whether a `model_id` refers to a custom, user-priced model. */ +export const isCustomModel = (id: string): boolean => + id.startsWith(`${CUSTOM_PROVIDER}${PROVIDER_SEPARATOR}`); + +/** A context-based price tier (e.g. higher rates once the prompt exceeds `size` tokens). */ +export interface ModelsDevCostTier { + input: number; + output: number; + cache_read?: number; + cache_write?: number; + tier: { type: string; size: number }; +} + +/** Per-token rates ($/M tokens) for a model. Only `input`/`output` are guaranteed. */ +export interface ModelsDevCost { + input: number; + output: number; + cache_read?: number; + cache_write?: number; + input_audio?: number; + output_audio?: number; + reasoning?: number; + tiers?: ModelsDevCostTier[]; + context_over_200k?: { + input: number; + output: number; + cache_read?: number; + cache_write?: number; + }; +} + /** Shape of a single model from the models.dev API */ export interface ModelsDevModel { id: string; name: string; release_date?: string; - cost: { - input: number; - output: number; - }; + cost: ModelsDevCost; } /** Shape of a provider from the models.dev API */ diff --git a/vite/src/views/products/features/credit-systems/components/AiCreditSchemaTable.tsx b/vite/src/views/products/features/credit-systems/components/AiCreditSchemaTable.tsx index 698ed2202..91b71df0c 100644 --- a/vite/src/views/products/features/credit-systems/components/AiCreditSchemaTable.tsx +++ b/vite/src/views/products/features/credit-systems/components/AiCreditSchemaTable.tsx @@ -1,4 +1,8 @@ -import type { ModelsDevProvider } from "@autumn/shared"; +import { + joinModelId, + type ModelsDevProvider, + splitModelId, +} from "@autumn/shared"; import { useStore } from "@tanstack/react-form"; import type { ColumnDef, Row } from "@tanstack/react-table"; import { InfoIcon, PlusIcon, X } from "lucide-react"; @@ -13,6 +17,7 @@ import { } from "@/components/v2/tooltips/Tooltip"; import { useProductTable } from "@/views/products/hooks/useProductTable"; import type { CreditSystemFormInstance } from "../hooks/useCreditSystemForm"; +import { addCustomModelMarkup } from "../utils/modelMarkupUtils"; import { AiModelSelectDropdown } from "./AiModelSelectDropdown"; import { CustomModelInput } from "./CustomModelInput"; import { EditableNumberCell } from "./EditableNumberCell"; @@ -90,10 +95,10 @@ export function AiCreditSchemaTable({ const data: ModelRow[] = useMemo( () => - modelFullIds.map((fullId) => { - const [, ...parts] = fullId.split("/"); - return { fullId, modelKey: parts.join("/") }; - }), + modelFullIds.map((fullId) => ({ + fullId, + modelKey: splitModelId(fullId).modelKey, + })), [modelFullIds.join(",")], ); @@ -111,8 +116,8 @@ export function AiCreditSchemaTable({ modelKey={modelKey} onRename={(newKey) => renameKey( - `${providerKey}/${modelKey}`, - `${providerKey}/${newKey}`, + joinModelId(providerKey, modelKey), + joinModelId(providerKey, newKey), ) } /> @@ -123,8 +128,8 @@ export function AiCreditSchemaTable({ value={modelKey} onValueChange={(newKey) => renameKey( - `${providerKey}/${modelKey}`, - `${providerKey}/${newKey}`, + joinModelId(providerKey, modelKey), + joinModelId(providerKey, newKey), ) } provider={provider} @@ -313,26 +318,18 @@ export function AiCreditSchemaTable({ onClick={() => form.setFieldValue("model_markups", (prev) => { if (isCustom) { - const existing = Object.keys(prev).filter((k) => - k.startsWith("custom/"), - ); - let i = 1; - while (existing.includes(`custom/model-${i}`)) i++; - return { - ...prev, - [`custom/model-${i}`]: { input_cost: 0, output_cost: 0 }, - }; + return addCustomModelMarkup(prev); } const usedKeys = new Set( Object.keys(prev) - .filter((k) => k.startsWith(`${providerKey}/`)) - .map((k) => k.slice(`${providerKey}/`.length)), + .filter((k) => splitModelId(k).provider === providerKey) + .map((k) => splitModelId(k).modelKey), ); const nextKey = Object.keys(provider.models).find( (k) => !usedKeys.has(k), ); if (!nextKey) return prev; - return { ...prev, [`${providerKey}/${nextKey}`]: {} }; + return { ...prev, [joinModelId(providerKey, nextKey)]: {} }; }) } className="flex items-center gap-1 w-full px-4 py-1.5 text-xs text-muted-foreground hover:text-foreground bg-interactive-secondary border-t border-border transition-colors" diff --git a/vite/src/views/products/features/credit-systems/components/CreditSystemSchema.tsx b/vite/src/views/products/features/credit-systems/components/CreditSystemSchema.tsx index 86136ab90..5590fba4e 100644 --- a/vite/src/views/products/features/credit-systems/components/CreditSystemSchema.tsx +++ b/vite/src/views/products/features/credit-systems/components/CreditSystemSchema.tsx @@ -1,4 +1,8 @@ -import { FeatureType, type ModelsDevProvider } from "@autumn/shared"; +import { + FeatureType, + joinModelId, + type ModelsDevProvider, +} from "@autumn/shared"; import { useStore } from "@tanstack/react-form"; import { useMemo } from "react"; import { GroupedTabButton } from "@/components/v2/buttons/GroupedTabButton"; @@ -44,7 +48,7 @@ function getDefaultModelMarkups( if (!latestModel) continue; const [modelKey] = latestModel; - result[`${providerKey}/${modelKey}`] = {}; + result[joinModelId(providerKey, modelKey)] = {}; } return result; } @@ -77,7 +81,9 @@ export function CreditSystemSchema({ form.setFieldValue("type", FeatureType.CreditSystem); form.setFieldValue("config", { ...form.state.values.config, - schema: [{ metered_feature_id: "", feature_amount: 1, credit_amount: 0 }], + schema: [ + { metered_feature_id: "", feature_amount: 1, credit_amount: 0 }, + ], }); form.setFieldValue("model_markups", {}); } diff --git a/vite/src/views/products/features/credit-systems/hooks/useAiProviders.ts b/vite/src/views/products/features/credit-systems/hooks/useAiProviders.ts index 08a8db63c..8c511fbb6 100644 --- a/vite/src/views/products/features/credit-systems/hooks/useAiProviders.ts +++ b/vite/src/views/products/features/credit-systems/hooks/useAiProviders.ts @@ -1,14 +1,22 @@ -import type { ModelsDevProvider } from "@autumn/shared"; +import { + joinModelId, + type ModelsDevProvider, + splitModelId, +} from "@autumn/shared"; import { useStore } from "@tanstack/react-form"; import { useMemo } from "react"; import { useModelsDevPricing } from "@/hooks/queries/useAiModelsQuery"; +import { addCustomModelMarkup } from "../utils/modelMarkupUtils"; import type { CreditSystemFormInstance } from "./useCreditSystemForm"; function groupByProvider(markups: Record) { const groups: Record = {}; for (const fullId of Object.keys(markups)) { - const [provider] = fullId.split("/"); - (groups[provider] ??= []).push(fullId); + const { provider } = splitModelId(fullId); + if (!provider) continue; + const group = groups[provider] ?? []; + group.push(fullId); + groups[provider] = group; } return groups; } @@ -55,21 +63,13 @@ export function useAiProviders(form: CreditSystemFormInstance) { const addProvider = (providerKey: string) => { form.setFieldValue("model_markups", (prev) => { if (providerKey === "custom") { - const existing = Object.keys(prev).filter((k) => - k.startsWith("custom/"), - ); - let i = 1; - while (existing.includes(`custom/model-${i}`)) i++; - return { - ...prev, - [`custom/model-${i}`]: { input_cost: 0, output_cost: 0 }, - }; + return addCustomModelMarkup(prev); } const provider = providers[providerKey]; if (!provider) return prev; const firstKey = Object.keys(provider.models)[0]; if (!firstKey) return prev; - return { ...prev, [`${providerKey}/${firstKey}`]: {} }; + return { ...prev, [joinModelId(providerKey, firstKey)]: {} }; }); }; @@ -96,7 +96,7 @@ export function useAiProviders(form: CreditSystemFormInstance) { form.setFieldValue("model_markups", (prev) => { const updated = { ...prev }; for (const k of Object.keys(updated)) { - if (k.split("/")[0] === providerKey) delete updated[k]; + if (splitModelId(k).provider === providerKey) delete updated[k]; } return updated; }); diff --git a/vite/src/views/products/features/credit-systems/utils/modelMarkupUtils.ts b/vite/src/views/products/features/credit-systems/utils/modelMarkupUtils.ts new file mode 100644 index 000000000..faf2fb7cd --- /dev/null +++ b/vite/src/views/products/features/credit-systems/utils/modelMarkupUtils.ts @@ -0,0 +1,24 @@ +import { + CUSTOM_PROVIDER, + isCustomModel, + joinModelId, + type ModelMarkups, +} from "@autumn/shared"; + +type ModelMarkupMap = NonNullable; + +/** Append a blank custom-model row, choosing the next free `custom/model-N` key. */ +export const addCustomModelMarkup = (prev: ModelMarkupMap): ModelMarkupMap => { + const existing = Object.keys(prev).filter((key) => isCustomModel(key)); + let index = 1; + while (existing.includes(joinModelId(CUSTOM_PROVIDER, `model-${index}`))) { + index++; + } + return { + ...prev, + [joinModelId(CUSTOM_PROVIDER, `model-${index}`)]: { + input_cost: 0, + output_cost: 0, + }, + }; +}; diff --git a/vite/src/views/products/features/credit-systems/utils/validateCreditSystem.ts b/vite/src/views/products/features/credit-systems/utils/validateCreditSystem.ts index 518a92b4a..5acb1a024 100644 --- a/vite/src/views/products/features/credit-systems/utils/validateCreditSystem.ts +++ b/vite/src/views/products/features/credit-systems/utils/validateCreditSystem.ts @@ -1,4 +1,9 @@ -import { type CreateFeature, FeatureType } from "@autumn/shared"; +import { + type CreateFeature, + FeatureType, + isCustomModel, + splitModelId, +} from "@autumn/shared"; export const validateCreditSystem = ( creditSystem: CreateFeature, @@ -16,9 +21,9 @@ export const validateCreditSystem = ( creditSystem.model_markups ?? {}, )) { if (!modelId) return "Select a model for each row"; - if (modelId.startsWith("custom/")) { - const customModelKey = modelId.slice("custom/".length); - if (!customModelKey) return "Custom model ID cannot be empty"; + if (isCustomModel(modelId)) { + const { modelKey } = splitModelId(modelId); + if (!modelKey) return "Custom model ID cannot be empty"; if (entry.input_cost == null || entry.output_cost == null) return "Custom models require input and output costs"; } diff --git a/vite/src/views/products/features/feature-list/CreditListColumns.tsx b/vite/src/views/products/features/feature-list/CreditListColumns.tsx index 34ba83739..9db46978b 100644 --- a/vite/src/views/products/features/feature-list/CreditListColumns.tsx +++ b/vite/src/views/products/features/feature-list/CreditListColumns.tsx @@ -1,4 +1,9 @@ -import { type Feature, FeatureType, type ModelsDevProvider } from "@autumn/shared"; +import { + type Feature, + FeatureType, + type ModelsDevProvider, + splitModelId, +} from "@autumn/shared"; import { CoinsIcon, CpuIcon } from "@phosphor-icons/react"; import type { ColumnDef, Row } from "@tanstack/react-table"; import { AdminHover } from "@/components/general/AdminHover"; @@ -10,9 +15,9 @@ function resolveModelName( fullId: string, providers: Record, ): string { - const [providerKey, ...modelParts] = fullId.split("/"); - const modelKey = modelParts.join("/"); - return providers[providerKey]?.models[modelKey]?.name ?? fullId; + const { provider, modelKey } = splitModelId(fullId); + if (!provider) return fullId; + return providers[provider]?.models[modelKey]?.name ?? fullId; } export const createCreditListColumns = ( @@ -93,7 +98,9 @@ export const createCreditListColumns = ( ) .join(", ") || "—"; return ( -
{featureIds}
+
+ {featureIds} +
); }, }, From adda630b5184e91faa1e81135676b08a04017725 Mon Sep 17 00:00:00 2001 From: Ridhwan Hussain <73362400+TheUntraceable@users.noreply.github.com.> Date: Fri, 5 Jun 2026 19:59:03 +0100 Subject: [PATCH 29/46] refactor: extract shared AI credit system helpers and consolidate type checks --- .../track/utils/getTokenTrackParams.ts | 6 +- .../internal/features/aiCreditSystemUtils.ts | 11 ++-- .../internal/features/creditSystemUtils.ts | 3 +- .../features/featureActions/updateFeature.ts | 59 +++++++++---------- server/src/internal/features/featureUtils.ts | 6 +- .../features/utils/constructFeatureUtils.ts | 8 +-- .../features/changes/V1.2_FeatureChange.ts | 3 +- shared/index.ts | 2 + shared/utils/agentTypes.ts | 21 +++---- .../featureUtils/apiFeatureToDbFeature.ts | 29 +++++---- .../featureUtils/buildAiCreditSystemConfig.ts | 16 +++++ .../classifyFeature/isAiCreditSystem.ts | 5 ++ .../classifyFeature/isAnyCreditSystem.ts | 3 +- shared/utils/featureUtils/index.ts | 3 + .../featureUtils/resolveInheritedMarkup.ts | 5 ++ shared/utils/productDisplayUtils.ts | 10 ++-- .../product/product-item/formatProductItem.ts | 4 +- .../components/CreateFeatureSheet.tsx | 25 +++----- .../components/UpdateFeatureSheet.tsx | 30 +++------- .../components/AiCreditSchemaTable.tsx | 16 ++--- .../components/ClassicCreditSchema.tsx | 5 +- .../components/CreditSystemSchema.tsx | 4 +- .../components/UpdateCreditSystemSheet.tsx | 23 +++----- .../credit-systems/hooks/useCreditSchema.ts | 4 +- .../credit-systems/hooks/useProviderMarkup.ts | 21 +++++++ .../utils/validateCreditSystem.ts | 6 +- .../feature-list/CreditListColumns.tsx | 4 +- .../utils/buildFeatureMutationParams.ts | 49 +++++++++++++++ .../features/utils/getFeatureIcon.tsx | 6 +- .../edit-plan-feature/BillingType.tsx | 7 +-- .../EditPlanFeatureSheet.tsx | 5 +- .../edit-plan-feature/IncludedUsage.tsx | 7 +-- .../new-feature/NewFeatureBehaviour.tsx | 3 +- .../plan-card/DummyPlanFeatureRow.tsx | 6 +- .../components/plan-card/PlanFeatureRow.tsx | 5 +- 35 files changed, 241 insertions(+), 179 deletions(-) create mode 100644 shared/utils/featureUtils/buildAiCreditSystemConfig.ts create mode 100644 shared/utils/featureUtils/classifyFeature/isAiCreditSystem.ts create mode 100644 shared/utils/featureUtils/resolveInheritedMarkup.ts create mode 100644 vite/src/views/products/features/credit-systems/hooks/useProviderMarkup.ts create mode 100644 vite/src/views/products/features/utils/buildFeatureMutationParams.ts diff --git a/server/src/internal/balances/track/utils/getTokenTrackParams.ts b/server/src/internal/balances/track/utils/getTokenTrackParams.ts index 59ddd9bfe..97a2996c1 100644 --- a/server/src/internal/balances/track/utils/getTokenTrackParams.ts +++ b/server/src/internal/balances/track/utils/getTokenTrackParams.ts @@ -1,9 +1,9 @@ import { ErrCode, type Feature, - FeatureType, fullCustomerToCustomerEntitlements, fullSubjectToFullCustomer, + isAiCreditSystem, RecaseError, type TrackParams, type TrackTokensParams, @@ -30,7 +30,7 @@ const resolveAiCreditFeatureById = ({ statusCode: 404, }); } - if (candidate.type !== FeatureType.AiCreditSystem) { + if (!isAiCreditSystem(candidate.type)) { throw new RecaseError({ message: `Feature ${featureId} is not an AI credit system feature`, code: ErrCode.InvalidRequest, @@ -78,7 +78,7 @@ const resolveAiCreditFeatureFromEntitlements = async ({ ...new Map( cusEnts .filter( - (ce) => ce.entitlement.feature.type === FeatureType.AiCreditSystem, + (ce) => isAiCreditSystem(ce.entitlement.feature.type), ) .map((ce) => [ce.entitlement.feature.id, ce.entitlement.feature]), ).values(), diff --git a/server/src/internal/features/aiCreditSystemUtils.ts b/server/src/internal/features/aiCreditSystemUtils.ts index 1ba507c68..77392e1f3 100644 --- a/server/src/internal/features/aiCreditSystemUtils.ts +++ b/server/src/internal/features/aiCreditSystemUtils.ts @@ -7,6 +7,7 @@ import { type ModelsDevModel, type ModelsDevProvider, RecaseError, + resolveInheritedMarkup, splitModelId, } from "@autumn/shared"; import { Decimal } from "decimal.js"; @@ -168,11 +169,13 @@ const resolveAiMarkup = ({ const providerMarkup = provider ? creditSystem.config?.provider_markups?.[provider]?.markup : undefined; - if (providerMarkup != null) { - return providerMarkup; - } - return creditSystem.config?.default_markup ?? 0; + return ( + resolveInheritedMarkup({ + providerMarkup, + defaultMarkup: creditSystem.config?.default_markup, + }) ?? 0 + ); }; export const getModelCreditCost = async ({ diff --git a/server/src/internal/features/creditSystemUtils.ts b/server/src/internal/features/creditSystemUtils.ts index ace6f4504..c195ade41 100644 --- a/server/src/internal/features/creditSystemUtils.ts +++ b/server/src/internal/features/creditSystemUtils.ts @@ -3,6 +3,7 @@ import { ErrCode, type Feature, FeatureType, + isAiCreditSystem, isAnyCreditSystem, RecaseError, } from "@autumn/shared"; @@ -94,7 +95,7 @@ export const getCreditCost = async ({ if (!isAnyCreditSystem(creditSystem.type)) { return amount; } - if (creditSystem.type === FeatureType.AiCreditSystem) { + if (isAiCreditSystem(creditSystem.type)) { if (!tokens || !modelName) { throw new RecaseError({ message: "modelName and tokens must be provided for AI credit systems", diff --git a/server/src/internal/features/featureActions/updateFeature.ts b/server/src/internal/features/featureActions/updateFeature.ts index 71c7fb5fb..f40846c7a 100644 --- a/server/src/internal/features/featureActions/updateFeature.ts +++ b/server/src/internal/features/featureActions/updateFeature.ts @@ -4,6 +4,7 @@ import { ErrCode, type Feature, FeatureType, + isAiCreditSystem, isAnyCreditSystem, type ModelMarkups, notNullish, @@ -31,13 +32,12 @@ interface UpdateFeatureParams { updates: Partial; } -const areModelMarkupsEqual = ({ - a, - b, -}: { - a: ModelMarkups; - b: ModelMarkups; -}): boolean => { +/** Generic keyed-record equality check with a caller-supplied per-entry comparison. */ +const areMarkupRecordsEqual = ( + a: Record | null | undefined, + b: Record | null | undefined, + entriesEqual: (aEntry: T, bEntry: T) => boolean, +): boolean => { const aIsAbsent = a == null; const bIsAbsent = b == null; if (aIsAbsent && bIsAbsent) return true; @@ -51,39 +51,38 @@ const areModelMarkupsEqual = ({ const aEntry = a[key]; const bEntry = b[key]; if (!bEntry) return false; - if (aEntry.markup !== bEntry.markup) return false; - if (aEntry.input_cost !== bEntry.input_cost) return false; - if (aEntry.output_cost !== bEntry.output_cost) return false; + if (!entriesEqual(aEntry, bEntry)) return false; } return true; }; +const areModelMarkupsEqual = ({ + a, + b, +}: { + a: ModelMarkups; + b: ModelMarkups; +}): boolean => + areMarkupRecordsEqual[string]>( + a, + b, + (aEntry, bEntry) => + aEntry.markup === bEntry.markup && + aEntry.input_cost === bEntry.input_cost && + aEntry.output_cost === bEntry.output_cost, + ); + const areProviderMarkupsEqual = ({ a, b, }: { a: CreditSystemConfig["provider_markups"]; b: CreditSystemConfig["provider_markups"]; -}): boolean => { - const aIsAbsent = a == null; - const bIsAbsent = b == null; - if (aIsAbsent && bIsAbsent) return true; - if (aIsAbsent || bIsAbsent) return false; - - const aKeys = Object.keys(a); - const bKeys = Object.keys(b); - if (aKeys.length !== bKeys.length) return false; - - for (const key of aKeys) { - const aEntry = a[key]; - const bEntry = b[key]; - if (!bEntry) return false; - if (aEntry.markup !== bEntry.markup) return false; - } - - return true; -}; +}): boolean => + areMarkupRecordsEqual< + NonNullable[string] + >(a, b, (aEntry, bEntry) => aEntry.markup === bEntry.markup); /** * Checks if the credit schema has changed between old and new config. @@ -295,7 +294,7 @@ export const updateFeature = async ({ }); const aiMarkupConfigChanged = - feature.type === FeatureType.AiCreditSystem && + isAiCreditSystem(feature.type) && updates.config != null && hasAiMarkupConfigChanged({ oldConfig: feature.config, diff --git a/server/src/internal/features/featureUtils.ts b/server/src/internal/features/featureUtils.ts index 6a3877ec6..cf8eb5311 100644 --- a/server/src/internal/features/featureUtils.ts +++ b/server/src/internal/features/featureUtils.ts @@ -7,6 +7,7 @@ import { FeatureType, FeatureUsageType, type FullCustomer, + isAiCreditSystem, isAllocatedPrice, type MeteredConfig, type UsagePriceConfig, @@ -45,10 +46,9 @@ export const validateCreditSystem = ( config: CreditSystemConfig, featureType: FeatureType = FeatureType.CreditSystem, ) => { - const isAiCreditSystem = featureType === FeatureType.AiCreditSystem; const schema = Array.isArray(config?.schema) ? config.schema : []; - if (!isAiCreditSystem && schema.length === 0) { + if (!isAiCreditSystem(featureType) && schema.length === 0) { throw new RecaseError({ message: `At least one metered feature is required for credit system`, code: ErrCode.InvalidFeature, @@ -56,7 +56,7 @@ export const validateCreditSystem = ( }); } - if (isAiCreditSystem && schema.length > 0) { + if (isAiCreditSystem(featureType) && schema.length > 0) { throw new RecaseError({ message: `AI credit systems are leaf features and cannot define a schema. Model rates live in model_markups.`, code: ErrCode.InvalidFeature, diff --git a/server/src/internal/features/utils/constructFeatureUtils.ts b/server/src/internal/features/utils/constructFeatureUtils.ts index 3fb8925ed..ca5562895 100644 --- a/server/src/internal/features/utils/constructFeatureUtils.ts +++ b/server/src/internal/features/utils/constructFeatureUtils.ts @@ -1,6 +1,7 @@ import { AggregateType, type AppEnv, + buildAiCreditSystemConfig, type Feature, FeatureType, FeatureUsageType, @@ -179,12 +180,7 @@ export const constructAiCreditSystem = ({ defaultMarkup?: number; providerMarkups?: ProviderMarkups; }) => { - const config = { - schema: [], - usage_type: FeatureUsageType.Single, - default_markup: defaultMarkup, - provider_markups: providerMarkups, - }; + const config = buildAiCreditSystemConfig({ defaultMarkup, providerMarkups }); const newFeature: Feature = { internal_id: generateId("fe"), diff --git a/shared/api/features/changes/V1.2_FeatureChange.ts b/shared/api/features/changes/V1.2_FeatureChange.ts index a03642547..5bfc72148 100644 --- a/shared/api/features/changes/V1.2_FeatureChange.ts +++ b/shared/api/features/changes/V1.2_FeatureChange.ts @@ -4,6 +4,7 @@ import { defineVersionChange, } from "@api/versionUtils/versionChangeUtils/VersionChange.js"; import { FeatureType } from "@models/featureModels/featureEnums.js"; +import { isAiCreditSystem } from "@utils/featureUtils/classifyFeature/isAiCreditSystem"; import type { z } from "zod/v4"; import { ApiFeatureV1Schema } from "../apiFeatureV1.js"; import { @@ -63,7 +64,7 @@ export const V1_2_FeatureChange = defineVersionChange({ v0Type = ApiFeatureType.Boolean; } else if (input.type === FeatureType.CreditSystem) { v0Type = ApiFeatureType.CreditSystem; - } else if (input.type === FeatureType.AiCreditSystem) { + } else if (isAiCreditSystem(input.type)) { v0Type = ApiFeatureType.AiCreditSystem; } else if (input.type === FeatureType.Metered) { v0Type = input.consumable diff --git a/shared/index.ts b/shared/index.ts index b93810450..4ced77c7e 100644 --- a/shared/index.ts +++ b/shared/index.ts @@ -212,6 +212,8 @@ export * from "./utils/cusEntUtils/balanceUtils/cusEntToUsageAllowed"; export * from "./utils/cusEntUtils/index"; // Utils export * from "./utils/displayUtils"; +export * from "./utils/featureUtils/buildAiCreditSystemConfig"; +export * from "./utils/featureUtils/resolveInheritedMarkup"; export * from "./utils/fullSubjectUtils"; export * from "./utils/index"; export * from "./utils/intervalUtils"; diff --git a/shared/utils/agentTypes.ts b/shared/utils/agentTypes.ts index dd338b179..8e77a2a2d 100644 --- a/shared/utils/agentTypes.ts +++ b/shared/utils/agentTypes.ts @@ -23,6 +23,8 @@ import { AppEnv } from "../models/genModels/genEnums.js"; import { Infinite } from "../models/productModels/productEnums.js"; import type { ProductItem } from "../models/productV2Models/productItemModels/productItemModels.js"; import type { ProductV2 } from "../models/productV2Models/productV2Models.js"; +import { isAiCreditSystem } from "@utils/featureUtils/classifyFeature/isAiCreditSystem"; +import { buildAiCreditSystemConfig } from "./featureUtils/buildAiCreditSystemConfig.js"; // ============ INTERFACES ============ @@ -127,14 +129,13 @@ export function agentFeatureToFeature(agentFeature: AgentFeature): Feature { })); } if (agentFeature.type === "ai_credit_system") { - config.schema = []; - config.usage_type = FeatureUsageType.Single; - if (agentFeature.default_markup != null) { - config.default_markup = agentFeature.default_markup; - } - if (agentFeature.provider_markups != null) { - config.provider_markups = agentFeature.provider_markups; - } + Object.assign( + config, + buildAiCreditSystemConfig({ + defaultMarkup: agentFeature.default_markup, + providerMarkups: agentFeature.provider_markups, + }), + ); } return { @@ -193,7 +194,7 @@ export function agentProductToProductV2(product: AgentProduct): ProductV2 { // ============ SHARED → AGENT CONVERTERS ============ function mapFeatureTypeToAgentType(feature: Feature): AgentFeatureType { - if (feature.type === FeatureType.AiCreditSystem) { + if (isAiCreditSystem(feature.type)) { return "ai_credit_system"; } @@ -242,7 +243,7 @@ export function featureToAgentFeature(feature: Feature): AgentFeature { }), ); } - if (feature.type === FeatureType.AiCreditSystem) { + if (isAiCreditSystem(feature.type)) { agentFeature.model_markups = feature.model_markups; agentFeature.default_markup = feature.config?.default_markup; agentFeature.provider_markups = feature.config?.provider_markups; diff --git a/shared/utils/featureUtils/apiFeatureToDbFeature.ts b/shared/utils/featureUtils/apiFeatureToDbFeature.ts index 70c7ec352..1df67b00a 100644 --- a/shared/utils/featureUtils/apiFeatureToDbFeature.ts +++ b/shared/utils/featureUtils/apiFeatureToDbFeature.ts @@ -8,6 +8,7 @@ import { FeatureUsageType, } from "@models/featureModels/featureEnums.js"; import type { Feature } from "@models/featureModels/featureModels.js"; +import { isAiCreditSystem } from "@utils/featureUtils/classifyFeature/isAiCreditSystem"; import { isAnyCreditSystem } from "./classifyFeature/isAnyCreditSystem.js"; import { AppEnv } from "@models/genModels/genEnums.js"; import type { ApiFeatureV1 } from "../../api/features/apiFeatureV1.js"; @@ -24,6 +25,7 @@ import { import type { CreditSchemaItem } from "../../models/featureModels/featureConfig/creditConfig.js"; import type { SharedContext } from "../../types/sharedContext.js"; import { notNullish, nullish } from "../utils.js"; +import { buildAiCreditSystemConfig } from "./buildAiCreditSystemConfig.js"; export const apiFeatureToDbFeature = ({ apiFeature, @@ -89,21 +91,19 @@ export const featureV1ToDbFeatureConfig = ({ const hasDefaultMarkup = "default_markup" in apiFeature; if ( - type === FeatureType.AiCreditSystem && - (apiFeature.type === FeatureType.AiCreditSystem || + isAiCreditSystem(type) && + (isAiCreditSystem(apiFeature.type) || hasDefaultMarkup || hasProviderMarkups) ) { - return { - schema: [], - usage_type: FeatureUsageType.Single, - default_markup: hasDefaultMarkup + return buildAiCreditSystemConfig({ + defaultMarkup: hasDefaultMarkup ? apiFeature.default_markup : originalFeature.config?.default_markup, - provider_markups: hasProviderMarkups + providerMarkups: hasProviderMarkups ? apiFeature.provider_markups : originalFeature.config?.provider_markups, - }; + }); } if (nullish(apiFeature.consumable) && nullish(apiFeature.credit_schema)) @@ -162,11 +162,14 @@ export const featureV1ToDbFeature = ({ : FeatureUsageType.Continuous; } - if (apiFeature.type === FeatureType.AiCreditSystem) { - newConfig.usage_type = FeatureUsageType.Single; - newConfig.schema = []; - newConfig.default_markup = apiFeature.default_markup; - newConfig.provider_markups = apiFeature.provider_markups; + if (isAiCreditSystem(apiFeature.type)) { + Object.assign( + newConfig, + buildAiCreditSystemConfig({ + defaultMarkup: apiFeature.default_markup, + providerMarkups: apiFeature.provider_markups, + }), + ); } if (apiFeature.credit_schema) { diff --git a/shared/utils/featureUtils/buildAiCreditSystemConfig.ts b/shared/utils/featureUtils/buildAiCreditSystemConfig.ts new file mode 100644 index 000000000..f461a17b3 --- /dev/null +++ b/shared/utils/featureUtils/buildAiCreditSystemConfig.ts @@ -0,0 +1,16 @@ +import type { + CreditSystemConfig, + ProviderMarkups, +} from "../../models/featureModels/featureConfig/creditConfig.js"; +import { FeatureUsageType } from "../../models/featureModels/featureEnums.js"; + +/** Single factory for the AiCreditSystem `config` shape. Callers pass already-resolved markup values; no fallback resolution happens here. */ +export const buildAiCreditSystemConfig = (args: { + defaultMarkup?: number | null; + providerMarkups?: ProviderMarkups; +}): CreditSystemConfig => ({ + schema: [], + usage_type: FeatureUsageType.Single, + default_markup: args.defaultMarkup ?? undefined, + provider_markups: args.providerMarkups, +}); diff --git a/shared/utils/featureUtils/classifyFeature/isAiCreditSystem.ts b/shared/utils/featureUtils/classifyFeature/isAiCreditSystem.ts new file mode 100644 index 000000000..02f295e3d --- /dev/null +++ b/shared/utils/featureUtils/classifyFeature/isAiCreditSystem.ts @@ -0,0 +1,5 @@ +import { FeatureType } from "@models/featureModels/featureEnums"; + +export const isAiCreditSystem = ( + type: FeatureType | undefined | null, +): boolean => type === FeatureType.AiCreditSystem; diff --git a/shared/utils/featureUtils/classifyFeature/isAnyCreditSystem.ts b/shared/utils/featureUtils/classifyFeature/isAnyCreditSystem.ts index 6f4bac496..eeeefe788 100644 --- a/shared/utils/featureUtils/classifyFeature/isAnyCreditSystem.ts +++ b/shared/utils/featureUtils/classifyFeature/isAnyCreditSystem.ts @@ -1,4 +1,5 @@ import { FeatureType } from "@models/featureModels/featureEnums"; +import { isAiCreditSystem } from "@utils/featureUtils/classifyFeature/isAiCreditSystem"; export const isAnyCreditSystem = (type: FeatureType): boolean => - type === FeatureType.CreditSystem || type === FeatureType.AiCreditSystem; + type === FeatureType.CreditSystem || isAiCreditSystem(type); diff --git a/shared/utils/featureUtils/index.ts b/shared/utils/featureUtils/index.ts index 184129007..3860239c2 100644 --- a/shared/utils/featureUtils/index.ts +++ b/shared/utils/featureUtils/index.ts @@ -1,3 +1,4 @@ +import { isAiCreditSystem } from "@utils/featureUtils/classifyFeature/isAiCreditSystem"; import { isAllocatedFeature } from "@utils/featureUtils/classifyFeature/isAllocatedFeature"; import { isAnyCreditSystem } from "@utils/featureUtils/classifyFeature/isAnyCreditSystem"; import { isConsumableFeature } from "@utils/featureUtils/classifyFeature/isConsumableFeature"; @@ -9,11 +10,13 @@ export * from "./convertFeatureUtils"; export * from "./creditSystemUtils"; export * from "./findFeatureUtils"; +export { isAiCreditSystem } from "@utils/featureUtils/classifyFeature/isAiCreditSystem"; export { isAnyCreditSystem } from "@utils/featureUtils/classifyFeature/isAnyCreditSystem"; export const featureUtils = { isConsumable: isConsumableFeature, isAllocated: isAllocatedFeature, + isAiCreditSystem, isAnyCreditSystem, find: { diff --git a/shared/utils/featureUtils/resolveInheritedMarkup.ts b/shared/utils/featureUtils/resolveInheritedMarkup.ts new file mode 100644 index 000000000..60bbfa416 --- /dev/null +++ b/shared/utils/featureUtils/resolveInheritedMarkup.ts @@ -0,0 +1,5 @@ +/** Inherited markup precedence: a provider-level markup wins, otherwise the global default. Returns undefined when neither is set. */ +export const resolveInheritedMarkup = (args: { + providerMarkup?: number | null; + defaultMarkup?: number | null; +}): number | undefined => args.providerMarkup ?? args.defaultMarkup ?? undefined; diff --git a/shared/utils/productDisplayUtils.ts b/shared/utils/productDisplayUtils.ts index 7d32f4dc2..f219e76c9 100644 --- a/shared/utils/productDisplayUtils.ts +++ b/shared/utils/productDisplayUtils.ts @@ -14,6 +14,7 @@ import { isFeaturePriceItem, isPriceItem, } from "./productV2Utils/productItemUtils/getItemType.js"; +import { isAiCreditSystem } from "@utils/featureUtils/classifyFeature/isAiCreditSystem"; import { notNullish, nullish } from "./utils.js"; // ============================================================================ @@ -54,7 +55,7 @@ const getIncludedUsageText = (item: ProductItem, feature: Feature): string => { if (item.included_usage === Infinite) { return `Unlimited ${featureName}`; } - if (feature.type === FeatureType.AiCreditSystem) { + if (isAiCreditSystem(feature.type)) { return `$${numberWithCommas(item.included_usage ?? 0)} of ${featureName}`; } if (nullish(item.included_usage) || item.included_usage === 0) { @@ -221,7 +222,6 @@ export const getFeaturePriceItemDisplay = ({ // Build included usage string (e.g., "100 credits") const includedUsage = item.included_usage as number | null; const hasIncludedUsage = notNullish(includedUsage) && includedUsage > 0; - const isAiCreditSystem = feature.type === FeatureType.AiCreditSystem; const includedFeatureName = getFeatureName({ feature, @@ -229,7 +229,7 @@ export const getFeaturePriceItemDisplay = ({ }); let includedUsageStr = ""; if (hasIncludedUsage) { - if (isAiCreditSystem) { + if (isAiCreditSystem(feature.type)) { includedUsageStr = `$${numberWithCommas(includedUsage)} of ${includedFeatureName}`; } else { includedUsageStr = `${numberWithCommas(includedUsage)} ${includedFeatureName}`; @@ -258,7 +258,7 @@ export const getFeaturePriceItemDisplay = ({ units: billingUnits, }); let perUnitStr: string; - if (isAiCreditSystem) { + if (isAiCreditSystem(feature.type)) { perUnitStr = billingUnits > 1 ? `$${numberWithCommas(billingUnits)} of ${billingFeatureName}` @@ -283,7 +283,7 @@ export const getFeaturePriceItemDisplay = ({ } // Format output based on what we have - if (isAiCreditSystem) { + if (isAiCreditSystem(feature.type)) { return { primary_text: includedUsageStr || "$0 included", secondary_text: "then charged based on model usage", diff --git a/vite/src/utils/product/product-item/formatProductItem.ts b/vite/src/utils/product/product-item/formatProductItem.ts index d3eff3e7d..0d2f992a0 100644 --- a/vite/src/utils/product/product-item/formatProductItem.ts +++ b/vite/src/utils/product/product-item/formatProductItem.ts @@ -5,6 +5,7 @@ import { formatAmount, formatInterval, Infinite, + isAiCreditSystem, type ProductItem, ProductItemType, TierBehavior, @@ -155,8 +156,7 @@ const getFeatureString = ({ intervalCount: item.interval_count ?? undefined, }); - const isAiCreditSystem = feature?.type === FeatureType.AiCreditSystem; - if (isAiCreditSystem) { + if (isAiCreditSystem(feature?.type)) { const amount = item.included_usage ?? 0; const formattedAmount = amount === 0 ? "$0.00" : `$${Number(amount).toFixed(2)}`; diff --git a/vite/src/views/products/features/components/CreateFeatureSheet.tsx b/vite/src/views/products/features/components/CreateFeatureSheet.tsx index bd4a714e6..469995c00 100644 --- a/vite/src/views/products/features/components/CreateFeatureSheet.tsx +++ b/vite/src/views/products/features/components/CreateFeatureSheet.tsx @@ -1,7 +1,5 @@ import { CreateFeatureSchema, - type CreditSchemaItem, - FeatureType, FeatureUsageType, isAnyCreditSystem, } from "@autumn/shared"; @@ -24,6 +22,7 @@ import { NewFeatureBehaviour } from "../../plan/components/new-feature/NewFeatur import { NewFeatureDetails } from "../../plan/components/new-feature/NewFeatureDetails"; import { NewFeatureType } from "../../plan/components/new-feature/NewFeatureType"; import { validateCreditSystem } from "../credit-systems/utils/validateCreditSystem"; +import { buildFeatureMarkupParams } from "../utils/buildFeatureMutationParams"; import { getDefaultFeature } from "../utils/defaultFeature"; function CreateFeatureSheet({ @@ -71,8 +70,6 @@ function CreateFeatureSheet({ setLoading(false); } else { try { - const isAiCreditSystem = feature.type === FeatureType.AiCreditSystem; - const { data: createdFeature } = await FeatureService.createFeature( axiosInstance, { @@ -80,19 +77,13 @@ function CreateFeatureSheet({ id: feature.id, type: feature.type, consumable: feature.config?.usage_type === FeatureUsageType.Single, - model_markups: feature.model_markups ?? undefined, - default_markup: isAiCreditSystem - ? feature.config?.default_markup - : undefined, - provider_markups: isAiCreditSystem - ? feature.config?.provider_markups - : undefined, - credit_schema: isAiCreditSystem - ? undefined - : feature.config?.schema?.map((x: CreditSchemaItem) => ({ - metered_feature_id: x.metered_feature_id, - credit_cost: x.credit_amount, - })), + ...buildFeatureMarkupParams({ + type: feature.type, + modelMarkups: feature.model_markups ?? undefined, + defaultMarkup: feature.config?.default_markup, + providerMarkups: feature.config?.provider_markups, + schema: feature.config?.schema, + }), event_names: feature.event_names, }, ); diff --git a/vite/src/views/products/features/components/UpdateFeatureSheet.tsx b/vite/src/views/products/features/components/UpdateFeatureSheet.tsx index b1097af6e..1541e5a61 100644 --- a/vite/src/views/products/features/components/UpdateFeatureSheet.tsx +++ b/vite/src/views/products/features/components/UpdateFeatureSheet.tsx @@ -1,9 +1,4 @@ -import { - type CreditSchemaItem, - type Feature, - FeatureType, - FeatureUsageType, -} from "@autumn/shared"; +import { type Feature, FeatureUsageType } from "@autumn/shared"; import type { AxiosError } from "axios"; import { useEffect, useState } from "react"; import { toast } from "sonner"; @@ -22,6 +17,7 @@ import { NewFeatureAdvanced } from "../../plan/components/new-feature/NewFeature import { NewFeatureBehaviour } from "../../plan/components/new-feature/NewFeatureBehaviour"; import { NewFeatureDetails } from "../../plan/components/new-feature/NewFeatureDetails"; import { NewFeatureType } from "../../plan/components/new-feature/NewFeatureType"; +import { buildFeatureMarkupParams } from "../utils/buildFeatureMutationParams"; interface UpdateFeatureSheetProps { open: boolean; @@ -58,29 +54,21 @@ function UpdateFeatureSheet({ setLoading(true); try { - const isAiCreditSystem = feature.type === FeatureType.AiCreditSystem; - await FeatureService.updateFeature(axiosInstance, selectedFeature.id, { ...feature, id: feature.id || undefined, name: feature.name || undefined, type: feature.type, consumable: feature.config?.usage_type === FeatureUsageType.Single, - model_markups: feature.model_markups ?? undefined, - default_markup: isAiCreditSystem - ? feature.config?.default_markup - : undefined, - provider_markups: isAiCreditSystem - ? feature.config?.provider_markups - : undefined, event_names: feature.event_names, display: undefined, - credit_schema: isAiCreditSystem - ? undefined - : feature.config?.schema?.map((item: CreditSchemaItem) => ({ - metered_feature_id: item.metered_feature_id, - credit_cost: item.credit_amount, - })), + ...buildFeatureMarkupParams({ + type: feature.type, + modelMarkups: feature.model_markups ?? undefined, + defaultMarkup: feature.config?.default_markup, + providerMarkups: feature.config?.provider_markups, + schema: feature.config?.schema, + }), }); await refetch(); diff --git a/vite/src/views/products/features/credit-systems/components/AiCreditSchemaTable.tsx b/vite/src/views/products/features/credit-systems/components/AiCreditSchemaTable.tsx index 91b71df0c..d4544cf27 100644 --- a/vite/src/views/products/features/credit-systems/components/AiCreditSchemaTable.tsx +++ b/vite/src/views/products/features/credit-systems/components/AiCreditSchemaTable.tsx @@ -3,7 +3,6 @@ import { type ModelsDevProvider, splitModelId, } from "@autumn/shared"; -import { useStore } from "@tanstack/react-form"; import type { ColumnDef, Row } from "@tanstack/react-table"; import { InfoIcon, PlusIcon, X } from "lucide-react"; import { useMemo } from "react"; @@ -17,6 +16,7 @@ import { } from "@/components/v2/tooltips/Tooltip"; import { useProductTable } from "@/views/products/hooks/useProductTable"; import type { CreditSystemFormInstance } from "../hooks/useCreditSystemForm"; +import { useProviderMarkup } from "../hooks/useProviderMarkup"; import { addCustomModelMarkup } from "../utils/modelMarkupUtils"; import { AiModelSelectDropdown } from "./AiModelSelectDropdown"; import { CustomModelInput } from "./CustomModelInput"; @@ -36,12 +36,7 @@ function MarkupCell({ fullId: string; providerKey: string; }) { - const defaultMarkup = useStore(form.store, (s) => s.values.defaultMarkup); - const providerMarkup = useStore( - form.store, - (s) => s.values.provider_markups[providerKey]?.markup, - ); - const inheritedMarkup = providerMarkup ?? defaultMarkup; + const { inheritedMarkup } = useProviderMarkup(form, providerKey); return ( s.values.defaultMarkup); - const providerMarkup = useStore( - form.store, - (s) => s.values.provider_markups[providerKey]?.markup, + const { defaultMarkup, providerMarkup } = useProviderMarkup( + form, + providerKey, ); const data: ModelRow[] = useMemo( diff --git a/vite/src/views/products/features/credit-systems/components/ClassicCreditSchema.tsx b/vite/src/views/products/features/credit-systems/components/ClassicCreditSchema.tsx index 7a5936a95..9c2817c2f 100644 --- a/vite/src/views/products/features/credit-systems/components/ClassicCreditSchema.tsx +++ b/vite/src/views/products/features/credit-systems/components/ClassicCreditSchema.tsx @@ -1,7 +1,7 @@ import { type CreditSchemaItem, type Feature, - FeatureType, + isAiCreditSystem, } from "@autumn/shared"; import { PlusIcon } from "@phosphor-icons/react"; import { X } from "lucide-react"; @@ -47,8 +47,7 @@ export function ClassicCreditSchema({ form }: ClassicCreditSchemaProps) { const selectedFeature = allSchemaCandidateFeatures.find( (f: Feature) => f.id === item.metered_feature_id, ); - const isAiChild = - selectedFeature?.type === FeatureType.AiCreditSystem; + const isAiChild = isAiCreditSystem(selectedFeature?.type); return (
s.values.type); - const mode: CreditSchemaMode = - type === FeatureType.AiCreditSystem ? "ai" : "classic"; + const mode: CreditSchemaMode = isAiCreditSystem(type) ? "ai" : "classic"; const handleModeChange = (newMode: string) => { if (newMode === "ai") { diff --git a/vite/src/views/products/features/credit-systems/components/UpdateCreditSystemSheet.tsx b/vite/src/views/products/features/credit-systems/components/UpdateCreditSystemSheet.tsx index b817cb45b..a81747c4f 100644 --- a/vite/src/views/products/features/credit-systems/components/UpdateCreditSystemSheet.tsx +++ b/vite/src/views/products/features/credit-systems/components/UpdateCreditSystemSheet.tsx @@ -1,5 +1,4 @@ import type { CreditSchemaItem, Feature } from "@autumn/shared"; -import { FeatureType } from "@autumn/shared"; import { useStore } from "@tanstack/react-form"; import type { AxiosError } from "axios"; import { toast } from "sonner"; @@ -13,6 +12,7 @@ import { useFeaturesQuery } from "@/hooks/queries/useFeaturesQuery"; import { FeatureService } from "@/services/FeatureService"; import { useAxiosInstance } from "@/services/useAxiosInstance"; import { getBackendErr } from "@/utils/genUtils"; +import { buildFeatureMarkupParams } from "../../utils/buildFeatureMutationParams"; import { useCreditSystemForm } from "../hooks/useCreditSystemForm"; import { validateCreditSystem } from "../utils/validateCreditSystem"; import { CreditSystemDetails } from "./CreditSystemDetails"; @@ -54,8 +54,6 @@ function UpdateCreditSystemSheet({ return; } - const isAiCreditSystem = values.type === FeatureType.AiCreditSystem; - await FeatureService.updateFeature( axiosInstance, selectedCreditSystem.id, @@ -63,18 +61,13 @@ function UpdateCreditSystemSheet({ id: values.id, name: values.name, type: values.type, - model_markups: isAiCreditSystem ? values.model_markups : undefined, - default_markup: isAiCreditSystem ? values.defaultMarkup : undefined, - provider_markups: isAiCreditSystem - ? values.provider_markups - : undefined, - credit_schema: isAiCreditSystem - ? undefined - : values.config?.schema?.map((x: CreditSchemaItem) => ({ - metered_feature_id: x.metered_feature_id, - credit_cost: - x.credit_amount != null ? Number(x.credit_amount) : 0, - })), + ...buildFeatureMarkupParams({ + type: values.type, + modelMarkups: values.model_markups, + defaultMarkup: values.defaultMarkup, + providerMarkups: values.provider_markups, + schema: values.config?.schema as CreditSchemaItem[] | undefined, + }), event_names: values.event_names, display: undefined, }, diff --git a/vite/src/views/products/features/credit-systems/hooks/useCreditSchema.ts b/vite/src/views/products/features/credit-systems/hooks/useCreditSchema.ts index fa805ac04..169eb12dd 100644 --- a/vite/src/views/products/features/credit-systems/hooks/useCreditSchema.ts +++ b/vite/src/views/products/features/credit-systems/hooks/useCreditSchema.ts @@ -1,5 +1,5 @@ import type { CreditSchemaItem, Feature } from "@autumn/shared"; -import { FeatureType } from "@autumn/shared"; +import { FeatureType, isAiCreditSystem } from "@autumn/shared"; import { useStore } from "@tanstack/react-form"; import { useMemo, useRef } from "react"; import { toast } from "sonner"; @@ -22,7 +22,7 @@ export function useCreditSchema(form: CreditSystemFormInstance) { const allSchemaCandidateFeatures = features.filter( (f: Feature) => - f.type === FeatureType.Metered || f.type === FeatureType.AiCreditSystem, + f.type === FeatureType.Metered || isAiCreditSystem(f.type), ); const handleSchemaChange = ( diff --git a/vite/src/views/products/features/credit-systems/hooks/useProviderMarkup.ts b/vite/src/views/products/features/credit-systems/hooks/useProviderMarkup.ts new file mode 100644 index 000000000..54013bfaf --- /dev/null +++ b/vite/src/views/products/features/credit-systems/hooks/useProviderMarkup.ts @@ -0,0 +1,21 @@ +import { resolveInheritedMarkup } from "@autumn/shared"; +import { useStore } from "@tanstack/react-form"; +import type { CreditSystemFormInstance } from "./useCreditSystemForm"; + +/** Centralizes the default/provider markup store selectors and their inherited-markup resolution for a single provider. */ +export const useProviderMarkup = ( + form: CreditSystemFormInstance, + providerKey: string, +) => { + const defaultMarkup = useStore(form.store, (s) => s.values.defaultMarkup); + const providerMarkup = useStore( + form.store, + (s) => s.values.provider_markups[providerKey]?.markup, + ); + const inheritedMarkup = resolveInheritedMarkup({ + providerMarkup, + defaultMarkup, + }); + + return { defaultMarkup, providerMarkup, inheritedMarkup }; +}; diff --git a/vite/src/views/products/features/credit-systems/utils/validateCreditSystem.ts b/vite/src/views/products/features/credit-systems/utils/validateCreditSystem.ts index 5acb1a024..d45d278b6 100644 --- a/vite/src/views/products/features/credit-systems/utils/validateCreditSystem.ts +++ b/vite/src/views/products/features/credit-systems/utils/validateCreditSystem.ts @@ -1,6 +1,6 @@ import { type CreateFeature, - FeatureType, + isAiCreditSystem, isCustomModel, splitModelId, } from "@autumn/shared"; @@ -12,9 +12,7 @@ export const validateCreditSystem = ( return "Please fill in all fields"; } - const isAiCreditSystem = creditSystem.type === FeatureType.AiCreditSystem; - - if (isAiCreditSystem) { + if (isAiCreditSystem(creditSystem.type)) { // No per-model rows is valid: such systems bill at the base cost, // adjusted by any provider-level or global default markup. for (const [modelId, entry] of Object.entries( diff --git a/vite/src/views/products/features/feature-list/CreditListColumns.tsx b/vite/src/views/products/features/feature-list/CreditListColumns.tsx index 9db46978b..bbf35d616 100644 --- a/vite/src/views/products/features/feature-list/CreditListColumns.tsx +++ b/vite/src/views/products/features/feature-list/CreditListColumns.tsx @@ -1,6 +1,6 @@ import { type Feature, - FeatureType, + isAiCreditSystem, type ModelsDevProvider, splitModelId, } from "@autumn/shared"; @@ -59,7 +59,7 @@ export const createCreditListColumns = ( size: 120, accessorKey: "type", cell: ({ row }: { row: Row }) => { - const isAi = row.original.type === FeatureType.AiCreditSystem; + const isAi = isAiCreditSystem(row.original.type); return (
{isAi ? ( diff --git a/vite/src/views/products/features/utils/buildFeatureMutationParams.ts b/vite/src/views/products/features/utils/buildFeatureMutationParams.ts new file mode 100644 index 000000000..6a1e6a89e --- /dev/null +++ b/vite/src/views/products/features/utils/buildFeatureMutationParams.ts @@ -0,0 +1,49 @@ +import { + type CreditSchemaItem, + type FeatureType, + isAiCreditSystem, + type ModelMarkups, + type ProviderMarkups, +} from "@autumn/shared"; + +interface BuildFeatureMarkupParamsArgs { + type: FeatureType; + modelMarkups?: ModelMarkups; + defaultMarkup?: number | null; + providerMarkups?: ProviderMarkups; + schema?: CreditSchemaItem[]; +} + +interface FeatureMarkupParams { + model_markups?: ModelMarkups; + default_markup?: number | null; + provider_markups?: ProviderMarkups; + credit_schema?: { metered_feature_id: string; credit_cost: number }[]; +} + +/** + * Centralizes the AI-vs-classic credit system field selection shared by the + * feature mutation sheets. AI credit systems carry markup fields and omit the + * credit schema; classic credit systems do the inverse. + */ +export const buildFeatureMarkupParams = ({ + type, + modelMarkups, + defaultMarkup, + providerMarkups, + schema, +}: BuildFeatureMarkupParamsArgs): FeatureMarkupParams => { + const ai = isAiCreditSystem(type); + return { + model_markups: ai ? modelMarkups : undefined, + default_markup: ai ? defaultMarkup : undefined, + provider_markups: ai ? providerMarkups : undefined, + credit_schema: ai + ? undefined + : schema?.map((item) => ({ + metered_feature_id: item.metered_feature_id, + credit_cost: + item.credit_amount != null ? Number(item.credit_amount) : 0, + })), + }; +}; diff --git a/vite/src/views/products/features/utils/getFeatureIcon.tsx b/vite/src/views/products/features/utils/getFeatureIcon.tsx index 6b5ce42a6..f1d958cde 100644 --- a/vite/src/views/products/features/utils/getFeatureIcon.tsx +++ b/vite/src/views/products/features/utils/getFeatureIcon.tsx @@ -2,6 +2,7 @@ import type { Feature, ProductItem } from "@autumn/shared"; import { FeatureType, FeatureUsageType, + isAiCreditSystem, ProductItemFeatureType, } from "@autumn/shared"; import { @@ -70,10 +71,7 @@ export const getFeatureIconConfig = ( } // Handle AI credit system - if ( - typeStr === FeatureType.AiCreditSystem || - typeStr === "ai_credit_system" - ) { + if (isAiCreditSystem(typeStr) || typeStr === "ai_credit_system") { return { icon: , color: "text-yellow-500", diff --git a/vite/src/views/products/plan/components/edit-plan-feature/BillingType.tsx b/vite/src/views/products/plan/components/edit-plan-feature/BillingType.tsx index 165307e2e..62ffb5d49 100644 --- a/vite/src/views/products/plan/components/edit-plan-feature/BillingType.tsx +++ b/vite/src/views/products/plan/components/edit-plan-feature/BillingType.tsx @@ -1,9 +1,9 @@ import { BillingInterval, - FeatureType, FeatureUsageType, getFeatureName, Infinite, + isAiCreditSystem, isContUseItem, isFeaturePriceItem, ProductItemInterval, @@ -74,7 +74,6 @@ export function BillingType() { }; const feature = features.find((f) => f.id === item.feature_id); - const isAiCreditSystem = feature?.type === FeatureType.AiCreditSystem; const featureName = getFeatureName({ feature, @@ -106,7 +105,7 @@ export function BillingType() {
Included
- {isAiCreditSystem + {isAiCreditSystem(feature?.type) ? "Set an included USD budget (eg, $10 per month)." : isConsumable ? `Set an included usage limit (eg, 100 ${featureName} per month).` @@ -128,7 +127,7 @@ export function BillingType() {
Priced
- {isAiCreditSystem + {isAiCreditSystem(feature?.type) ? "Bill model usage at the markup you set in USD after the included budget is used." : isConsumable ? `Charge a price for usage (eg, $0.05 per ${singleFeatureName}).` diff --git a/vite/src/views/products/plan/components/edit-plan-feature/EditPlanFeatureSheet.tsx b/vite/src/views/products/plan/components/edit-plan-feature/EditPlanFeatureSheet.tsx index c6563308d..98f901f99 100644 --- a/vite/src/views/products/plan/components/edit-plan-feature/EditPlanFeatureSheet.tsx +++ b/vite/src/views/products/plan/components/edit-plan-feature/EditPlanFeatureSheet.tsx @@ -1,4 +1,4 @@ -import { FeatureType, TierBehavior } from "@autumn/shared"; +import { FeatureType, isAiCreditSystem, TierBehavior } from "@autumn/shared"; import { PencilSimpleIcon } from "@phosphor-icons/react"; import { useState } from "react"; import { IconButton } from "@/components/v2/buttons/IconButton"; @@ -116,7 +116,6 @@ export function EditPlanFeatureSheet({ const feature = getFeature(item?.feature_id ?? "", features); const isFeaturePrice = isFeaturePriceItem(item); - const isAiCreditSystem = feature?.type === FeatureType.AiCreditSystem; // Allow confirming a priced feature that has a $0 tier (valid zero-price config) const isZeroPriceItem = @@ -167,7 +166,7 @@ export function EditPlanFeatureSheet({ - {isFeaturePrice && !isAiCreditSystem && ( + {isFeaturePrice && !isAiCreditSystem(feature?.type) && ( 1 ? ( diff --git a/vite/src/views/products/plan/components/edit-plan-feature/IncludedUsage.tsx b/vite/src/views/products/plan/components/edit-plan-feature/IncludedUsage.tsx index 57ed3b92e..f310287bf 100644 --- a/vite/src/views/products/plan/components/edit-plan-feature/IncludedUsage.tsx +++ b/vite/src/views/products/plan/components/edit-plan-feature/IncludedUsage.tsx @@ -13,9 +13,9 @@ import { billingToItemInterval, EntInterval, entToItemInterval, - FeatureType, getFeatureName, Infinite, + isAiCreditSystem, isContUseItem, } from "@autumn/shared"; import { InfinityIcon } from "@phosphor-icons/react"; @@ -31,7 +31,6 @@ export function IncludedUsage() { const isFeaturePrice = isFeaturePriceItem(item); const feature = features.find((f) => f.id === item.feature_id); - const isAiCreditSystem = feature?.type === FeatureType.AiCreditSystem; // Helper function to get the display value for the input const getInputValue = () => { @@ -49,13 +48,13 @@ export function IncludedUsage() {
- {isAiCreditSystem + {isAiCreditSystem(feature?.type) ? `USD budget ${isFeaturePrice ? "granted before billing" : "allocated to this plan"}` : <>Quantity of {getFeatureName({ feature, plural: true })}{isFeaturePrice ? " granted before billing" : " that can be used"} }
- {isAiCreditSystem ? ( + {isAiCreditSystem(feature?.type) ? ( $ { - const isAi = values.type === FeatureType.AiCreditSystem; + const isAi = isAiCreditSystem(values.type); setFeature({ ...feature, diff --git a/vite/src/views/products/plan/components/plan-card/DummyPlanFeatureRow.tsx b/vite/src/views/products/plan/components/plan-card/DummyPlanFeatureRow.tsx index 7bc9cb3ae..78e184def 100644 --- a/vite/src/views/products/plan/components/plan-card/DummyPlanFeatureRow.tsx +++ b/vite/src/views/products/plan/components/plan-card/DummyPlanFeatureRow.tsx @@ -1,4 +1,4 @@ -import { FeatureType, FeatureUsageType } from "@autumn/shared"; +import { FeatureType, FeatureUsageType, isAiCreditSystem } from "@autumn/shared"; import { BoxArrowDownIcon } from "@phosphor-icons/react"; import { useFeatureStore } from "@/hooks/stores/useFeatureStore"; import { cn } from "@/lib/utils"; @@ -44,13 +44,11 @@ export const DummyPlanFeatureRow = () => { return "Chat Messages"; }; - const isAiCreditSystem = feature.type === FeatureType.AiCreditSystem; - // Build display text based on feature type const getDisplayText = () => { const name = hasName ? featureName : getPlaceholderName(); - if (isAiCreditSystem) { + if (isAiCreditSystem(feature.type)) { return { primary: `$10.00 of ${name}`, secondary: "" }; } diff --git a/vite/src/views/products/plan/components/plan-card/PlanFeatureRow.tsx b/vite/src/views/products/plan/components/plan-card/PlanFeatureRow.tsx index cf6144f1e..e3de7e631 100644 --- a/vite/src/views/products/plan/components/plan-card/PlanFeatureRow.tsx +++ b/vite/src/views/products/plan/components/plan-card/PlanFeatureRow.tsx @@ -1,6 +1,6 @@ /** biome-ignore-all lint/a11y/noStaticElementInteractions: needed */ /** biome-ignore-all lint/a11y/useSemanticElements: needed */ -import { type ProductItem, FeatureType } from "@autumn/shared"; +import { type ProductItem, isAiCreditSystem } from "@autumn/shared"; import { getProductItemDisplay } from "@autumn/shared"; import { TrashIcon } from "@phosphor-icons/react"; import { useEffect, useRef, useState } from "react"; @@ -68,7 +68,6 @@ export const PlanFeatureRow = ({ const feature = features.find((f) => f.id === item.feature_id); const hasFeatureName = feature?.name && feature.name.trim() !== ""; - const isAiCreditSystem = feature?.type === FeatureType.AiCreditSystem; const displayText = hasFeatureName ? display.primary_text @@ -207,7 +206,7 @@ export const PlanFeatureRow = ({ {displayText} - {!isAiCreditSystem && display.secondary_text && ( + {!isAiCreditSystem(feature?.type) && display.secondary_text && ( {display.secondary_text} )}

From 22f1260bf7ed1c2d8e88073a1aeb3deb6f8bb646 Mon Sep 17 00:00:00 2001 From: Charlie Lamb Date: Tue, 9 Jun 2026 18:46:20 +0100 Subject: [PATCH 30/46] chore: minor ai credit system cleanup --- .gitignore | 1 + .../external-providers/ai-sdk.mdx | 12 +- bun.lock | 3 + packages/ai-sdk/package.json | 1 + packages/ai-sdk/src/index.ts | 3 +- packages/ai-sdk/tests/unit/index.test.ts | 144 ++++++++++++++++++ .../src/lib/transforms/sdkToApi/feature.ts | 2 +- packages/openapi/tsconfig.json | 2 +- server/src/init.ts | 2 - .../priceToStripePrepaidV2Tiers.ts | 33 ++-- .../components/AiCreditSchema.tsx | 73 +++++---- .../components/CreditSystemSchema.tsx | 56 +------ .../components/new-feature/NewFeatureType.tsx | 3 +- 13 files changed, 227 insertions(+), 108 deletions(-) create mode 100644 packages/ai-sdk/tests/unit/index.test.ts diff --git a/.gitignore b/.gitignore index 58a0b8b47..9e04e6a10 100644 --- a/.gitignore +++ b/.gitignore @@ -22,6 +22,7 @@ supabase.sh tests/ !server/tests !packages/mcp/tests +!packages/ai-sdk/tests !apps/leaf/tests !vite/tests .secrets diff --git a/apps/docs/mintlify/documentation/external-providers/ai-sdk.mdx b/apps/docs/mintlify/documentation/external-providers/ai-sdk.mdx index bb33e85f0..e9bb2b92a 100644 --- a/apps/docs/mintlify/documentation/external-providers/ai-sdk.mdx +++ b/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, diff --git a/bun.lock b/bun.lock index f8d001f5a..6d4d8dcdf 100644 --- a/bun.lock +++ b/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=="], diff --git a/packages/ai-sdk/package.json b/packages/ai-sdk/package.json index f8adf149d..e7f964477 100644 --- a/packages/ai-sdk/package.json +++ b/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" }, diff --git a/packages/ai-sdk/src/index.ts b/packages/ai-sdk/src/index.ts index ce2f5a9ff..e655582ff 100644 --- a/packages/ai-sdk/src/index.ts +++ b/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, diff --git a/packages/ai-sdk/tests/unit/index.test.ts b/packages/ai-sdk/tests/unit/index.test.ts new file mode 100644 index 000000000..fd28f5863 --- /dev/null +++ b/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; +}; + +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, + }, + ]); + }); +}); diff --git a/packages/atmn/src/lib/transforms/sdkToApi/feature.ts b/packages/atmn/src/lib/transforms/sdkToApi/feature.ts index a67702701..2c53db6f5 100644 --- a/packages/atmn/src/lib/transforms/sdkToApi/feature.ts +++ b/packages/atmn/src/lib/transforms/sdkToApi/feature.ts @@ -12,7 +12,7 @@ export interface ApiFeatureParams { credit_cost: number; }>; model_markups?: Record; diff --git a/packages/openapi/tsconfig.json b/packages/openapi/tsconfig.json index 6d08724fc..fcbd4453d 100644 --- a/packages/openapi/tsconfig.json +++ b/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"] } diff --git a/server/src/init.ts b/server/src/init.ts index 707bd94f9..254807594 100644 --- a/server/src/init.ts +++ b/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()]); diff --git a/shared/utils/productUtils/priceUtils/convertPrice/priceToStripePrepaidV2Tiers.ts b/shared/utils/productUtils/priceUtils/convertPrice/priceToStripePrepaidV2Tiers.ts index 81a9e2d3f..a58047276 100644 --- a/shared/utils/productUtils/priceUtils/convertPrice/priceToStripePrepaidV2Tiers.ts +++ b/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; }; diff --git a/vite/src/views/products/features/credit-systems/components/AiCreditSchema.tsx b/vite/src/views/products/features/credit-systems/components/AiCreditSchema.tsx index edaf869d8..b2cd1d3ef 100644 --- a/vite/src/views/products/features/credit-systems/components/AiCreditSchema.tsx +++ b/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) { />
-
- {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 && ( +
+ {activeProviderKeys.map((providerKey) => { + const provider = providers[providerKey]; + const modelFullIds = providerGroups[providerKey] ?? []; + const providerName = + provider?.name ?? + providerKey.charAt(0).toUpperCase() + providerKey.slice(1); - return ( - - ); - })} -
+ return ( + + ); + })} +
+ )}
e.stopPropagation()} > - Add Provider + + Add Provider Override + + + + + + Add specific markup overrides for certain providers/models. + + + { - if (!releaseDate) return -1; - const timestamp = Date.parse(releaseDate); - return Number.isNaN(timestamp) ? -1 : timestamp; -}; - -function getDefaultModelMarkups( - providers: Record, -): Record { - const result: Record = {}; - 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", {}); } }; diff --git a/vite/src/views/products/plan/components/new-feature/NewFeatureType.tsx b/vite/src/views/products/plan/components/new-feature/NewFeatureType.tsx index 3e06d28c3..f046fe991 100644 --- a/vite/src/views/products/plan/components/new-feature/NewFeatureType.tsx +++ b/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({
{ setFeature({ ...feature, From 2c43fce1d17bbf3b7a6e71891f1dd84fe8034249 Mon Sep 17 00:00:00 2001 From: Charlie Lamb Date: Wed, 10 Jun 2026 11:11:10 +0100 Subject: [PATCH 31/46] chore: increase test coverage --- .../track/basic/track-tokens-limits.test.ts | 254 ++++++++++++++++++ .../track/basic/track-tokens-paid.test.ts | 202 ++++++++++++++ .../basic/track-tokens-resolution.test.ts | 151 +++++++++++ .../balances/track/basic/track-tokens.test.ts | 194 ++++++++++++- .../unit/features/get-model-pricing.test.ts | 115 ++++++++ server/tests/utils/fixtures/items.ts | 14 + 6 files changed, 929 insertions(+), 1 deletion(-) create mode 100644 server/tests/integration/balances/track/basic/track-tokens-limits.test.ts create mode 100644 server/tests/integration/balances/track/basic/track-tokens-paid.test.ts create mode 100644 server/tests/integration/balances/track/basic/track-tokens-resolution.test.ts create mode 100644 server/tests/unit/features/get-model-pricing.test.ts diff --git a/server/tests/integration/balances/track/basic/track-tokens-limits.test.ts b/server/tests/integration/balances/track/basic/track-tokens-limits.test.ts new file mode 100644 index 000000000..35afcd6dd --- /dev/null +++ b/server/tests/integration/balances/track/basic/track-tokens-limits.test.ts @@ -0,0 +1,254 @@ +import { expect, test } from "bun:test"; + +import type { ApiCustomerV5, TrackResponseV3 } from "@autumn/shared"; +import { ErrCode } from "@autumn/shared"; +import { expectBalanceCorrect } from "@tests/integration/utils/expectBalanceCorrect.js"; +import { TestFeature } from "@tests/setup/v2Features.js"; +import { expectAutumnError } from "@tests/utils/expectUtils/expectErrUtils.js"; +import { items } from "@tests/utils/fixtures/items.js"; +import { products } from "@tests/utils/fixtures/products.js"; +import { initScenario, s } from "@tests/utils/testInitUtils/initScenario.js"; +import chalk from "chalk"; + +// custom/internal-model: input_cost=5 $/M, output_cost=15 $/M, markup=0% +// in=5000/out=2500 -> 0.0625; in=10000/out=5000 -> 0.125 + +// ═══════════════════════════════════════════════════════════════════ +// TRACK-TOKENS-LIM-1: default behavior caps deduction at zero balance +// ═══════════════════════════════════════════════════════════════════ + +test.concurrent( + `${chalk.yellowBright("track-tokens-lim-1: default behavior caps token deduction at zero balance")}`, + async () => { + const aiCreditsItem = items.free({ + featureId: TestFeature.AiCredits, + includedUsage: 0.1, + }); + const freeProd = products.base({ id: "free", items: [aiCreditsItem] }); + + const { customerId, autumnV2_2 } = await initScenario({ + customerId: "track-tokens-lim-1", + setup: [ + s.customer({ testClock: false }), + s.products({ list: [freeProd] }), + ], + actions: [s.attach({ productId: freeProd.id })], + }); + + // First track: cost 0.0625 fits within the 0.1 balance + const trackRes1: TrackResponseV3 = await autumnV2_2.post("/track_tokens", { + customer_id: customerId, + feature_id: TestFeature.AiCredits, + model_id: "custom/internal-model", + input_tokens: 5000, + output_tokens: 2500, + }); + expect(trackRes1.value).toBeCloseTo(0.0625, 10); + + // Second track: cost 0.125 exceeds the remaining 0.0375 — capped at zero + await autumnV2_2.post("/track_tokens", { + customer_id: customerId, + feature_id: TestFeature.AiCredits, + model_id: "custom/internal-model", + input_tokens: 10000, + output_tokens: 5000, + }); + + const customer = await autumnV2_2.customers.get(customerId); + expectBalanceCorrect({ + customer, + featureId: TestFeature.AiCredits, + remaining: 0, + usage: 0.1, + }); + }, +); + +// ═══════════════════════════════════════════════════════════════════ +// TRACK-TOKENS-LIM-2: overage_behavior "reject" errors, balance intact +// ═══════════════════════════════════════════════════════════════════ + +test.concurrent( + `${chalk.yellowBright("track-tokens-lim-2: overage_behavior reject errors with InsufficientBalance")}`, + async () => { + const aiCreditsItem = items.free({ + featureId: TestFeature.AiCredits, + includedUsage: 0.1, + }); + const freeProd = products.base({ id: "free", items: [aiCreditsItem] }); + + const { customerId, autumnV2_2 } = await initScenario({ + customerId: "track-tokens-lim-2", + setup: [ + s.customer({ testClock: false }), + s.products({ list: [freeProd] }), + ], + actions: [s.attach({ productId: freeProd.id })], + }); + + await expectAutumnError({ + errCode: ErrCode.InsufficientBalance, + func: () => + autumnV2_2.post("/track_tokens", { + customer_id: customerId, + feature_id: TestFeature.AiCredits, + model_id: "custom/internal-model", + input_tokens: 10000, + output_tokens: 5000, + overage_behavior: "reject", + }), + }); + + const customer = await autumnV2_2.customers.get(customerId); + expectBalanceCorrect({ + customer, + featureId: TestFeature.AiCredits, + remaining: 0.1, + usage: 0, + }); + }, +); + +// ═══════════════════════════════════════════════════════════════════ +// TRACK-TOKENS-LIM-3: explicit overage_behavior "cap" deducts up to zero +// ═══════════════════════════════════════════════════════════════════ + +test.concurrent( + `${chalk.yellowBright("track-tokens-lim-3: explicit overage_behavior cap deducts up to zero")}`, + async () => { + const aiCreditsItem = items.free({ + featureId: TestFeature.AiCredits, + includedUsage: 0.1, + }); + const freeProd = products.base({ id: "free", items: [aiCreditsItem] }); + + const { customerId, autumnV2_2 } = await initScenario({ + customerId: "track-tokens-lim-3", + setup: [ + s.customer({ testClock: false }), + s.products({ list: [freeProd] }), + ], + actions: [s.attach({ productId: freeProd.id })], + }); + + await autumnV2_2.post("/track_tokens", { + customer_id: customerId, + feature_id: TestFeature.AiCredits, + model_id: "custom/internal-model", + input_tokens: 10000, + output_tokens: 5000, + overage_behavior: "cap", + }); + + const customer = await autumnV2_2.customers.get(customerId); + expectBalanceCorrect({ + customer, + featureId: TestFeature.AiCredits, + remaining: 0, + usage: 0.1, + }); + }, +); + +// ═══════════════════════════════════════════════════════════════════ +// TRACK-TOKENS-LIM-4: unlimited balance never rejects or deducts +// ═══════════════════════════════════════════════════════════════════ + +test.concurrent( + `${chalk.yellowBright("track-tokens-lim-4: unlimited AI credit balance never rejects or deducts")}`, + async () => { + const aiCreditsItem = items.unlimited({ featureId: TestFeature.AiCredits }); + const freeProd = products.base({ id: "free", items: [aiCreditsItem] }); + + const { customerId, autumnV2_2 } = await initScenario({ + customerId: "track-tokens-lim-4", + setup: [ + s.customer({ testClock: false }), + s.products({ list: [freeProd] }), + ], + actions: [s.attach({ productId: freeProd.id })], + }); + + const trackRes: TrackResponseV3 = await autumnV2_2.post("/track_tokens", { + customer_id: customerId, + feature_id: TestFeature.AiCredits, + model_id: "custom/internal-model", + input_tokens: 10000, + output_tokens: 5000, + }); + + expect(trackRes.value).toBeCloseTo(0.125, 10); + expect(trackRes.balance).toMatchObject({ + feature_id: TestFeature.AiCredits, + unlimited: true, + usage: 0, + }); + + // Second track: still no deduction, never rejected + await autumnV2_2.post("/track_tokens", { + customer_id: customerId, + feature_id: TestFeature.AiCredits, + model_id: "custom/internal-model", + input_tokens: 20000, + output_tokens: 10000, + }); + + const customer = await autumnV2_2.customers.get(customerId); + expect(customer.balances[TestFeature.AiCredits]).toMatchObject({ + unlimited: true, + usage: 0, + }); + }, +); + +// ═══════════════════════════════════════════════════════════════════ +// TRACK-TOKENS-LIM-5: duplicate idempotency_key rejected, deducts once +// ═══════════════════════════════════════════════════════════════════ + +test.concurrent( + `${chalk.yellowBright("track-tokens-lim-5: duplicate idempotency_key rejected, deducts once")}`, + async () => { + const aiCreditsItem = items.free({ + featureId: TestFeature.AiCredits, + includedUsage: 1000, + }); + const freeProd = products.base({ id: "free", items: [aiCreditsItem] }); + + const { customerId, autumnV2_2 } = await initScenario({ + customerId: "track-tokens-lim-5", + setup: [ + s.customer({ testClock: false }), + s.products({ list: [freeProd] }), + ], + actions: [s.attach({ productId: freeProd.id })], + }); + + const body = { + customer_id: customerId, + feature_id: TestFeature.AiCredits, + model_id: "custom/internal-model", + input_tokens: 10000, + output_tokens: 5000, + idempotency_key: `track-tokens-idem-${Date.now().toString(36)}`, + }; + + const trackRes: TrackResponseV3 = await autumnV2_2.post( + "/track_tokens", + body, + ); + expect(trackRes.value).toBeCloseTo(0.125, 10); + + await expectAutumnError({ + errCode: ErrCode.DuplicateIdempotencyKey, + func: () => autumnV2_2.post("/track_tokens", body), + }); + + const customer = await autumnV2_2.customers.get(customerId); + expectBalanceCorrect({ + customer, + featureId: TestFeature.AiCredits, + remaining: 999.875, + usage: 0.125, + }); + }, +); diff --git a/server/tests/integration/balances/track/basic/track-tokens-paid.test.ts b/server/tests/integration/balances/track/basic/track-tokens-paid.test.ts new file mode 100644 index 000000000..ec1240354 --- /dev/null +++ b/server/tests/integration/balances/track/basic/track-tokens-paid.test.ts @@ -0,0 +1,202 @@ +import { expect, test } from "bun:test"; + +import type { + ApiCustomerV3, + ApiCustomerV5, + TrackResponseV3, +} from "@autumn/shared"; +import { ErrCode } from "@autumn/shared"; +import { expectCustomerInvoiceCorrect } from "@tests/integration/billing/utils/expectCustomerInvoiceCorrect.js"; +import { expectBalanceCorrect } from "@tests/integration/utils/expectBalanceCorrect.js"; +import { TestFeature } from "@tests/setup/v2Features.js"; +import { expectAutumnError } from "@tests/utils/expectUtils/expectErrUtils.js"; +import { items } from "@tests/utils/fixtures/items.js"; +import { products } from "@tests/utils/fixtures/products.js"; +import { timeout } from "@tests/utils/genUtils.js"; +import { advanceToNextInvoice } from "@tests/utils/testAttachUtils/testAttachUtils.js"; +import ctx from "@tests/utils/testInitUtils/createTestContext.js"; +import { initScenario, s } from "@tests/utils/testInitUtils/initScenario.js"; +import chalk from "chalk"; + +// custom/internal-model: input_cost=5 $/M, output_cost=15 $/M, markup=0% + +// ═══════════════════════════════════════════════════════════════════ +// TRACK-TOKENS-PAID-1: prepaid AI credits deduct through purchased balance +// ═══════════════════════════════════════════════════════════════════ + +test.concurrent( + `${chalk.yellowBright("track-tokens-paid-1: prepaid AI credits deduct through purchased balance")}`, + async () => { + const prepaidItem = items.prepaid({ + featureId: TestFeature.AiCredits, + price: 1, + billingUnits: 1, + includedUsage: 2, + }); + const prepaidProduct = products.pro({ + id: "prepaid-ai", + items: [prepaidItem], + }); + + const { customerId, autumnV2_2 } = await initScenario({ + customerId: "track-tokens-paid-1", + setup: [ + s.customer({ testClock: false, paymentMethod: "success" }), + s.products({ list: [prepaidProduct] }), + ], + actions: [ + s.attach({ + productId: prepaidProduct.id, + options: [{ feature_id: TestFeature.AiCredits, quantity: 3 }], + }), + ], + }); + + // 2 included + 3 purchased = 5 + const customerBefore = + await autumnV2_2.customers.get(customerId); + expectBalanceCorrect({ + customer: customerBefore, + featureId: TestFeature.AiCredits, + granted: 5, + remaining: 5, + }); + + // (5*100000 + 15*100000) / 1e6 = $2.00 + const trackRes1: TrackResponseV3 = await autumnV2_2.post("/track_tokens", { + customer_id: customerId, + feature_id: TestFeature.AiCredits, + model_id: "custom/internal-model", + input_tokens: 100000, + output_tokens: 100000, + }); + expect(trackRes1.value).toBeCloseTo(2, 10); + + // Cost $4 > remaining 3 with reject — errors, balance unchanged + await expectAutumnError({ + errCode: ErrCode.InsufficientBalance, + func: () => + autumnV2_2.post("/track_tokens", { + customer_id: customerId, + feature_id: TestFeature.AiCredits, + model_id: "custom/internal-model", + input_tokens: 200000, + output_tokens: 200000, + overage_behavior: "reject", + }), + }); + + const customerMid = + await autumnV2_2.customers.get(customerId); + expectBalanceCorrect({ + customer: customerMid, + featureId: TestFeature.AiCredits, + remaining: 3, + usage: 2, + }); + + // Cost $3.00 drains the remaining balance exactly + const trackRes2: TrackResponseV3 = await autumnV2_2.post("/track_tokens", { + customer_id: customerId, + feature_id: TestFeature.AiCredits, + model_id: "custom/internal-model", + input_tokens: 150000, + output_tokens: 150000, + }); + expect(trackRes2.value).toBeCloseTo(3, 10); + + // Cached vs DB agreement (mutation-log sync is async) + await timeout(6000); + const customerNonCached = await autumnV2_2.customers.get( + customerId, + { skip_cache: "true" }, + ); + expectBalanceCorrect({ + customer: customerNonCached, + featureId: TestFeature.AiCredits, + granted: 5, + remaining: 0, + usage: 5, + }); + }, +); + +// ═══════════════════════════════════════════════════════════════════ +// TRACK-TOKENS-PAID-2: consumable AI credit overage lands on the invoice +// ═══════════════════════════════════════════════════════════════════ + +test.concurrent( + `${chalk.yellowBright("track-tokens-paid-2: consumable AI credit overage lands on the renewal invoice")}`, + async () => { + const consumableItem = items.consumable({ + featureId: TestFeature.AiCredits, + includedUsage: 1, + price: 1, + billingUnits: 1, + }); + const proProduct = products.pro({ + id: "consumable-ai", + items: [consumableItem], + }); + + const { customerId, autumnV1, autumnV2_2, testClockId } = + await initScenario({ + customerId: "track-tokens-paid-2", + setup: [ + s.customer({ paymentMethod: "success" }), + s.products({ list: [proProduct] }), + ], + actions: [s.attach({ productId: proProduct.id })], + }); + + // (5*200000 + 15*200000) / 1e6 = $4.00 exactly + const trackRes: TrackResponseV3 = await autumnV2_2.post("/track_tokens", { + customer_id: customerId, + feature_id: TestFeature.AiCredits, + model_id: "custom/internal-model", + input_tokens: 200000, + output_tokens: 200000, + }); + expect(trackRes.value).toBeCloseTo(4, 10); + + const customerMid = + await autumnV2_2.customers.get(customerId); + expectBalanceCorrect({ + customer: customerMid, + featureId: TestFeature.AiCredits, + remaining: 0, + usage: 4, + }); + + await timeout(2000); + await advanceToNextInvoice({ + stripeCli: ctx.stripeCli, + testClockId: testClockId!, + withPause: true, + }); + + // Renewal invoice: $20 pro base + 3 overage units × $1 = $23. + // Invoice lands via Stripe webhook — poll briefly before asserting. + for (let attempt = 0; attempt < 5; attempt++) { + const customer = await autumnV1.customers.get(customerId); + if ((customer.invoices?.length ?? 0) >= 2) break; + await timeout(10000); + } + await expectCustomerInvoiceCorrect({ + customerId, + count: 2, + latestTotal: 23, + latestInvoiceProductId: proProduct.id, + }); + + // Balance resets for the new cycle + const customerReset = + await autumnV2_2.customers.get(customerId); + expectBalanceCorrect({ + customer: customerReset, + featureId: TestFeature.AiCredits, + remaining: 1, + usage: 0, + }); + }, +); diff --git a/server/tests/integration/balances/track/basic/track-tokens-resolution.test.ts b/server/tests/integration/balances/track/basic/track-tokens-resolution.test.ts new file mode 100644 index 000000000..691cf5944 --- /dev/null +++ b/server/tests/integration/balances/track/basic/track-tokens-resolution.test.ts @@ -0,0 +1,151 @@ +import { expect, test } from "bun:test"; + +import type { + ApiCustomerV5, + ApiEntityV2, + TrackResponseV3, +} from "@autumn/shared"; +import { ApiVersion, FeatureType } from "@autumn/shared"; +import { expectBalanceCorrect } from "@tests/integration/utils/expectBalanceCorrect.js"; +import { TestFeature } from "@tests/setup/v2Features.js"; +import { items } from "@tests/utils/fixtures/items.js"; +import { products } from "@tests/utils/fixtures/products.js"; +import { initScenario, s } from "@tests/utils/testInitUtils/initScenario.js"; +import chalk from "chalk"; +import { AutumnInt } from "@/external/autumn/autumnCli.js"; + +// custom/internal-model: input_cost=5 $/M, output_cost=15 $/M, markup=0% + +// ═══════════════════════════════════════════════════════════════════ +// TRACK-TOKENS-RES-1: entity_id deducts entity balance via auto-resolution +// ═══════════════════════════════════════════════════════════════════ + +test.concurrent( + `${chalk.yellowBright("track-tokens-res-1: entity_id deducts entity balance via auto-resolution")}`, + async () => { + const aiCreditsItem = items.free({ + featureId: TestFeature.AiCredits, + includedUsage: 100, + entityFeatureId: TestFeature.Users, + }); + const freeProd = products.base({ id: "free", items: [aiCreditsItem] }); + + const { customerId, autumnV2_2, entities } = await initScenario({ + customerId: "track-tokens-res-1", + setup: [ + s.customer({ testClock: false }), + s.products({ list: [freeProd] }), + s.entities({ count: 2, featureId: TestFeature.Users }), + ], + actions: [s.attach({ productId: freeProd.id })], + }); + + // No feature_id — exercises AI credit auto-resolution with entity scoping + const trackRes: TrackResponseV3 = await autumnV2_2.post("/track_tokens", { + customer_id: customerId, + entity_id: entities[0].id, + model_id: "custom/internal-model", + input_tokens: 10000, + output_tokens: 5000, + }); + + expect(trackRes.customer_id).toBe(customerId); + expect(trackRes.value).toBeCloseTo(0.125, 10); + + const entity0 = await autumnV2_2.entities.get( + customerId, + entities[0].id, + ); + expectBalanceCorrect({ + customer: entity0, + featureId: TestFeature.AiCredits, + remaining: 99.875, + usage: 0.125, + }); + + const entity1 = await autumnV2_2.entities.get( + customerId, + entities[1].id, + ); + expectBalanceCorrect({ + customer: entity1, + featureId: TestFeature.AiCredits, + remaining: 100, + usage: 0, + }); + }, +); + +// ═══════════════════════════════════════════════════════════════════ +// TRACK-TOKENS-RES-2: updated model markup applies to subsequent tracks +// ═══════════════════════════════════════════════════════════════════ + +test.concurrent( + `${chalk.yellowBright("track-tokens-res-2: updated model markup applies to subsequent tracks")}`, + async () => { + const autumn = new AutumnInt({ version: ApiVersion.V2_2 }); + + // Throwaway feature — never mutate the shared AiCredits fixtures + const featureId = `ai_credits_mut_${Date.now()}_${Math.random() + .toString(36) + .slice(2, 8)}`; + await autumn.post("/features.create", { + feature_id: featureId, + name: "AI Credits Mutable", + type: FeatureType.AiCreditSystem, + model_markups: { + "custom/mut-model": { markup: 0, input_cost: 10, output_cost: 20 }, + }, + }); + + const aiCreditsItem = items.free({ featureId, includedUsage: 1000 }); + const freeProd = products.base({ id: "free-mut", items: [aiCreditsItem] }); + + const { customerId, autumnV2_2 } = await initScenario({ + customerId: "track-tokens-res-2", + setup: [ + s.customer({ testClock: false }), + s.products({ list: [freeProd] }), + ], + actions: [s.attach({ productId: freeProd.id })], + }); + + const trackBody = { + customer_id: customerId, + feature_id: featureId, + model_id: "custom/mut-model", + input_tokens: 10000, + output_tokens: 5000, + }; + + // Markup 0 → base cost (10*10000 + 20*5000)/1e6 = 0.2 + const trackRes1: TrackResponseV3 = await autumnV2_2.post( + "/track_tokens", + trackBody, + ); + expect(trackRes1.value).toBeCloseTo(0.2, 10); + + // Bump the model markup to 100% + await autumn.post("/features.update", { + feature_id: featureId, + model_markups: { + "custom/mut-model": { markup: 100, input_cost: 10, output_cost: 20 }, + }, + }); + + // Explicit feature_id resolves from freshly loaded org features → 0.4 + const trackRes2: TrackResponseV3 = await autumnV2_2.post( + "/track_tokens", + trackBody, + ); + expect(trackRes2.value).toBeCloseTo(0.4, 10); + + const customer = await autumnV2_2.customers.get(customerId); + expectBalanceCorrect({ + customer, + featureId, + remaining: 999.4, + usage: 0.6, + }); + }, +); diff --git a/server/tests/integration/balances/track/basic/track-tokens.test.ts b/server/tests/integration/balances/track/basic/track-tokens.test.ts index d500f4e8a..6db39c9fa 100644 --- a/server/tests/integration/balances/track/basic/track-tokens.test.ts +++ b/server/tests/integration/balances/track/basic/track-tokens.test.ts @@ -1,7 +1,15 @@ import { expect, test } from "bun:test"; -import type { ApiCustomerV3, TrackResponseV2 } from "@autumn/shared"; +import type { + ApiCustomerV3, + ApiCustomerV5, + TrackResponseV2, + TrackResponseV3, +} from "@autumn/shared"; +import { ErrCode } from "@autumn/shared"; +import { expectBalanceCorrect } from "@tests/integration/utils/expectBalanceCorrect.js"; import { TestFeature } from "@tests/setup/v2Features.js"; +import { expectAutumnError } from "@tests/utils/expectUtils/expectErrUtils.js"; import { items } from "@tests/utils/fixtures/items.js"; import { products } from "@tests/utils/fixtures/products.js"; import { initScenario, s } from "@tests/utils/testInitUtils/initScenario.js"; @@ -390,3 +398,187 @@ test.concurrent( }); }, ); + +// ═══════════════════════════════════════════════════════════════════ +// TRACK-TOKENS-6: custom/* model without configured costs errors +// ═══════════════════════════════════════════════════════════════════ + +test.concurrent( + `${chalk.yellowBright("track-tokens-6: custom model missing input_cost/output_cost errors")}`, + async () => { + const aiCreditsItem = items.free({ + featureId: TestFeature.AiCredits, + includedUsage: 1000, + }); + const freeProd = products.base({ + id: "free", + items: [aiCreditsItem], + }); + + const { customerId, autumnV2_2 } = await initScenario({ + customerId: "track-tokens-6", + setup: [ + s.customer({ testClock: false }), + s.products({ list: [freeProd] }), + ], + actions: [s.attach({ productId: freeProd.id })], + }); + + await expectAutumnError({ + errCode: ErrCode.InvalidRequest, + errMessage: "missing input_cost or output_cost", + func: () => + autumnV2_2.post("/track_tokens", { + customer_id: customerId, + feature_id: TestFeature.AiCredits, + model_id: "custom/unconfigured-model", + input_tokens: 100, + output_tokens: 50, + }), + }); + + const customer = await autumnV2_2.customers.get(customerId); + expectBalanceCorrect({ + customer, + featureId: TestFeature.AiCredits, + remaining: 1000, + usage: 0, + }); + }, +); + +// ═══════════════════════════════════════════════════════════════════ +// TRACK-TOKENS-7: cache/audio/reasoning pools forwarded end-to-end +// ═══════════════════════════════════════════════════════════════════ + +test.concurrent( + `${chalk.yellowBright("track-tokens-7: cache/audio/reasoning token pools are billed end-to-end")}`, + async () => { + const aiCreditsItem = items.free({ + featureId: TestFeature.AiCredits, + includedUsage: 1000, + }); + const freeProd = products.base({ + id: "free", + items: [aiCreditsItem], + }); + + const { customerId, autumnV2_2, ctx } = await initScenario({ + customerId: "track-tokens-7", + setup: [ + s.customer({ testClock: false }), + s.products({ list: [freeProd] }), + ], + actions: [s.attach({ productId: freeProd.id })], + }); + + const aiCreditFeature = ctx.features.find( + (f) => f.id === TestFeature.AiCredits, + ); + if (!aiCreditFeature) { + throw new Error(`${TestFeature.AiCredits} feature not found`); + } + + // Total input (input + cache pools) stays far below the 200k tier threshold + const modelId = "anthropic/claude-sonnet-4-20250514"; + const pools = { + input: 10000, + output: 5000, + cacheRead: 20000, + cacheWrite: 8000, + audioInput: 1000, + audioOutput: 1000, + reasoning: 4000, + }; + + const expectedCost = await getCreditCost({ + featureId: aiCreditFeature.id, + creditSystem: aiCreditFeature, + modelName: modelId, + tokens: pools, + }); + + // Pools must increase the bill vs text-only — otherwise the assertion + // below couldn't tell whether the HTTP layer forwarded them at all. + const textOnlyCost = await getCreditCost({ + featureId: aiCreditFeature.id, + creditSystem: aiCreditFeature, + modelName: modelId, + tokens: { input: pools.input, output: pools.output }, + }); + expect(expectedCost).toBeGreaterThan(textOnlyCost); + + const trackRes: TrackResponseV3 = await autumnV2_2.post("/track_tokens", { + customer_id: customerId, + feature_id: TestFeature.AiCredits, + model_id: modelId, + input_tokens: pools.input, + output_tokens: pools.output, + cache_read_tokens: pools.cacheRead, + cache_write_tokens: pools.cacheWrite, + audio_input_tokens: pools.audioInput, + audio_output_tokens: pools.audioOutput, + reasoning_tokens: pools.reasoning, + }); + + expect(trackRes.value).toBeCloseTo(expectedCost, 10); + + const customer = await autumnV2_2.customers.get(customerId); + expectBalanceCorrect({ + customer, + featureId: TestFeature.AiCredits, + remaining: new Decimal(1000).minus(expectedCost).toNumber(), + usage: expectedCost, + }); + }, +); + +// ═══════════════════════════════════════════════════════════════════ +// TRACK-TOKENS-8: custom models bill input/output only +// ═══════════════════════════════════════════════════════════════════ + +test.concurrent( + `${chalk.yellowBright("track-tokens-8: custom models ignore cache/audio/reasoning pools")}`, + async () => { + const aiCreditsItem = items.free({ + featureId: TestFeature.AiCredits, + includedUsage: 1000, + }); + const freeProd = products.base({ + id: "free", + items: [aiCreditsItem], + }); + + const { customerId, autumnV2_2 } = await initScenario({ + customerId: "track-tokens-8", + setup: [ + s.customer({ testClock: false }), + s.products({ list: [freeProd] }), + ], + actions: [s.attach({ productId: freeProd.id })], + }); + + // custom/internal-model: input_cost=5 $/M, output_cost=15 $/M, markup=0% + // Pool tokens are dropped for custom models, so cost is text-only. + const expectedCost = new Decimal(5) + .mul(10000) + .add(new Decimal(15).mul(5000)) + .div(1_000_000) + .toNumber(); // 0.125 + + const trackRes: TrackResponseV3 = await autumnV2_2.post("/track_tokens", { + customer_id: customerId, + feature_id: TestFeature.AiCredits, + model_id: "custom/internal-model", + input_tokens: 10000, + output_tokens: 5000, + cache_read_tokens: 20000, + cache_write_tokens: 8000, + audio_input_tokens: 1000, + audio_output_tokens: 1000, + reasoning_tokens: 4000, + }); + + expect(trackRes.value).toBeCloseTo(expectedCost, 10); + }, +); diff --git a/server/tests/unit/features/get-model-pricing.test.ts b/server/tests/unit/features/get-model-pricing.test.ts new file mode 100644 index 000000000..67ac21dfe --- /dev/null +++ b/server/tests/unit/features/get-model-pricing.test.ts @@ -0,0 +1,115 @@ +import { afterAll, afterEach, expect, mock, test } from "bun:test"; +import { ErrCode } from "@autumn/shared"; + +// Map-backed CacheManager stub — getModelsDevPricing's cache key is shared +// with the dev server, so the real Redis must never be touched here. +const store = new Map(); +const setJsonCalls: { key: string; value: unknown; ttl?: number }[] = []; + +mock.module("@/utils/cacheUtils/CacheManager.js", () => ({ + CacheManager: { + getJson: async (key: string) => store.get(key) ?? null, + setJson: async (key: string, value: unknown, ttl?: number) => { + setJsonCalls.push({ key, value, ttl }); + store.set(key, value); + }, + }, +})); + +const { getModelsDevPricing } = await import( + "@/internal/features/utils/getModelPricing.js" +); + +const PRIMARY_KEY = "models_dev_pricing"; +const STALE_KEY = "models_dev_pricing_stale"; + +const pricingData = { + anthropic: { id: "anthropic", name: "Anthropic", models: {} }, +}; +const stalePricingData = { + openai: { id: "openai", name: "OpenAI", models: {} }, +}; + +const realFetch = globalThis.fetch; +let fetchCalls = 0; + +const stubFetch = (impl: () => Promise) => { + globalThis.fetch = Object.assign( + async () => { + fetchCalls++; + return await impl(); + }, + { preconnect: realFetch.preconnect }, + ); +}; + +afterEach(() => { + store.clear(); + setJsonCalls.length = 0; + fetchCalls = 0; + globalThis.fetch = realFetch; +}); + +afterAll(() => { + mock.restore(); + globalThis.fetch = realFetch; +}); + +test("primary cache hit returns cached data without fetching", async () => { + store.set(PRIMARY_KEY, pricingData); + stubFetch(() => { + throw new Error("should not fetch"); + }); + + const result = await getModelsDevPricing(); + + expect(result).toEqual(pricingData); + expect(fetchCalls).toBe(0); +}); + +test("cache miss fetches and populates primary + stale caches", async () => { + stubFetch(async () => Response.json(pricingData)); + + const result = await getModelsDevPricing(); + + expect(result).toEqual(pricingData); + expect(fetchCalls).toBe(1); + + // Cache writes are fire-and-forget — flush microtasks before asserting + await Bun.sleep(0); + expect(setJsonCalls).toEqual([ + { key: PRIMARY_KEY, value: pricingData, ttl: 60 * 60 * 3 }, + { key: STALE_KEY, value: pricingData, ttl: 60 * 60 * 24 * 3 }, + ]); +}); + +test("non-ok response falls back to the stale cache", async () => { + store.set(STALE_KEY, stalePricingData); + stubFetch(async () => new Response("oops", { status: 500 })); + + const result = await getModelsDevPricing(); + + expect(result).toEqual(stalePricingData); +}); + +test("fetch network error falls back to the stale cache", async () => { + store.set(STALE_KEY, stalePricingData); + stubFetch(() => { + throw new Error("network down"); + }); + + const result = await getModelsDevPricing(); + + expect(result).toEqual(stalePricingData); +}); + +test("fetch failure with no stale cache throws InternalError", async () => { + stubFetch(() => { + throw new Error("network down"); + }); + + await expect(getModelsDevPricing()).rejects.toMatchObject({ + code: ErrCode.InternalError, + message: "Failed to fetch models.dev pricing and no cache available", + }); +}); diff --git a/server/tests/utils/fixtures/items.ts b/server/tests/utils/fixtures/items.ts index 39e0b7ca7..131504a13 100644 --- a/server/tests/utils/fixtures/items.ts +++ b/server/tests/utils/fixtures/items.ts @@ -51,13 +51,16 @@ const adminRights = () => const free = ({ featureId, includedUsage = 100, + entityFeatureId, }: { featureId: string; includedUsage?: number; + entityFeatureId?: string; }): LimitedItem => constructFeatureItem({ featureId, includedUsage, + entityFeatureId, }) as LimitedItem; /** @@ -163,6 +166,16 @@ const monthlyCredits = ({ rolloverConfig, }) as LimitedItem; +/** + * Generic unlimited feature - no usage cap + * @param featureId - Feature ID + */ +const unlimited = ({ featureId }: { featureId: string }) => + constructFeatureItem({ + featureId, + unlimited: true, + }); + /** * Unlimited messages - no usage cap * @returns Unlimited messages feature item @@ -783,6 +796,7 @@ export const items = { freeUsers, freeAllocatedUsers, freeAllocatedWorkflows, + unlimited, unlimitedMessages, weeklyMessages, lifetimeMessages, From f25838d21c2aef8479da5e62f02cd21146b47f9e Mon Sep 17 00:00:00 2001 From: Charlie Lamb Date: Wed, 10 Jun 2026 13:05:45 +0100 Subject: [PATCH 32/46] chore: improve types and atmn --- packages/atmn/src/commands/push/push.ts | 26 ++++ .../src/lib/transforms/apiToSdk/feature.ts | 2 + .../src/lib/transforms/sdkToApi/feature.ts | 32 +++-- .../src/lib/transforms/sdkToCode/feature.ts | 14 +- .../actions/generateAndUpdateAgentRules.ts | 8 +- .../track/utils/buildAiCreditCostProperty.ts | 2 +- .../track/utils/getTokenTrackParams.ts | 12 +- .../utils/deduction/computeCreditCosts.ts | 26 ++-- .../balances/utils/types/featureDeduction.ts | 11 +- .../internal/features/creditSystemUtils.ts | 23 ++-- .../features/utils/getModelPricing.ts | 6 +- .../track/basic/track-tokens-orbs.test.ts | 73 +++++++++-- .../track/basic/track-tokens-replay.test.ts | 120 ++++++++++++++++++ .../generate-and-update-agent-rules.test.ts | 82 ++++++++++++ .../track/handle-track-tokens.test.ts | 9 +- .../unit/features/get-credit-cost.test.ts | 70 ++++++++++ .../unit/features/get-model-pricing.test.ts | 17 +++ 17 files changed, 467 insertions(+), 66 deletions(-) create mode 100644 server/tests/integration/balances/track/basic/track-tokens-replay.test.ts create mode 100644 server/tests/unit/agent/generate-and-update-agent-rules.test.ts create mode 100644 server/tests/unit/features/get-credit-cost.test.ts diff --git a/packages/atmn/src/commands/push/push.ts b/packages/atmn/src/commands/push/push.ts index fe3480669..a3384ce97 100644 --- a/packages/atmn/src/commands/push/push.ts +++ b/packages/atmn/src/commands/push/push.ts @@ -349,6 +349,32 @@ function normalizeFeatureForCompare(f: Feature): Record { })); } + if (f.type === "ai_credit_system") { + const ai = f as Extract; + if (ai.modelMarkups && Object.keys(ai.modelMarkups).length > 0) { + result.modelMarkups = Object.fromEntries( + Object.entries(ai.modelMarkups) + .sort(([a], [b]) => a.localeCompare(b)) + .map(([modelId, entry]) => [ + modelId, + { + markup: entry.markup, + inputCost: entry.inputCost, + outputCost: entry.outputCost, + }, + ]), + ); + } + if (ai.defaultMarkup != null) result.defaultMarkup = ai.defaultMarkup; + if (ai.providerMarkups && Object.keys(ai.providerMarkups).length > 0) { + result.providerMarkups = Object.fromEntries( + Object.entries(ai.providerMarkups).sort(([a], [b]) => + a.localeCompare(b), + ), + ); + } + } + return result; } diff --git a/packages/atmn/src/lib/transforms/apiToSdk/feature.ts b/packages/atmn/src/lib/transforms/apiToSdk/feature.ts index e4ba94ffb..793025736 100644 --- a/packages/atmn/src/lib/transforms/apiToSdk/feature.ts +++ b/packages/atmn/src/lib/transforms/apiToSdk/feature.ts @@ -60,6 +60,8 @@ export const featureTransformer = createTransformer({ ...BASE_COMPUTE, type: () => "ai_credit_system" as const, modelMarkups: (api) => mapModelMarkups(api), + defaultMarkup: (api) => api.default_markup ?? undefined, + providerMarkups: (api) => api.provider_markups ?? undefined, }, }, diff --git a/packages/atmn/src/lib/transforms/sdkToApi/feature.ts b/packages/atmn/src/lib/transforms/sdkToApi/feature.ts index 2c53db6f5..d963639e3 100644 --- a/packages/atmn/src/lib/transforms/sdkToApi/feature.ts +++ b/packages/atmn/src/lib/transforms/sdkToApi/feature.ts @@ -16,6 +16,8 @@ export interface ApiFeatureParams { input_cost?: number; output_cost?: number; }>; + default_markup?: number; + provider_markups?: Record; } export function transformFeatureToApi(feature: Feature): ApiFeatureParams { @@ -44,17 +46,25 @@ export function transformFeatureToApi(feature: Feature): ApiFeatureParams { })); } - if (feature.type === "ai_credit_system" && feature.modelMarkups) { - base.model_markups = Object.fromEntries( - Object.entries(feature.modelMarkups).map(([modelId, entry]) => [ - modelId, - { - markup: entry.markup, - input_cost: entry.inputCost, - output_cost: entry.outputCost, - }, - ]) - ); + if (feature.type === "ai_credit_system") { + if (feature.modelMarkups) { + base.model_markups = Object.fromEntries( + Object.entries(feature.modelMarkups).map(([modelId, entry]) => [ + modelId, + { + markup: entry.markup, + input_cost: entry.inputCost, + output_cost: entry.outputCost, + }, + ]) + ); + } + if (feature.defaultMarkup !== undefined) { + base.default_markup = feature.defaultMarkup; + } + if (feature.providerMarkups) { + base.provider_markups = feature.providerMarkups; + } } return base; diff --git a/packages/atmn/src/lib/transforms/sdkToCode/feature.ts b/packages/atmn/src/lib/transforms/sdkToCode/feature.ts index 54a728f17..eddbd30b1 100644 --- a/packages/atmn/src/lib/transforms/sdkToCode/feature.ts +++ b/packages/atmn/src/lib/transforms/sdkToCode/feature.ts @@ -42,9 +42,17 @@ export function buildFeatureCode(feature: Feature, varNameOverride?: string): st lines.push(`\tcreditSchema: ${formatValue(feature.creditSchema)},`); } - // Add modelMarkups for ai_credit_system features - if (feature.type === "ai_credit_system" && feature.modelMarkups) { - lines.push(`\tmodelMarkups: ${formatValue(feature.modelMarkups)},`); + // Add markup config for ai_credit_system features + if (feature.type === "ai_credit_system") { + if (feature.modelMarkups) { + lines.push(`\tmodelMarkups: ${formatValue(feature.modelMarkups)},`); + } + if (feature.defaultMarkup !== undefined) { + lines.push(`\tdefaultMarkup: ${feature.defaultMarkup},`); + } + if (feature.providerMarkups) { + lines.push(`\tproviderMarkups: ${formatValue(feature.providerMarkups)},`); + } } lines.push(`});`); diff --git a/server/src/internal/agent/rules/actions/generateAndUpdateAgentRules.ts b/server/src/internal/agent/rules/actions/generateAndUpdateAgentRules.ts index 34cd35d26..f3acc2ee4 100644 --- a/server/src/internal/agent/rules/actions/generateAndUpdateAgentRules.ts +++ b/server/src/internal/agent/rules/actions/generateAndUpdateAgentRules.ts @@ -11,13 +11,17 @@ export const generateAndUpdateAgentRules = async ({ endTime?: string; startTime?: string; }) => { - const generated = await generateAgentRules({ ctx, endTime, startTime }); + const [generated, existing] = await Promise.all([ + generateAgentRules({ ctx, endTime, startTime }), + agentRulesRepo.get({ db: ctx.db, orgId: ctx.org.id }), + ]); + // Generation only derives entity/credit rules; never overwrite user-written notes. const rules = await agentRulesRepo.upsert({ db: ctx.db, metadata: generated.metadata, orgId: ctx.org.id, orgSlug: ctx.org.slug, - rules: generated.rules, + rules: { ...generated.rules, notes: existing.notes }, }); return { diff --git a/server/src/internal/balances/track/utils/buildAiCreditCostProperty.ts b/server/src/internal/balances/track/utils/buildAiCreditCostProperty.ts index 66156e522..673225f5d 100644 --- a/server/src/internal/balances/track/utils/buildAiCreditCostProperty.ts +++ b/server/src/internal/balances/track/utils/buildAiCreditCostProperty.ts @@ -14,7 +14,7 @@ export const buildAiCreditCostProperty = ({ featureDeductions: FeatureDeduction[]; entries: Array<{ featureId: string; amount: number }>; }): Record | undefined => { - const aiDeduction = featureDeductions.find((d) => d.tokenUsage); + const aiDeduction = featureDeductions.find((d) => d.tokens); if (!aiDeduction) return; const creditCost: Record = {}; diff --git a/server/src/internal/balances/track/utils/getTokenTrackParams.ts b/server/src/internal/balances/track/utils/getTokenTrackParams.ts index 97a2996c1..16a9aedc5 100644 --- a/server/src/internal/balances/track/utils/getTokenTrackParams.ts +++ b/server/src/internal/balances/track/utils/getTokenTrackParams.ts @@ -139,12 +139,14 @@ export const getTokenTrackParams = async ({ { feature: aiCreditFeature, deduction: 1, - tokenUsage: { - modelName: input.model_id, - inputTokens: input.input_tokens, - outputTokens: input.output_tokens, + tokens: { + usage: { + modelName: input.model_id, + inputTokens: input.input_tokens, + outputTokens: input.output_tokens, + }, + cost, }, - precomputedCreditCost: cost, }, ]; diff --git a/server/src/internal/balances/utils/deduction/computeCreditCosts.ts b/server/src/internal/balances/utils/deduction/computeCreditCosts.ts index 0ffe954a9..442ef83ae 100644 --- a/server/src/internal/balances/utils/deduction/computeCreditCosts.ts +++ b/server/src/internal/balances/utils/deduction/computeCreditCosts.ts @@ -8,8 +8,8 @@ export type CreditCostLookup = (entitlementId: string) => number; /** * Computes the credit cost for each customer entitlement and returns a lookup - * function. Uses precomputedCreditCost when available (token tracking), - * otherwise calls getCreditCost per entitlement (credit system schema lookups). + * function. Token deductions carry their USD cost from the API layer; all other + * costs come from credit system schema ratios. */ export const computeCreditCosts = async ({ cusEnts, @@ -20,33 +20,23 @@ export const computeCreditCosts = async ({ }): Promise => { const costMap = new Map(); - const tokens = deduction.tokenUsage - ? { - input: deduction.tokenUsage.inputTokens, - output: deduction.tokenUsage.outputTokens, - } - : undefined; - await Promise.all( cusEnts.map(async (ce) => { - // Precomputed cost (from /track/tokens) is in the AI credit feature's - // native unit (USD). It applies 1:1 to that feature's own entitlement, - // but parent credit systems still need their schema ratio applied — - // fall through to getCreditCost with amount = precomputed cost. + // A token deduction's cost is in the AI feature's native unit (USD): it + // applies 1:1 to its own entitlement, while parent credit systems apply + // their schema ratio to it via getCreditCost's amount. if ( - deduction.precomputedCreditCost != null && + deduction.tokens && ce.entitlement.feature.id === deduction.feature.id ) { - costMap.set(ce.id, deduction.precomputedCreditCost); + costMap.set(ce.id, deduction.tokens.cost); return; } const creditCost = await getCreditCost({ featureId: deduction.feature.id, creditSystem: ce.entitlement.feature, - amount: deduction.precomputedCreditCost, - modelName: deduction.tokenUsage?.modelName, - tokens, + amount: deduction.tokens?.cost, }); costMap.set(ce.id, creditCost); }), diff --git a/server/src/internal/balances/utils/types/featureDeduction.ts b/server/src/internal/balances/utils/types/featureDeduction.ts index ae6b7755d..8448112a5 100644 --- a/server/src/internal/balances/utils/types/featureDeduction.ts +++ b/server/src/internal/balances/utils/types/featureDeduction.ts @@ -7,13 +7,18 @@ export type TokenUsage = { outputTokens: number; }; +/** Token usage and its USD cost are priced together at the API layer — one cannot exist without the other. */ +export type TokenDeduction = { + usage: TokenUsage; + cost: number; +}; + export type FeatureDeduction = { feature: Feature; deduction: number; targetBalance?: number; - tokenUsage?: TokenUsage; - /** Pre-computed dollar cost; if set, the deduction layer skips its own getCreditCost call. */ - precomputedCreditCost?: number; + /** Present only for track_tokens deductions; standard deductions omit it. */ + tokens?: TokenDeduction; lock?: LockParams; lockReceipt?: LockReceipt; lockReceiptKey?: string; diff --git a/server/src/internal/features/creditSystemUtils.ts b/server/src/internal/features/creditSystemUtils.ts index c195ade41..19bb0b3c1 100644 --- a/server/src/internal/features/creditSystemUtils.ts +++ b/server/src/internal/features/creditSystemUtils.ts @@ -96,17 +96,22 @@ export const getCreditCost = async ({ return amount; } if (isAiCreditSystem(creditSystem.type)) { - if (!tokens || !modelName) { - throw new RecaseError({ - message: "modelName and tokens must be provided for AI credit systems", - code: ErrCode.InvalidRequest, - statusCode: 400, + if (tokens && modelName) { + return await getModelCreditCost({ + modelName, + creditSystem, + ...tokens, }); } - return await getModelCreditCost({ - modelName, - creditSystem, - ...tokens, + // No token context (plain /track values, balance updates, queued replays): + // the feature's own balance is already in USD, so the value maps 1:1. + if (featureId === creditSystem.id) { + return amount; + } + throw new RecaseError({ + message: "modelName and tokens must be provided for AI credit systems", + code: ErrCode.InvalidRequest, + statusCode: 400, }); } // If tracking the credit system feature itself, 1:1 mapping diff --git a/server/src/internal/features/utils/getModelPricing.ts b/server/src/internal/features/utils/getModelPricing.ts index d45d23cd6..0b4cdf0a9 100644 --- a/server/src/internal/features/utils/getModelPricing.ts +++ b/server/src/internal/features/utils/getModelPricing.ts @@ -7,9 +7,13 @@ const CACHE_KEY = "models_dev_pricing"; const STALE_KEY = `${CACHE_KEY}_stale`; const TTL_PRIMARY = 60 * 60 * 3; const TTL_STALE = 60 * 60 * 24 * 3; +// Runs inside the track request path — a hanging models.dev must not hang tracks. +const FETCH_TIMEOUT_MS = 5000; const fetchFromSource = async (): Promise => { - const response = await fetch("https://models.dev/api.json"); + const response = await fetch("https://models.dev/api.json", { + signal: AbortSignal.timeout(FETCH_TIMEOUT_MS), + }); if (!response.ok) { throw new InternalError({ message: `models.dev returned ${response.status}`, diff --git a/server/tests/integration/balances/track/basic/track-tokens-orbs.test.ts b/server/tests/integration/balances/track/basic/track-tokens-orbs.test.ts index 36dcde8b8..77f2a4576 100644 --- a/server/tests/integration/balances/track/basic/track-tokens-orbs.test.ts +++ b/server/tests/integration/balances/track/basic/track-tokens-orbs.test.ts @@ -10,13 +10,14 @@ import { Decimal } from "decimal.js"; // ═══════════════════════════════════════════════════════════════════ // TRACK-TOKENS-ORBS: AI credit system nested inside a parent credit system -// Verifies that a single /track/tokens call deducts USD from the AI credit -// feature AND deducts the ratio-mapped amount from any parent credit -// system whose schema references it. +// +// Parent credit systems are overflow pools (same semantics as classic +// metered → credits deduction order): a token track drains the AI credit +// balance first, and only the overflow is ratio-mapped onto the parent. // ═══════════════════════════════════════════════════════════════════ test.concurrent( - `${chalk.yellowBright("track-tokens-orbs: AI credit system inside parent credit system deducts both balances")}`, + `${chalk.yellowBright("track-tokens-orbs-1: AI balance covers the cost — parent orbs untouched")}`, async () => { const aiCreditsItem = items.free({ featureId: TestFeature.AiCredits, @@ -24,7 +25,7 @@ test.concurrent( }); const orbsItem = items.free({ featureId: TestFeature.Orbs, - includedUsage: 50_000, // 50,000 orbs + includedUsage: 50_000, // orbs schema: 1000 orbs per $1 of AI usage }); const freeProd = products.base({ id: "free", @@ -32,7 +33,7 @@ test.concurrent( }); const { customerId, autumnV1, autumnV2 } = await initScenario({ - customerId: "track-tokens-orbs", + customerId: "track-tokens-orbs-1", setup: [ s.customer({ testClock: false }), s.products({ list: [freeProd] }), @@ -49,9 +50,6 @@ test.concurrent( .div(1_000_000) .toNumber(); // 0.125 - // Orbs schema: 1000 orbs per $1 of AI usage - const expectedOrbsCost = new Decimal(expectedUsdCost).mul(1000).toNumber(); // 125 - const trackRes: TrackResponseV2 = await autumnV2.post("/track_tokens", { customer_id: customerId, feature_id: TestFeature.AiCredits, @@ -71,10 +69,61 @@ test.concurrent( usage: expectedUsdCost, }); - // Parent orbs balance dropped by USD cost × 1000 + // AI balance covered the full cost, so the parent overflow pool is untouched expect(customer.features[TestFeature.Orbs]).toMatchObject({ - balance: new Decimal(50_000).minus(expectedOrbsCost).toNumber(), - usage: expectedOrbsCost, + balance: 50_000, + usage: 0, + }); + }, +); + +test.concurrent( + `${chalk.yellowBright("track-tokens-orbs-2: cost exceeding AI balance overflows into parent orbs at the schema ratio")}`, + async () => { + const aiCreditsItem = items.free({ + featureId: TestFeature.AiCredits, + includedUsage: 100, // $100 of AI usage + }); + const orbsItem = items.free({ + featureId: TestFeature.Orbs, + includedUsage: 50_000, + }); + const freeProd = products.base({ + id: "free", + items: [aiCreditsItem, orbsItem], + }); + + const { customerId, autumnV1, autumnV2 } = await initScenario({ + customerId: "track-tokens-orbs-2", + setup: [ + s.customer({ testClock: false }), + s.products({ list: [freeProd] }), + ], + actions: [s.attach({ productId: freeProd.id })], + }); + + // (5 * 24M) / 1M = $120 > the $100 AI balance + const trackRes: TrackResponseV2 = await autumnV2.post("/track_tokens", { + customer_id: customerId, + feature_id: TestFeature.AiCredits, + model_id: "custom/internal-model", + input_tokens: 24_000_000, + output_tokens: 0, + }); + expect(trackRes.value).toBeCloseTo(120, 10); + + const customer = await autumnV1.customers.get(customerId); + + // AI pool fully drained + expect(customer.features[TestFeature.AiCredits]).toMatchObject({ + balance: 0, + usage: 100, + }); + + // $20 overflow lands on orbs at 1000 orbs per $1 + expect(customer.features[TestFeature.Orbs]).toMatchObject({ + balance: new Decimal(50_000).minus(20_000).toNumber(), + usage: 20_000, }); }, ); diff --git a/server/tests/integration/balances/track/basic/track-tokens-replay.test.ts b/server/tests/integration/balances/track/basic/track-tokens-replay.test.ts new file mode 100644 index 000000000..23e501696 --- /dev/null +++ b/server/tests/integration/balances/track/basic/track-tokens-replay.test.ts @@ -0,0 +1,120 @@ +import { expect, test } from "bun:test"; + +import type { ApiCustomerV3 } from "@autumn/shared"; +import { ApiVersion, ApiVersionClass } from "@autumn/shared"; +import { TestFeature } from "@tests/setup/v2Features.js"; +import { items } from "@tests/utils/fixtures/items.js"; +import { products } from "@tests/utils/fixtures/products.js"; +import { initScenario, s } from "@tests/utils/testInitUtils/initScenario.js"; +import chalk from "chalk"; +import { Decimal } from "decimal.js"; +import { runQueuedTrack } from "@/internal/balances/track/runQueuedTrack.js"; + +// ═══════════════════════════════════════════════════════════════════ +// TRACK-TOKENS-REPLAY: queued replay + plain value tracks on AI credit features +// +// When Redis fails open, track_tokens queues only the TrackParams body — the +// token context (FeatureDeduction.tokens) is not serialized. The +// replay worker rebuilds deductions from {feature_id, value}, so the USD value +// must deduct 1:1 from the AI credit balance, exactly like the original token +// track would have. Parent credit systems are overflow pools: untouched while +// the AI balance covers the deduction (same as live track_tokens behavior). +// ═══════════════════════════════════════════════════════════════════ + +test.concurrent( + `${chalk.yellowBright("track-tokens-replay-1: queued replay body deducts AI credits 1:1")}`, + async () => { + const aiCreditsItem = items.free({ + featureId: TestFeature.AiCredits, + includedUsage: 100, // $100 of AI usage + }); + const orbsItem = items.free({ + featureId: TestFeature.Orbs, + includedUsage: 50_000, // orbs schema: 1000 orbs per $1 of AI usage + }); + const freeProd = products.base({ + id: "free", + items: [aiCreditsItem, orbsItem], + }); + + const { customerId, autumnV1, ctx } = await initScenario({ + customerId: "track-tokens-replay-1", + setup: [ + s.customer({ testClock: false }), + s.products({ list: [freeProd] }), + ], + actions: [s.attach({ productId: freeProd.id })], + }); + + // The USD cost computed by the original track_tokens call; only this + // survives in the queued body. + const usdCost = 0.125; + + await runQueuedTrack({ + ctx: { ...ctx, apiVersion: new ApiVersionClass(ApiVersion.V2_1) }, + body: { + customer_id: customerId, + feature_id: TestFeature.AiCredits, + value: usdCost, + idempotency_key: `replay-${crypto.randomUUID()}`, + }, + apiVersion: ApiVersion.V2_1, + }); + + const customer = await autumnV1.customers.get(customerId); + + expect(customer.features[TestFeature.AiCredits]).toMatchObject({ + balance: new Decimal(100).minus(usdCost).toNumber(), + usage: usdCost, + }); + expect(customer.features[TestFeature.Orbs]).toMatchObject({ + balance: 50_000, + usage: 0, + }); + }, +); + +test.concurrent( + `${chalk.yellowBright("track-tokens-replay-2: plain /track with a USD value deducts an AI credit balance 1:1")}`, + async () => { + const aiCreditsItem = items.free({ + featureId: TestFeature.AiCredits, + includedUsage: 100, + }); + const orbsItem = items.free({ + featureId: TestFeature.Orbs, + includedUsage: 50_000, + }); + const freeProd = products.base({ + id: "free", + items: [aiCreditsItem, orbsItem], + }); + + const { customerId, autumnV1, autumnV2 } = await initScenario({ + customerId: "track-tokens-replay-2", + setup: [ + s.customer({ testClock: false }), + s.products({ list: [freeProd] }), + ], + actions: [s.attach({ productId: freeProd.id })], + }); + + const usdValue = 5; + await autumnV2.post("/track", { + customer_id: customerId, + feature_id: TestFeature.AiCredits, + value: usdValue, + }); + + const customer = await autumnV1.customers.get(customerId); + + expect(customer.features[TestFeature.AiCredits]).toMatchObject({ + balance: new Decimal(100).minus(usdValue).toNumber(), + usage: usdValue, + }); + expect(customer.features[TestFeature.Orbs]).toMatchObject({ + balance: 50_000, + usage: 0, + }); + }, +); diff --git a/server/tests/unit/agent/generate-and-update-agent-rules.test.ts b/server/tests/unit/agent/generate-and-update-agent-rules.test.ts new file mode 100644 index 000000000..9eff79834 --- /dev/null +++ b/server/tests/unit/agent/generate-and-update-agent-rules.test.ts @@ -0,0 +1,82 @@ +import { beforeEach, describe, expect, mock, test } from "bun:test"; +import type { AutumnContext } from "@/honoUtils/HonoEnv.js"; + +const generatedRules = { + entity_rules: { attach_to_entities: true, entity_feature_id: "deployments" }, + credit_rules: { credit_feature_id: "credits" }, + notes: "", +}; + +const mockState = { + existingNotes: "", + upsertCalls: [] as Record[], +}; + +mock.module( + "@/internal/agent/workflows/generateAgentRules/generateAgentRules.js", + () => ({ + generateAgentRules: async () => ({ + rules: generatedRules, + metadata: { generated_from: "axiom" }, + unconfigured: false, + }), + }), +); + +mock.module("@/internal/agent/rules/repos/index.js", () => ({ + agentRulesRepo: { + get: async () => ({ + entity_rules: { attach_to_entities: false, entity_feature_id: "" }, + credit_rules: { credit_feature_id: "" }, + notes: mockState.existingNotes, + metadata: {}, + org_id: "org_test", + org_slug: "test", + updated_at: null, + }), + upsert: async (args: { rules: typeof generatedRules }) => { + mockState.upsertCalls.push(args); + return { ...args.rules, metadata: {}, org_id: "org_test" }; + }, + }, +})); + +const { generateAndUpdateAgentRules } = await import( + "@/internal/agent/rules/actions/generateAndUpdateAgentRules.js" +); + +const ctx = { + db: {}, + org: { id: "org_test", slug: "test" }, +} as unknown as AutumnContext; + +describe("generateAndUpdateAgentRules", () => { + beforeEach(() => { + mockState.existingNotes = ""; + mockState.upsertCalls = []; + }); + + test("preserves existing user notes when applying generated rules", async () => { + mockState.existingNotes = "Always attach add-ons at the customer level."; + + const result = await generateAndUpdateAgentRules({ ctx }); + + expect(mockState.upsertCalls).toHaveLength(1); + expect(mockState.upsertCalls[0]).toMatchObject({ + rules: { + entity_rules: generatedRules.entity_rules, + credit_rules: generatedRules.credit_rules, + notes: "Always attach add-ons at the customer level.", + }, + }); + expect(result.notes).toBe("Always attach add-ons at the customer level."); + }); + + test("keeps notes empty when none were saved", async () => { + await generateAndUpdateAgentRules({ ctx }); + + expect(mockState.upsertCalls[0]).toMatchObject({ + rules: { notes: "" }, + }); + }); +}); diff --git a/server/tests/unit/balances/track/handle-track-tokens.test.ts b/server/tests/unit/balances/track/handle-track-tokens.test.ts index decca2c78..81f01e3e4 100644 --- a/server/tests/unit/balances/track/handle-track-tokens.test.ts +++ b/server/tests/unit/balances/track/handle-track-tokens.test.ts @@ -19,7 +19,14 @@ const featureDeductions = [ { feature: { id: "ai_credits" }, deduction: 1, - precomputedCreditCost: 3.5, + tokens: { + usage: { + modelName: "openai/gpt-4.1", + inputTokens: 100, + outputTokens: 50, + }, + cost: 3.5, + }, }, ]; diff --git a/server/tests/unit/features/get-credit-cost.test.ts b/server/tests/unit/features/get-credit-cost.test.ts new file mode 100644 index 000000000..2d97d89f1 --- /dev/null +++ b/server/tests/unit/features/get-credit-cost.test.ts @@ -0,0 +1,70 @@ +import { describe, expect, test } from "bun:test"; +import { + ErrCode, + type Feature, + FeatureType, + FeatureUsageType, +} from "@autumn/shared"; +import { getCreditCost } from "@/internal/features/creditSystemUtils.js"; + +// Uses custom/* models so pricing resolves offline (no models.dev fetch). +const CUSTOM_MODEL = "custom/foo"; + +const aiCreditFeature: Feature = { + internal_id: "fe_ai_credits", + org_id: "org_test", + created_at: Date.now(), + env: "sandbox" as Feature["env"], + id: "ai_credits", + name: "AI Credits", + type: FeatureType.AiCreditSystem, + config: { schema: [], usage_type: FeatureUsageType.Single }, + archived: false, + event_names: [], + model_markups: { + [CUSTOM_MODEL]: { markup: 0, input_cost: 1000, output_cost: 2000 }, + }, +}; + +describe("getCreditCost — AI credit system without token context", () => { + test("self feature with no tokens maps 1:1 (plain /track values, queued replays)", async () => { + const cost = await getCreditCost({ + featureId: aiCreditFeature.id, + creditSystem: aiCreditFeature, + amount: 5.25, + }); + expect(cost).toBe(5.25); + }); + + test("self feature with no tokens defaults to a per-unit cost of 1", async () => { + const cost = await getCreditCost({ + featureId: aiCreditFeature.id, + creditSystem: aiCreditFeature, + }); + expect(cost).toBe(1); + }); + + test("self feature WITH tokens still prices through the model (not 1:1)", async () => { + const cost = await getCreditCost({ + featureId: aiCreditFeature.id, + creditSystem: aiCreditFeature, + modelName: CUSTOM_MODEL, + tokens: { input: 1000, output: 500 }, + }); + // (1000 * 1000 + 2000 * 500) / 1_000_000 = 2.0 + expect(cost).toBeCloseTo(2.0, 10); + }); + + test("non-self feature with no tokens throws", async () => { + expect( + getCreditCost({ + featureId: "some_other_feature", + creditSystem: aiCreditFeature, + amount: 5, + }), + ).rejects.toMatchObject({ + code: ErrCode.InvalidRequest, + message: expect.stringContaining("modelName and tokens"), + }); + }); +}); diff --git a/server/tests/unit/features/get-model-pricing.test.ts b/server/tests/unit/features/get-model-pricing.test.ts index 67ac21dfe..9b06aedeb 100644 --- a/server/tests/unit/features/get-model-pricing.test.ts +++ b/server/tests/unit/features/get-model-pricing.test.ts @@ -113,3 +113,20 @@ test("fetch failure with no stale cache throws InternalError", async () => { message: "Failed to fetch models.dev pricing and no cache available", }); }); + +test("fetch carries an abort timeout so a hanging models.dev cannot hang tracks", async () => { + let capturedSignal: AbortSignal | undefined; + globalThis.fetch = Object.assign( + async (_input: unknown, init?: RequestInit) => { + fetchCalls++; + capturedSignal = init?.signal ?? undefined; + return Response.json(pricingData); + }, + { preconnect: realFetch.preconnect }, + ) as typeof fetch; + + await getModelsDevPricing(); + + expect(capturedSignal).toBeInstanceOf(AbortSignal); + expect(capturedSignal?.aborted).toBe(false); +}); From 402b73f71638f9b44f42c6151405f4ddbcd5f6f6 Mon Sep 17 00:00:00 2001 From: Charlie Lamb Date: Wed, 10 Jun 2026 13:27:14 +0100 Subject: [PATCH 33/46] chore: cleanup async diff --- .../track/utils/getTokenTrackParams.ts | 27 ++++----- .../utils/deduction/computeCreditCosts.ts | 42 ++++++-------- .../deduction/executePostgresDeduction.ts | 2 +- .../utils/deduction/executeRedisDeduction.ts | 2 +- .../deduction/prepareFeatureDeduction.ts | 6 +- .../deductionV2/executePostgresDeductionV2.ts | 2 +- .../deductionV2/executeRedisDeductionV2.ts | 2 +- .../deductionV2/prepareFeatureDeductionV2.ts | 6 +- .../internal/features/creditSystemUtils.ts | 40 ++++--------- .../ai-markup-resolution.test.ts | 17 +++--- server/tests/advanced/usage/usage2.test.ts | 2 +- server/tests/advanced/usage/usage3.test.ts | 2 +- server/tests/advanced/usage/usage4.test.ts | 2 +- .../credit-systems/credit-systems1.test.ts | 6 +- .../check/send-event/send-event3.test.ts | 6 +- .../check/send-event/send-event4.test.ts | 2 +- .../auto-topup-credit-systems.test.ts | 6 +- .../check-entity-product-spend-limit.test.ts | 2 +- .../check-per-entity-spend-limit.test.ts | 2 +- .../check-with-lock-credit-system.test.ts | 12 ++-- .../track/basic/track-credit-system.test.ts | 20 +++---- .../track/basic/track-deductions.test.ts | 2 +- .../balances/track/basic/track-tokens.test.ts | 57 +++++++++---------- .../track-overage-allowed-consumable.test.ts | 2 +- .../track-customer-spend-limit.test.ts | 2 +- .../track-entity-product-spend-limit.test.ts | 2 +- .../track-per-entity-spend-limit.test.ts | 2 +- .../track-postgres-entity-spend-limit.test.ts | 4 +- .../unit/features/get-credit-cost.test.ts | 51 ++++++++++------- 29 files changed, 156 insertions(+), 174 deletions(-) diff --git a/server/src/internal/balances/track/utils/getTokenTrackParams.ts b/server/src/internal/balances/track/utils/getTokenTrackParams.ts index 16a9aedc5..56d00e36b 100644 --- a/server/src/internal/balances/track/utils/getTokenTrackParams.ts +++ b/server/src/internal/balances/track/utils/getTokenTrackParams.ts @@ -11,7 +11,7 @@ import { import type { AutumnContext } from "@/honoUtils/HonoEnv.js"; import { getOrSetCachedFullSubject } from "@/internal/customers/cache/fullSubject/actions/getOrSetCachedFullSubject.js"; import { getOrSetCachedFullCustomer } from "@/internal/customers/cusUtils/fullCustomerCacheUtils/getOrSetCachedFullCustomer.js"; -import { getCreditCost } from "@/internal/features/creditSystemUtils.js"; +import { getModelCreditCost } from "@/internal/features/aiCreditSystemUtils.js"; import { isFullSubjectRolloutEnabled } from "@/internal/misc/rollouts/fullSubjectRolloutUtils.js"; import type { FeatureDeduction } from "../../utils/types/featureDeduction.js"; @@ -77,9 +77,7 @@ const resolveAiCreditFeatureFromEntitlements = async ({ const aiCreditFeatures = [ ...new Map( cusEnts - .filter( - (ce) => isAiCreditSystem(ce.entitlement.feature.type), - ) + .filter((ce) => isAiCreditSystem(ce.entitlement.feature.type)) .map((ce) => [ce.entitlement.feature.id, ce.entitlement.feature]), ).values(), ]; @@ -120,19 +118,16 @@ export const getTokenTrackParams = async ({ entityId: input.entity_id, }); - const cost = await getCreditCost({ - featureId: aiCreditFeature.id, - creditSystem: aiCreditFeature, + const cost = await getModelCreditCost({ modelName: input.model_id, - tokens: { - input: input.input_tokens, - output: input.output_tokens, - cacheRead: input.cache_read_tokens, - cacheWrite: input.cache_write_tokens, - audioInput: input.audio_input_tokens, - audioOutput: input.audio_output_tokens, - reasoning: input.reasoning_tokens, - }, + creditSystem: aiCreditFeature, + input: input.input_tokens, + output: input.output_tokens, + cacheRead: input.cache_read_tokens, + cacheWrite: input.cache_write_tokens, + audioInput: input.audio_input_tokens, + audioOutput: input.audio_output_tokens, + reasoning: input.reasoning_tokens, }); const featureDeductions: FeatureDeduction[] = [ diff --git a/server/src/internal/balances/utils/deduction/computeCreditCosts.ts b/server/src/internal/balances/utils/deduction/computeCreditCosts.ts index 442ef83ae..d1d34a04a 100644 --- a/server/src/internal/balances/utils/deduction/computeCreditCosts.ts +++ b/server/src/internal/balances/utils/deduction/computeCreditCosts.ts @@ -6,41 +6,35 @@ const DEFAULT_CREDIT_COST = 1; export type CreditCostLookup = (entitlementId: string) => number; -/** - * Computes the credit cost for each customer entitlement and returns a lookup - * function. Token deductions carry their USD cost from the API layer; all other - * costs come from credit system schema ratios. - */ -export const computeCreditCosts = async ({ +/** Per-entitlement credit cost lookup. Pure schema math — no I/O. */ +export const computeCreditCosts = ({ cusEnts, deduction, }: { cusEnts: FullCusEntWithFullCusProduct[]; deduction: FeatureDeduction; -}): Promise => { +}): CreditCostLookup => { const costMap = new Map(); - await Promise.all( - cusEnts.map(async (ce) => { - // A token deduction's cost is in the AI feature's native unit (USD): it - // applies 1:1 to its own entitlement, while parent credit systems apply - // their schema ratio to it via getCreditCost's amount. - if ( - deduction.tokens && - ce.entitlement.feature.id === deduction.feature.id - ) { - costMap.set(ce.id, deduction.tokens.cost); - return; - } + for (const ce of cusEnts) { + // Token cost is USD: 1:1 on its own ent; parents apply their ratio to it. + if ( + deduction.tokens && + ce.entitlement.feature.id === deduction.feature.id + ) { + costMap.set(ce.id, deduction.tokens.cost); + continue; + } - const creditCost = await getCreditCost({ + costMap.set( + ce.id, + getCreditCost({ featureId: deduction.feature.id, creditSystem: ce.entitlement.feature, amount: deduction.tokens?.cost, - }); - costMap.set(ce.id, creditCost); - }), - ); + }), + ); + } return (entitlementId) => costMap.get(entitlementId) ?? DEFAULT_CREDIT_COST; }; diff --git a/server/src/internal/balances/utils/deduction/executePostgresDeduction.ts b/server/src/internal/balances/utils/deduction/executePostgresDeduction.ts index 05c015f1e..d0f3a84ac 100644 --- a/server/src/internal/balances/utils/deduction/executePostgresDeduction.ts +++ b/server/src/internal/balances/utils/deduction/executePostgresDeduction.ts @@ -111,7 +111,7 @@ export const executePostgresDeduction = async ({ customerEntitlements, unlimitedFeatureIds, lock: preparedLock, - } = await prepareFeatureDeduction({ + } = prepareFeatureDeduction({ ctx, fullCustomer, deduction, diff --git a/server/src/internal/balances/utils/deduction/executeRedisDeduction.ts b/server/src/internal/balances/utils/deduction/executeRedisDeduction.ts index 43a5040b7..9ac4e797a 100644 --- a/server/src/internal/balances/utils/deduction/executeRedisDeduction.ts +++ b/server/src/internal/balances/utils/deduction/executeRedisDeduction.ts @@ -109,7 +109,7 @@ export const executeRedisDeduction = async ({ customerEntitlements, unlimitedFeatureIds, lock: preparedLock, - } = await prepareFeatureDeduction({ + } = prepareFeatureDeduction({ ctx, fullCustomer, deduction, diff --git a/server/src/internal/balances/utils/deduction/prepareFeatureDeduction.ts b/server/src/internal/balances/utils/deduction/prepareFeatureDeduction.ts index a03e6f3b5..bdb599fbf 100644 --- a/server/src/internal/balances/utils/deduction/prepareFeatureDeduction.ts +++ b/server/src/internal/balances/utils/deduction/prepareFeatureDeduction.ts @@ -27,7 +27,7 @@ import type { FeatureDeduction } from "../types/featureDeduction.js"; * Prepares all the inputs needed to execute a deduction for a single feature. * Shared by both Redis (Lua) and Postgres (SQL) deduction paths. */ -export const prepareFeatureDeduction = async ({ +export const prepareFeatureDeduction = ({ ctx, fullCustomer, deduction, @@ -37,7 +37,7 @@ export const prepareFeatureDeduction = async ({ fullCustomer: FullCustomer; deduction: FeatureDeduction; options?: DeductionOptions; -}): Promise => { +}): PreparedFeatureDeduction => { const { org } = ctx; const { env } = ctx; const { feature, lock, targetBalance } = deduction; @@ -101,7 +101,7 @@ export const prepareFeatureDeduction = async ({ .map((ce) => ce.entitlement.feature.id), ); - const getCreditCostForEnt = await computeCreditCosts({ cusEnts, deduction }); + const getCreditCostForEnt = computeCreditCosts({ cusEnts, deduction }); // Build input for each customer entitlement const customerEntitlementDeductions: CustomerEntitlementDeduction[] = diff --git a/server/src/internal/balances/utils/deductionV2/executePostgresDeductionV2.ts b/server/src/internal/balances/utils/deductionV2/executePostgresDeductionV2.ts index 894076631..99ebecc17 100644 --- a/server/src/internal/balances/utils/deductionV2/executePostgresDeductionV2.ts +++ b/server/src/internal/balances/utils/deductionV2/executePostgresDeductionV2.ts @@ -112,7 +112,7 @@ export const executePostgresDeductionV2 = async ({ unlimitedFeatureIds, unlimitedCusEnt, lock: preparedLock, - } = await prepareFeatureDeductionV2({ + } = prepareFeatureDeductionV2({ ctx, fullSubject, deduction, diff --git a/server/src/internal/balances/utils/deductionV2/executeRedisDeductionV2.ts b/server/src/internal/balances/utils/deductionV2/executeRedisDeductionV2.ts index d40a508dc..8d1d2f11e 100644 --- a/server/src/internal/balances/utils/deductionV2/executeRedisDeductionV2.ts +++ b/server/src/internal/balances/utils/deductionV2/executeRedisDeductionV2.ts @@ -123,7 +123,7 @@ export const executeRedisDeductionV2 = async ({ unlimitedFeatureIds, unlimitedCusEnt, lock: preparedLock, - } = await prepareFeatureDeductionV2({ + } = prepareFeatureDeductionV2({ ctx, fullSubject, deduction, diff --git a/server/src/internal/balances/utils/deductionV2/prepareFeatureDeductionV2.ts b/server/src/internal/balances/utils/deductionV2/prepareFeatureDeductionV2.ts index 1d09dd5a7..e7c2502ac 100644 --- a/server/src/internal/balances/utils/deductionV2/prepareFeatureDeductionV2.ts +++ b/server/src/internal/balances/utils/deductionV2/prepareFeatureDeductionV2.ts @@ -29,7 +29,7 @@ import type { FeatureDeduction } from "../types/featureDeduction.js"; * Prepares all the inputs needed to execute a deduction for a single feature. * Mirrors the legacy helper, but reads from FullSubject. */ -export const prepareFeatureDeductionV2 = async ({ +export const prepareFeatureDeductionV2 = ({ ctx, fullSubject, deduction, @@ -39,7 +39,7 @@ export const prepareFeatureDeductionV2 = async ({ fullSubject: FullSubject; deduction: FeatureDeduction; options?: DeductionOptions; -}): Promise => { +}): PreparedFeatureDeduction => { const { org, env } = ctx; const { feature, lock, targetBalance } = deduction; const { overageBehaviour = "cap", customerEntitlementFilters } = options; @@ -115,7 +115,7 @@ export const prepareFeatureDeductionV2 = async ({ .map((customerEntitlement) => customerEntitlement.entitlement.feature.id), ); - const getCreditCostForEnt = await computeCreditCosts({ + const getCreditCostForEnt = computeCreditCosts({ cusEnts: customerEntitlements, deduction, }); diff --git a/server/src/internal/features/creditSystemUtils.ts b/server/src/internal/features/creditSystemUtils.ts index 19bb0b3c1..e9ddee717 100644 --- a/server/src/internal/features/creditSystemUtils.ts +++ b/server/src/internal/features/creditSystemUtils.ts @@ -8,10 +8,6 @@ import { RecaseError, } from "@autumn/shared"; import { Decimal } from "decimal.js"; -import { - getModelCreditCost, - type TokenInput, -} from "@/internal/features/aiCreditSystemUtils.js"; const creditSystemContainsFeature = ({ creditSystem, @@ -79,45 +75,31 @@ export const featureToCreditSystem = ({ return amount; }; -export const getCreditCost = async ({ +/** Sync credit-schema math; token pricing (models.dev I/O) lives in getModelCreditCost. */ +export const getCreditCost = ({ featureId, creditSystem, amount = 1, - tokens, - modelName, }: { featureId: string; creditSystem: Feature; amount?: number; - modelName?: string; - tokens?: TokenInput; }) => { if (!isAnyCreditSystem(creditSystem.type)) { return amount; } - if (isAiCreditSystem(creditSystem.type)) { - if (tokens && modelName) { - return await getModelCreditCost({ - modelName, - creditSystem, - ...tokens, - }); - } - // No token context (plain /track values, balance updates, queued replays): - // the feature's own balance is already in USD, so the value maps 1:1. - if (featureId === creditSystem.id) { - return amount; - } - throw new RecaseError({ - message: "modelName and tokens must be provided for AI credit systems", - code: ErrCode.InvalidRequest, - statusCode: 400, - }); - } - // If tracking the credit system feature itself, 1:1 mapping + // Own balance is in the system's native unit (USD for AI), so values map 1:1. if (featureId === creditSystem.id) { return amount; } + if (isAiCreditSystem(creditSystem.type)) { + throw new RecaseError({ + message: `AI credit system ${creditSystem.id} has no schema; only its own feature can be priced here. Use getModelCreditCost for token pricing.`, + code: ErrCode.InvalidRequest, + statusCode: 400, + data: { featureId, creditSystemId: creditSystem.id }, + }); + } const schema: CreditSchemaItem[] = creditSystem.config.schema; for (const schemaItem of schema) { if (schemaItem.metered_feature_id === featureId) { diff --git a/server/tests/advanced/creditSystems/ai-markup-resolution.test.ts b/server/tests/advanced/creditSystems/ai-markup-resolution.test.ts index 390b3307b..39a9b710c 100644 --- a/server/tests/advanced/creditSystems/ai-markup-resolution.test.ts +++ b/server/tests/advanced/creditSystems/ai-markup-resolution.test.ts @@ -6,11 +6,11 @@ import { type ModelMarkups, type ProviderMarkups, } from "@autumn/shared"; -import { getCreditCost } from "@/internal/features/creditSystemUtils.js"; +import { getModelCreditCost } from "@/internal/features/aiCreditSystemUtils.js"; -// Custom models carry their own input/output costs, so getCreditCost resolves -// them without hitting the models.dev pricing fetch — ideal for unit-testing -// the tiered markup resolution (model > provider > global > none). +// Custom models carry their own input/output costs, so getModelCreditCost +// resolves them without hitting the models.dev pricing fetch — ideal for +// unit-testing the tiered markup resolution (model > provider > global > none). const CUSTOM_MODEL = "custom/foo"; const TOKENS = { input: 1000, output: 500 }; // base cost = (1000 * 1000 + 500 * 2000) / 1_000_000 = 2.0 @@ -46,14 +46,13 @@ const makeAiCredit = ({ }); const cost = (creditSystem: Feature) => - getCreditCost({ - featureId: "ai_credits", - creditSystem, - tokens: TOKENS, + getModelCreditCost({ modelName: CUSTOM_MODEL, + creditSystem, + ...TOKENS, }); -describe("getCreditCost — tiered AI markup resolution", () => { +describe("getModelCreditCost — tiered AI markup resolution", () => { test("per-model markup wins over provider and global", async () => { const creditSystem = makeAiCredit({ model_markups: { diff --git a/server/tests/advanced/usage/usage2.test.ts b/server/tests/advanced/usage/usage2.test.ts index 75344a9dd..886ce2e4d 100644 --- a/server/tests/advanced/usage/usage2.test.ts +++ b/server/tests/advanced/usage/usage2.test.ts @@ -104,7 +104,7 @@ describe(`${chalk.yellowBright("usage2: Testing basic usage product")}`, () => { .toNumber(); const featureId = i % 2 === 0 ? TestFeature.Action1 : TestFeature.Action2; - const creditsUsed = await getCreditCost({ + const creditsUsed = getCreditCost({ creditSystem: creditsFeature, featureId: featureId, amount: randomVal, diff --git a/server/tests/advanced/usage/usage3.test.ts b/server/tests/advanced/usage/usage3.test.ts index 6a48c6ebe..c1474e470 100644 --- a/server/tests/advanced/usage/usage3.test.ts +++ b/server/tests/advanced/usage/usage3.test.ts @@ -121,7 +121,7 @@ describe(`${chalk.yellowBright( .toNumber(); const featureId = i % 2 === 0 ? TestFeature.Action1 : TestFeature.Action2; - const creditsUsed = await getCreditCost({ + const creditsUsed = getCreditCost({ creditSystem: creditsFeature, featureId: featureId, amount: randomVal, diff --git a/server/tests/advanced/usage/usage4.test.ts b/server/tests/advanced/usage/usage4.test.ts index 7199ecc5b..3a3d6f1a8 100644 --- a/server/tests/advanced/usage/usage4.test.ts +++ b/server/tests/advanced/usage/usage4.test.ts @@ -105,7 +105,7 @@ describe(`${chalk.yellowBright("usage4: GPU starter annual")}`, () => { .toNumber(); const featureId = i % 2 === 0 ? TestFeature.Action1 : TestFeature.Action2; - const creditsUsed = await getCreditCost({ + const creditsUsed = getCreditCost({ creditSystem: creditsFeature, featureId: featureId, amount: randomVal, diff --git a/server/tests/balances/check/credit-systems/credit-systems1.test.ts b/server/tests/balances/check/credit-systems/credit-systems1.test.ts index 6fe182bee..22f1a2fef 100644 --- a/server/tests/balances/check/credit-systems/credit-systems1.test.ts +++ b/server/tests/balances/check/credit-systems/credit-systems1.test.ts @@ -71,7 +71,7 @@ describe(`${chalk.yellowBright("credit-systems1: test /check on action that uses required_balance: requiredActionUnits, })) as unknown as CheckResponseV2; - const creditCost = await getCreditCost({ + const creditCost = getCreditCost({ featureId: action, creditSystem: creditFeature!, amount: requiredActionUnits, @@ -172,7 +172,7 @@ describe(`${chalk.yellowBright("credit-systems1: test /check on action that uses required_balance: requiredAction1Units, })) as unknown as CheckResponseV0; - const meteredCost = await getCreditCost({ + const meteredCost = getCreditCost({ featureId: TestFeature.Action1, creditSystem: creditFeature!, amount: requiredAction1Units, @@ -196,7 +196,7 @@ describe(`${chalk.yellowBright("credit-systems1: test /check on action that uses required_balance: requiredAction2Units, })) as unknown as CheckResponseV0; - const meteredCost = await getCreditCost({ + const meteredCost = getCreditCost({ featureId: TestFeature.Action2, creditSystem: creditFeature!, amount: requiredAction2Units, diff --git a/server/tests/balances/check/send-event/send-event3.test.ts b/server/tests/balances/check/send-event/send-event3.test.ts index 4d84de2ca..866be0b0b 100644 --- a/server/tests/balances/check/send-event/send-event3.test.ts +++ b/server/tests/balances/check/send-event/send-event3.test.ts @@ -93,7 +93,7 @@ describe(`${chalk.yellowBright("send-event3: Testing check with track, credit sy send_event: true, })) as unknown as CheckResponseV2; - const creditCost = await getCreditCost({ + const creditCost = getCreditCost({ featureId: TestFeature.Action1, creditSystem: creditFeature!, amount: 10, @@ -156,7 +156,7 @@ describe(`${chalk.yellowBright("send-event3: Testing check with track, credit sy allowed: true, customer_id: customerId, feature_id: TestFeature.Credits, - required_balance: await getCreditCost({ + required_balance: getCreditCost({ featureId: TestFeature.Action1, creditSystem: creditFeature!, amount: 10, @@ -167,7 +167,7 @@ describe(`${chalk.yellowBright("send-event3: Testing check with track, credit sy test("should check with track and deduct from credits", async () => { const value = 2.5; - const creditCost = await getCreditCost({ + const creditCost = getCreditCost({ featureId: TestFeature.Action1, creditSystem: creditFeature!, amount: value, diff --git a/server/tests/balances/check/send-event/send-event4.test.ts b/server/tests/balances/check/send-event/send-event4.test.ts index 05ce02f20..aa1a91409 100644 --- a/server/tests/balances/check/send-event/send-event4.test.ts +++ b/server/tests/balances/check/send-event/send-event4.test.ts @@ -82,7 +82,7 @@ describe(`${chalk.yellowBright("send-event4: Testing check with track, unlimited send_event: true, }); - const requiredBalance = await getCreditCost({ + const requiredBalance = getCreditCost({ featureId: TestFeature.Action1, creditSystem: ctx.features.find((f) => f.id === TestFeature.Credits)!, amount: 1000, diff --git a/server/tests/integration/balances/auto-topup/auto-topup-credit-systems.test.ts b/server/tests/integration/balances/auto-topup/auto-topup-credit-systems.test.ts index 5336d752e..85399f694 100644 --- a/server/tests/integration/balances/auto-topup/auto-topup-credit-systems.test.ts +++ b/server/tests/integration/balances/auto-topup/auto-topup-credit-systems.test.ts @@ -77,7 +77,7 @@ test.concurrent(`${chalk.yellowBright("auto-topup cs1: action track depletes cre // Track 845 units → 845 × 0.2 = 169 credits deducted // Balance: 200 - 169 = 31 → strictly above threshold (30) → does NOT trigger // (exact threshold uses <= in code, so landing on 30 would fire auto top-up) - const action1Cost = await getCreditCost({ + const action1Cost = getCreditCost({ featureId: TestFeature.Action1, creditSystem: creditFeature!, amount: 845, @@ -98,7 +98,7 @@ test.concurrent(`${chalk.yellowBright("auto-topup cs1: action track depletes cre // Track 10 units of action1 → 10 × 0.2 = 2 credits // Balance: 31 - 2 = 29 → 29 <= threshold → auto top-up fires → 29 + 100 = 129 - const action1CostSmall = await getCreditCost({ + const action1CostSmall = getCreditCost({ featureId: TestFeature.Action1, creditSystem: creditFeature!, amount: 10, @@ -169,7 +169,7 @@ test.concurrent(`${chalk.yellowBright("auto-topup cs2: action track depletes cre // Action1 costs 0.2 credits per unit // Track 900 units of action1 → 900 × 0.2 = 180 credits deducted // Balance: 200 - 180 = 20 → auto top-up fires - const action1Cost = await getCreditCost({ + const action1Cost = getCreditCost({ featureId: TestFeature.Action1, creditSystem: creditFeature!, amount: 900, diff --git a/server/tests/integration/balances/check/spend-limit/check-entity-product-spend-limit.test.ts b/server/tests/integration/balances/check/spend-limit/check-entity-product-spend-limit.test.ts index 988e87ed8..94487bbe0 100644 --- a/server/tests/integration/balances/check/spend-limit/check-entity-product-spend-limit.test.ts +++ b/server/tests/integration/balances/check/spend-limit/check-entity-product-spend-limit.test.ts @@ -362,7 +362,7 @@ test.concurrent(`${chalk.yellowBright("check-entity-product-spend-limit5: credit const creditsFeature = ctx.features.find( (f) => f.id === TestFeature.Credits, )!; - const action1CreditCost = await getCreditCost({ + const action1CreditCost = getCreditCost({ featureId: TestFeature.Action1, creditSystem: creditsFeature, amount: 1, diff --git a/server/tests/integration/balances/check/spend-limit/check-per-entity-spend-limit.test.ts b/server/tests/integration/balances/check/spend-limit/check-per-entity-spend-limit.test.ts index dab02bfdf..83d0f53c8 100644 --- a/server/tests/integration/balances/check/spend-limit/check-per-entity-spend-limit.test.ts +++ b/server/tests/integration/balances/check/spend-limit/check-per-entity-spend-limit.test.ts @@ -210,7 +210,7 @@ test.concurrent(`${chalk.yellowBright("check-per-entity-spend-limit4: credit-sys const creditsFeature = ctx.features.find( (f) => f.id === TestFeature.Credits, )!; - const action1CreditCost = await getCreditCost({ + const action1CreditCost = getCreditCost({ featureId: TestFeature.Action1, creditSystem: creditsFeature, amount: 1, diff --git a/server/tests/integration/balances/lock/check-with-lock-credit-system.test.ts b/server/tests/integration/balances/lock/check-with-lock-credit-system.test.ts index b60c09c95..3c4598022 100644 --- a/server/tests/integration/balances/lock/check-with-lock-credit-system.test.ts +++ b/server/tests/integration/balances/lock/check-with-lock-credit-system.test.ts @@ -131,12 +131,12 @@ test.concurrent(`${chalk.yellowBright("lock-credit CS-2: cross-boundary lock=8 c }); const creditFeature = features.find((f) => f.id === TestFeature.Credits)!; - const lockCreditCost = await getCreditCost({ + const lockCreditCost = getCreditCost({ featureId: TestFeature.Action1, creditSystem: creditFeature, amount: 3, // overflow during lock: 8 - 5 remaining = 3 }); - const extraCreditCost = await getCreditCost({ + const extraCreditCost = getCreditCost({ featureId: TestFeature.Action1, creditSystem: creditFeature, amount: 4, // confirm delta: 12 - 8 = 4 more units @@ -289,7 +289,7 @@ test.concurrent(`${chalk.yellowBright("lock-credit CS-4: lock within action1, co // Lock deducted 10 from action1 (→90). Confirm delta=+105: // exhaust remaining 90 from action1 (→0), then 15 overflow → 15×0.2=3 credits. const creditFeature = features.find((f) => f.id === TestFeature.Credits)!; - const overflowCreditCost = await getCreditCost({ + const overflowCreditCost = getCreditCost({ featureId: TestFeature.Action1, creditSystem: creditFeature, amount: 15, @@ -628,12 +628,12 @@ test.concurrent(`${chalk.yellowBright("lock-credit CS-9: cross-boundary lock=8 c // Lock deducted: 5 from action1 + 3 overflow (0.6 credits). // Confirm delta = 20 - 8 = 12 more units, action1 is already 0, all go to credits: 12×0.2=2.4. const creditFeature = features.find((f) => f.id === TestFeature.Credits)!; - const lockOverflowCost = await getCreditCost({ + const lockOverflowCost = getCreditCost({ featureId: TestFeature.Action1, creditSystem: creditFeature, amount: 3, }); - const confirmExtraCost = await getCreditCost({ + const confirmExtraCost = getCreditCost({ featureId: TestFeature.Action1, creditSystem: creditFeature, amount: 12, @@ -721,7 +721,7 @@ test.concurrent(`${chalk.yellowBright("lock-credit CS-10: confirm no override_va // Balances unchanged from what the lock left const creditFeature = features.find((f) => f.id === TestFeature.Credits)!; - const lockOverflowCost = await getCreditCost({ + const lockOverflowCost = getCreditCost({ featureId: TestFeature.Action1, creditSystem: creditFeature, amount: 3, // overflow during lock: 8 - 5 remaining = 3 diff --git a/server/tests/integration/balances/track/basic/track-credit-system.test.ts b/server/tests/integration/balances/track/basic/track-credit-system.test.ts index a84933213..7cd345c7d 100644 --- a/server/tests/integration/balances/track/basic/track-credit-system.test.ts +++ b/server/tests/integration/balances/track/basic/track-credit-system.test.ts @@ -88,7 +88,7 @@ test.concurrent(`${chalk.yellowBright("track-credit-system2: track metered featu expect(customerBefore.features[TestFeature.Credits].balance).toBe(200); const action1Value = 50.25; - const expectedAction1CreditCost = await getCreditCost({ + const expectedAction1CreditCost = getCreditCost({ featureId: TestFeature.Action1, creditSystem: creditFeature!, amount: action1Value, @@ -108,7 +108,7 @@ test.concurrent(`${chalk.yellowBright("track-credit-system2: track metered featu }); const action2Value = 33.67; - const expectedAction2CreditCost = await getCreditCost({ + const expectedAction2CreditCost = getCreditCost({ featureId: TestFeature.Action2, creditSystem: creditFeature!, amount: action2Value, @@ -209,7 +209,7 @@ test.concurrent(`${chalk.yellowBright("track-credit-system3: test deduction orde const deduct2 = 80; const remainingAction1 = 100 - deduct1; const overflowAmount = deduct2 - remainingAction1; - const creditCostForOverflow = await getCreditCost({ + const creditCostForOverflow = getCreditCost({ featureId: TestFeature.Action1, creditSystem: creditFeature!, amount: overflowAmount, @@ -244,7 +244,7 @@ test.concurrent(`${chalk.yellowBright("track-credit-system3: test deduction orde const creditsBefore = customer2.features[TestFeature.Credits].balance; const deduct3 = 50.75; - const creditCost3 = await getCreditCost({ + const creditCost3 = getCreditCost({ featureId: TestFeature.Action1, creditSystem: creditFeature!, amount: deduct3, @@ -367,13 +367,13 @@ test.concurrent(`${chalk.yellowBright("track-credit-system4: test deduction with const overflowAction1 = deduct2 - remainingAction1; const overflowAction3 = deduct2 - remainingAction3; - const creditCostAction1 = await getCreditCost({ + const creditCostAction1 = getCreditCost({ featureId: TestFeature.Action1, creditSystem: creditFeature!, amount: overflowAction1, }); - const creditCostAction3 = await getCreditCost({ + const creditCostAction3 = getCreditCost({ featureId: TestFeature.Action3, creditSystem: credit2Feature!, amount: overflowAction3, @@ -419,13 +419,13 @@ test.concurrent(`${chalk.yellowBright("track-credit-system4: test deduction with const deduct3 = 40.25; - const creditCostAction1Deduct3 = await getCreditCost({ + const creditCostAction1Deduct3 = getCreditCost({ featureId: TestFeature.Action1, creditSystem: creditFeature!, amount: deduct3, }); - const creditCostAction3Deduct3 = await getCreditCost({ + const creditCostAction3Deduct3 = getCreditCost({ featureId: TestFeature.Action3, creditSystem: credit2Feature!, amount: deduct3, @@ -556,7 +556,7 @@ test.concurrent(`${chalk.yellowBright("track-credit-system5: test deduction orde const deduct2 = 80; const remainingAction1 = 100 - deduct1; const overflowAmount = deduct2 - remainingAction1; - const creditCostForOverflow = await getCreditCost({ + const creditCostForOverflow = getCreditCost({ featureId: TestFeature.Action1, creditSystem: creditFeature!, amount: overflowAmount, @@ -596,7 +596,7 @@ test.concurrent(`${chalk.yellowBright("track-credit-system5: test deduction orde const creditsBefore = customer2.features[TestFeature.Credits].balance; const deduct3 = 50.75; - const creditCost3 = await getCreditCost({ + const creditCost3 = getCreditCost({ featureId: TestFeature.Action1, creditSystem: creditFeature!, amount: deduct3, diff --git a/server/tests/integration/balances/track/basic/track-deductions.test.ts b/server/tests/integration/balances/track/basic/track-deductions.test.ts index 7a4b6e6c8..55eb828be 100644 --- a/server/tests/integration/balances/track/basic/track-deductions.test.ts +++ b/server/tests/integration/balances/track/basic/track-deductions.test.ts @@ -306,7 +306,7 @@ test.concurrent( }); const overflowAmount = 50; - const expectedCreditCost = await getCreditCost({ + const expectedCreditCost = getCreditCost({ featureId: TestFeature.Action1, creditSystem: creditFeature, amount: overflowAmount, diff --git a/server/tests/integration/balances/track/basic/track-tokens.test.ts b/server/tests/integration/balances/track/basic/track-tokens.test.ts index 6db39c9fa..fb891fac8 100644 --- a/server/tests/integration/balances/track/basic/track-tokens.test.ts +++ b/server/tests/integration/balances/track/basic/track-tokens.test.ts @@ -15,7 +15,7 @@ import { products } from "@tests/utils/fixtures/products.js"; import { initScenario, s } from "@tests/utils/testInitUtils/initScenario.js"; import chalk from "chalk"; import { Decimal } from "decimal.js"; -import { getCreditCost } from "@/internal/features/creditSystemUtils.js"; +import { getModelCreditCost } from "@/internal/features/aiCreditSystemUtils.js"; // ═══════════════════════════════════════════════════════════════════ // TRACK-TOKENS-1: Basic trackTokens with models.dev pricing @@ -57,11 +57,11 @@ test.concurrent( const outputTokens = 500; const modelId = "anthropic/claude-sonnet-4-20250514"; - const expectedCost = await getCreditCost({ - featureId: aiCreditFeature.id, - creditSystem: aiCreditFeature, + const expectedCost = await getModelCreditCost({ modelName: modelId, - tokens: { input: inputTokens, output: outputTokens }, + creditSystem: aiCreditFeature, + input: inputTokens, + output: outputTokens, }); const trackRes: TrackResponseV2 = await autumnV2.post("/track_tokens", { @@ -140,11 +140,11 @@ test.concurrent( const outputTokens = 1000; const modelId = "anthropic/claude-sonnet-4-20250514"; - const expectedCost = await getCreditCost({ - featureId: aiCreditFeature.id, - creditSystem: aiCreditFeature, + const expectedCost = await getModelCreditCost({ modelName: modelId, - tokens: { input: inputTokens, output: outputTokens }, + creditSystem: aiCreditFeature, + input: inputTokens, + output: outputTokens, }); const trackRes: TrackResponseV2 = await autumnV2.post("/track_tokens", { @@ -284,11 +284,11 @@ test.concurrent( const outputTokens = 10000; const modelId = "anthropic/claude-3-5-haiku-20241022"; - const expectedCost = await getCreditCost({ - featureId: aiCreditFeature.id, - creditSystem: aiCreditFeature, + const expectedCost = await getModelCreditCost({ modelName: modelId, - tokens: { input: inputTokens, output: outputTokens }, + creditSystem: aiCreditFeature, + input: inputTokens, + output: outputTokens, }); const trackRes: TrackResponseV2 = await autumnV2.post("/track_tokens", { @@ -358,11 +358,11 @@ test.concurrent( } // First track: custom/internal-model (input_cost=5, output_cost=15, markup=0%) - const cost1 = await getCreditCost({ - featureId: aiCreditFeature.id, - creditSystem: aiCreditFeature, + const cost1 = await getModelCreditCost({ modelName: "custom/internal-model", - tokens: { input: 5000, output: 2000 }, + creditSystem: aiCreditFeature, + input: 5000, + output: 2000, }); await autumnV2.post("/track_tokens", { @@ -374,11 +374,11 @@ test.concurrent( }); // Second track: custom/marked-up-model (input_cost=10, output_cost=30, markup=50%) - const cost2 = await getCreditCost({ - featureId: aiCreditFeature.id, - creditSystem: aiCreditFeature, + const cost2 = await getModelCreditCost({ modelName: "custom/marked-up-model", - tokens: { input: 3000, output: 1000 }, + creditSystem: aiCreditFeature, + input: 3000, + output: 1000, }); await autumnV2.post("/track_tokens", { @@ -491,20 +491,19 @@ test.concurrent( reasoning: 4000, }; - const expectedCost = await getCreditCost({ - featureId: aiCreditFeature.id, - creditSystem: aiCreditFeature, + const expectedCost = await getModelCreditCost({ modelName: modelId, - tokens: pools, + creditSystem: aiCreditFeature, + ...pools, }); // Pools must increase the bill vs text-only — otherwise the assertion // below couldn't tell whether the HTTP layer forwarded them at all. - const textOnlyCost = await getCreditCost({ - featureId: aiCreditFeature.id, - creditSystem: aiCreditFeature, + const textOnlyCost = await getModelCreditCost({ modelName: modelId, - tokens: { input: pools.input, output: pools.output }, + creditSystem: aiCreditFeature, + input: pools.input, + output: pools.output, }); expect(expectedCost).toBeGreaterThan(textOnlyCost); diff --git a/server/tests/integration/balances/track/overage-allowed/track-overage-allowed-consumable.test.ts b/server/tests/integration/balances/track/overage-allowed/track-overage-allowed-consumable.test.ts index 636513eb9..6e3242fad 100644 --- a/server/tests/integration/balances/track/overage-allowed/track-overage-allowed-consumable.test.ts +++ b/server/tests/integration/balances/track/overage-allowed/track-overage-allowed-consumable.test.ts @@ -368,7 +368,7 @@ test.concurrent(`${chalk.yellowBright("track-consumable-overage-8: credit system (feature) => feature.id === TestFeature.Credits, )!; - const action1CreditCost = await getCreditCost({ + const action1CreditCost = getCreditCost({ featureId: TestFeature.Action1, creditSystem: creditsFeature, amount: 1, diff --git a/server/tests/integration/balances/track/spend-limit/track-customer-spend-limit.test.ts b/server/tests/integration/balances/track/spend-limit/track-customer-spend-limit.test.ts index 25c11d2ec..64e0b3622 100644 --- a/server/tests/integration/balances/track/spend-limit/track-customer-spend-limit.test.ts +++ b/server/tests/integration/balances/track/spend-limit/track-customer-spend-limit.test.ts @@ -430,7 +430,7 @@ test.concurrent(`${chalk.yellowBright("track-customer-spend-limit6: credit-syste const creditsFeature = ctx.features.find( (feature) => feature.id === TestFeature.Credits, )!; - const action1CreditCost = await getCreditCost({ + const action1CreditCost = getCreditCost({ featureId: TestFeature.Action1, creditSystem: creditsFeature, amount: 1, diff --git a/server/tests/integration/balances/track/spend-limit/track-entity-product-spend-limit.test.ts b/server/tests/integration/balances/track/spend-limit/track-entity-product-spend-limit.test.ts index 79fae21fa..85c5ab273 100644 --- a/server/tests/integration/balances/track/spend-limit/track-entity-product-spend-limit.test.ts +++ b/server/tests/integration/balances/track/spend-limit/track-entity-product-spend-limit.test.ts @@ -463,7 +463,7 @@ test.concurrent(`${chalk.yellowBright("track-entity-product-spend-limit5: credit const creditsFeature = ctx.features.find( (feature) => feature.id === TestFeature.Credits, )!; - const action1CreditCost = await getCreditCost({ + const action1CreditCost = getCreditCost({ featureId: TestFeature.Action1, creditSystem: creditsFeature, amount: 1, diff --git a/server/tests/integration/balances/track/spend-limit/track-per-entity-spend-limit.test.ts b/server/tests/integration/balances/track/spend-limit/track-per-entity-spend-limit.test.ts index 89822ddad..e61f687bd 100644 --- a/server/tests/integration/balances/track/spend-limit/track-per-entity-spend-limit.test.ts +++ b/server/tests/integration/balances/track/spend-limit/track-per-entity-spend-limit.test.ts @@ -489,7 +489,7 @@ test.concurrent(`${chalk.yellowBright("track-per-entity-spend-limit5: credit-sys const creditsFeature = ctx.features.find( (feature) => feature.id === TestFeature.Credits, )!; - const action1CreditCost = await getCreditCost({ + const action1CreditCost = getCreditCost({ featureId: TestFeature.Action1, creditSystem: creditsFeature, amount: 1, diff --git a/server/tests/integration/balances/track/spend-limit/track-postgres-entity-spend-limit.test.ts b/server/tests/integration/balances/track/spend-limit/track-postgres-entity-spend-limit.test.ts index 72ebef0b4..00422b328 100644 --- a/server/tests/integration/balances/track/spend-limit/track-postgres-entity-spend-limit.test.ts +++ b/server/tests/integration/balances/track/spend-limit/track-postgres-entity-spend-limit.test.ts @@ -216,7 +216,7 @@ test.concurrent(`${chalk.yellowBright("track-postgres-entity-spend-limit3: credi const creditsFeature = ctx.features.find( (feature) => feature.id === TestFeature.Credits, )!; - const action1CreditCost = await getCreditCost({ + const action1CreditCost = getCreditCost({ featureId: TestFeature.Action1, creditSystem: creditsFeature, amount: 1, @@ -334,7 +334,7 @@ test.concurrent(`${chalk.yellowBright("track-postgres-entity-spend-limit4: prepa const creditsFeature = ctx.features.find( (feature) => feature.id === TestFeature.Credits, )!; - const action1CreditCost = await getCreditCost({ + const action1CreditCost = getCreditCost({ featureId: TestFeature.Action1, creditSystem: creditsFeature, amount: 1, diff --git a/server/tests/unit/features/get-credit-cost.test.ts b/server/tests/unit/features/get-credit-cost.test.ts index 2d97d89f1..7fb7c5553 100644 --- a/server/tests/unit/features/get-credit-cost.test.ts +++ b/server/tests/unit/features/get-credit-cost.test.ts @@ -5,6 +5,7 @@ import { FeatureType, FeatureUsageType, } from "@autumn/shared"; +import { getModelCreditCost } from "@/internal/features/aiCreditSystemUtils.js"; import { getCreditCost } from "@/internal/features/creditSystemUtils.js"; // Uses custom/* models so pricing resolves offline (no models.dev fetch). @@ -26,9 +27,9 @@ const aiCreditFeature: Feature = { }, }; -describe("getCreditCost — AI credit system without token context", () => { - test("self feature with no tokens maps 1:1 (plain /track values, queued replays)", async () => { - const cost = await getCreditCost({ +describe("getCreditCost — AI credit system schema math", () => { + test("self feature maps 1:1 (plain /track values, queued replays)", () => { + const cost = getCreditCost({ featureId: aiCreditFeature.id, creditSystem: aiCreditFeature, amount: 5.25, @@ -36,35 +37,47 @@ describe("getCreditCost — AI credit system without token context", () => { expect(cost).toBe(5.25); }); - test("self feature with no tokens defaults to a per-unit cost of 1", async () => { - const cost = await getCreditCost({ + test("self feature defaults to a per-unit cost of 1", () => { + const cost = getCreditCost({ featureId: aiCreditFeature.id, creditSystem: aiCreditFeature, }); expect(cost).toBe(1); }); - test("self feature WITH tokens still prices through the model (not 1:1)", async () => { - const cost = await getCreditCost({ - featureId: aiCreditFeature.id, - creditSystem: aiCreditFeature, - modelName: CUSTOM_MODEL, - tokens: { input: 1000, output: 500 }, - }); - // (1000 * 1000 + 2000 * 500) / 1_000_000 = 2.0 - expect(cost).toBeCloseTo(2.0, 10); - }); - - test("non-self feature with no tokens throws", async () => { - expect( + test("non-self feature throws — AI credit systems have no schema", () => { + expect(() => getCreditCost({ featureId: "some_other_feature", creditSystem: aiCreditFeature, amount: 5, }), + ).toThrow(/no schema/); + }); +}); + +describe("getModelCreditCost — token pricing", () => { + test("prices through the model markup config", async () => { + const cost = await getModelCreditCost({ + modelName: CUSTOM_MODEL, + creditSystem: aiCreditFeature, + input: 1000, + output: 500, + }); + // (1000 * 1000 + 2000 * 500) / 1_000_000 = 2.0 + expect(cost).toBeCloseTo(2.0, 10); + }); + + test("custom model without configured costs throws", async () => { + expect( + getModelCreditCost({ + modelName: "custom/unconfigured", + creditSystem: aiCreditFeature, + input: 100, + output: 50, + }), ).rejects.toMatchObject({ code: ErrCode.InvalidRequest, - message: expect.stringContaining("modelName and tokens"), }); }); }); From 6d747ba4389c9c73bbdbe15e9b1cdef0545ed68b Mon Sep 17 00:00:00 2001 From: Charlie Lamb Date: Wed, 10 Jun 2026 13:41:24 +0100 Subject: [PATCH 34/46] chore: cleanup out of scope diff --- .../actions/generateAndUpdateAgentRules.ts | 8 +- .../utils/deduction/computeCreditCosts.ts | 29 +++++-- .../generate-and-update-agent-rules.test.ts | 82 ------------------ .../balances/compute-credit-costs.test.ts | 84 +++++++++++++++++++ 4 files changed, 107 insertions(+), 96 deletions(-) delete mode 100644 server/tests/unit/agent/generate-and-update-agent-rules.test.ts create mode 100644 server/tests/unit/balances/compute-credit-costs.test.ts diff --git a/server/src/internal/agent/rules/actions/generateAndUpdateAgentRules.ts b/server/src/internal/agent/rules/actions/generateAndUpdateAgentRules.ts index f3acc2ee4..34cd35d26 100644 --- a/server/src/internal/agent/rules/actions/generateAndUpdateAgentRules.ts +++ b/server/src/internal/agent/rules/actions/generateAndUpdateAgentRules.ts @@ -11,17 +11,13 @@ export const generateAndUpdateAgentRules = async ({ endTime?: string; startTime?: string; }) => { - const [generated, existing] = await Promise.all([ - generateAgentRules({ ctx, endTime, startTime }), - agentRulesRepo.get({ db: ctx.db, orgId: ctx.org.id }), - ]); - // Generation only derives entity/credit rules; never overwrite user-written notes. + const generated = await generateAgentRules({ ctx, endTime, startTime }); const rules = await agentRulesRepo.upsert({ db: ctx.db, metadata: generated.metadata, orgId: ctx.org.id, orgSlug: ctx.org.slug, - rules: { ...generated.rules, notes: existing.notes }, + rules: generated.rules, }); return { diff --git a/server/src/internal/balances/utils/deduction/computeCreditCosts.ts b/server/src/internal/balances/utils/deduction/computeCreditCosts.ts index d1d34a04a..998f12a97 100644 --- a/server/src/internal/balances/utils/deduction/computeCreditCosts.ts +++ b/server/src/internal/balances/utils/deduction/computeCreditCosts.ts @@ -1,4 +1,5 @@ import type { FullCusEntWithFullCusProduct } from "@autumn/shared"; +import { logger } from "@/external/logtail/logtailUtils.js"; import { getCreditCost } from "@/internal/features/creditSystemUtils.js"; import type { FeatureDeduction } from "../types/featureDeduction.js"; @@ -26,14 +27,26 @@ export const computeCreditCosts = ({ continue; } - costMap.set( - ce.id, - getCreditCost({ - featureId: deduction.feature.id, - creditSystem: ce.entitlement.feature, - amount: deduction.tokens?.cost, - }), - ); + try { + costMap.set( + ce.id, + getCreditCost({ + featureId: deduction.feature.id, + creditSystem: ce.entitlement.feature, + amount: deduction.tokens?.cost, + }), + ); + } catch (error) { + // Cached cusEnt schemas can briefly trail a feature update; deduct at + // 1:1 rather than failing the track. + logger.warn("[computeCreditCosts] falling back to credit cost 1", { + feature_id: deduction.feature.id, + credit_system_id: ce.entitlement.feature.id, + customer_entitlement_id: ce.id, + error: String(error), + }); + costMap.set(ce.id, DEFAULT_CREDIT_COST); + } } return (entitlementId) => costMap.get(entitlementId) ?? DEFAULT_CREDIT_COST; diff --git a/server/tests/unit/agent/generate-and-update-agent-rules.test.ts b/server/tests/unit/agent/generate-and-update-agent-rules.test.ts deleted file mode 100644 index 9eff79834..000000000 --- a/server/tests/unit/agent/generate-and-update-agent-rules.test.ts +++ /dev/null @@ -1,82 +0,0 @@ -import { beforeEach, describe, expect, mock, test } from "bun:test"; -import type { AutumnContext } from "@/honoUtils/HonoEnv.js"; - -const generatedRules = { - entity_rules: { attach_to_entities: true, entity_feature_id: "deployments" }, - credit_rules: { credit_feature_id: "credits" }, - notes: "", -}; - -const mockState = { - existingNotes: "", - upsertCalls: [] as Record[], -}; - -mock.module( - "@/internal/agent/workflows/generateAgentRules/generateAgentRules.js", - () => ({ - generateAgentRules: async () => ({ - rules: generatedRules, - metadata: { generated_from: "axiom" }, - unconfigured: false, - }), - }), -); - -mock.module("@/internal/agent/rules/repos/index.js", () => ({ - agentRulesRepo: { - get: async () => ({ - entity_rules: { attach_to_entities: false, entity_feature_id: "" }, - credit_rules: { credit_feature_id: "" }, - notes: mockState.existingNotes, - metadata: {}, - org_id: "org_test", - org_slug: "test", - updated_at: null, - }), - upsert: async (args: { rules: typeof generatedRules }) => { - mockState.upsertCalls.push(args); - return { ...args.rules, metadata: {}, org_id: "org_test" }; - }, - }, -})); - -const { generateAndUpdateAgentRules } = await import( - "@/internal/agent/rules/actions/generateAndUpdateAgentRules.js" -); - -const ctx = { - db: {}, - org: { id: "org_test", slug: "test" }, -} as unknown as AutumnContext; - -describe("generateAndUpdateAgentRules", () => { - beforeEach(() => { - mockState.existingNotes = ""; - mockState.upsertCalls = []; - }); - - test("preserves existing user notes when applying generated rules", async () => { - mockState.existingNotes = "Always attach add-ons at the customer level."; - - const result = await generateAndUpdateAgentRules({ ctx }); - - expect(mockState.upsertCalls).toHaveLength(1); - expect(mockState.upsertCalls[0]).toMatchObject({ - rules: { - entity_rules: generatedRules.entity_rules, - credit_rules: generatedRules.credit_rules, - notes: "Always attach add-ons at the customer level.", - }, - }); - expect(result.notes).toBe("Always attach add-ons at the customer level."); - }); - - test("keeps notes empty when none were saved", async () => { - await generateAndUpdateAgentRules({ ctx }); - - expect(mockState.upsertCalls[0]).toMatchObject({ - rules: { notes: "" }, - }); - }); -}); diff --git a/server/tests/unit/balances/compute-credit-costs.test.ts b/server/tests/unit/balances/compute-credit-costs.test.ts new file mode 100644 index 000000000..38df654d1 --- /dev/null +++ b/server/tests/unit/balances/compute-credit-costs.test.ts @@ -0,0 +1,84 @@ +import { describe, expect, test } from "bun:test"; +import { + type Feature, + FeatureType, + FeatureUsageType, + type FullCusEntWithFullCusProduct, +} from "@autumn/shared"; +import { computeCreditCosts } from "@/internal/balances/utils/deduction/computeCreditCosts.js"; +import type { FeatureDeduction } from "@/internal/balances/utils/types/featureDeduction.js"; + +const makeFeature = ( + id: string, + type: FeatureType, + schema: { metered_feature_id: string; credit_amount: number }[] = [], +): Feature => ({ + internal_id: `fe_${id}`, + org_id: "org_test", + created_at: 0, + env: "sandbox" as Feature["env"], + id, + name: id, + type, + config: { schema, usage_type: FeatureUsageType.Single }, + archived: false, + event_names: [], + model_markups: null, +}); + +const makeCusEnt = (id: string, feature: Feature) => + ({ id, entitlement: { feature } }) as FullCusEntWithFullCusProduct; + +const messages = makeFeature("messages", FeatureType.Metered); +const credits = makeFeature("credits", FeatureType.CreditSystem, [ + { metered_feature_id: "messages", credit_amount: 0.2 }, +]); +// Simulates a stale cached snapshot whose schema no longer includes "messages". +const staleCredits = makeFeature("credits", FeatureType.CreditSystem, [ + { metered_feature_id: "other_feature", credit_amount: 5 }, +]); + +describe("computeCreditCosts", () => { + test("applies schema ratios for parent credit systems", () => { + const deduction: FeatureDeduction = { feature: messages, deduction: 10 }; + const lookup = computeCreditCosts({ + cusEnts: [makeCusEnt("ce_msg", messages), makeCusEnt("ce_cred", credits)], + deduction, + }); + + expect(lookup("ce_msg")).toBe(1); + expect(lookup("ce_cred")).toBe(0.2); + }); + + test("token deductions use their USD cost 1:1 and ratio-map to parents", () => { + const aiCredits = makeFeature("ai_credits", FeatureType.AiCreditSystem); + const orbs = makeFeature("orbs", FeatureType.CreditSystem, [ + { metered_feature_id: "ai_credits", credit_amount: 1000 }, + ]); + const deduction: FeatureDeduction = { + feature: aiCredits, + deduction: 1, + tokens: { + usage: { modelName: "custom/m", inputTokens: 1, outputTokens: 1 }, + cost: 0.125, + }, + }; + const lookup = computeCreditCosts({ + cusEnts: [makeCusEnt("ce_ai", aiCredits), makeCusEnt("ce_orbs", orbs)], + deduction, + }); + + expect(lookup("ce_ai")).toBe(0.125); + expect(lookup("ce_orbs")).toBe(125); + }); + + test("stale schema snapshot falls back to 1 instead of failing the track", () => { + const deduction: FeatureDeduction = { feature: messages, deduction: 10 }; + const lookup = computeCreditCosts({ + cusEnts: [makeCusEnt("ce_stale", staleCredits)], + deduction, + }); + + expect(lookup("ce_stale")).toBe(1); + }); +}); From afbce358debe9c7f4776114843ac956a3ae3dbdd Mon Sep 17 00:00:00 2001 From: Charlie Lamb Date: Wed, 10 Jun 2026 14:05:39 +0100 Subject: [PATCH 35/46] chore: ai-sdk cleanup --- bun.lock | 3 + packages/ai-sdk/package.json | 3 + packages/ai-sdk/src/index.ts | 186 ++++++----------------- packages/ai-sdk/src/usage.ts | 117 ++++++++++++++ packages/ai-sdk/tests/unit/usage.test.ts | 93 ++++++++++++ 5 files changed, 261 insertions(+), 141 deletions(-) create mode 100644 packages/ai-sdk/src/usage.ts create mode 100644 packages/ai-sdk/tests/unit/usage.test.ts diff --git a/bun.lock b/bun.lock index 6d4d8dcdf..e3c8202ff 100644 --- a/bun.lock +++ b/bun.lock @@ -181,6 +181,9 @@ "packages/ai-sdk": { "name": "@useautumn/ai-sdk", "version": "0.0.1", + "dependencies": { + "@ai-sdk/provider": "^3.0.0", + }, "devDependencies": { "@types/node": "^24.9.1", "tsup": "^8.4.0", diff --git a/packages/ai-sdk/package.json b/packages/ai-sdk/package.json index e7f964477..6907b1fb9 100644 --- a/packages/ai-sdk/package.json +++ b/packages/ai-sdk/package.json @@ -24,6 +24,9 @@ "build": "rm -rf dist && tsup", "prepublishOnly": "bun run build" }, + "dependencies": { + "@ai-sdk/provider": "^3.0.0" + }, "peerDependencies": { "ai": "^6.0.116", "autumn-js": "*" diff --git a/packages/ai-sdk/src/index.ts b/packages/ai-sdk/src/index.ts index e655582ff..e786c48be 100644 --- a/packages/ai-sdk/src/index.ts +++ b/packages/ai-sdk/src/index.ts @@ -1,51 +1,39 @@ -import type { LanguageModelV3, LanguageModelV3Usage } from "@ai-sdk/provider"; -import { - type LanguageModelMiddleware, - 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"; +import type { LanguageModelV3 } from "@ai-sdk/provider"; +import { type LanguageModelMiddleware, wrapLanguageModel } from "ai"; +import { normalizeUsage, type TokenPools, type UsageLike } from "./usage.js"; -// Standalone published package: must not import from the internal @autumn/shared workspace. -const PROVIDER_SEPARATOR = "/"; +export type { TokenPools, UsageLike } from "./usage.js"; -type NestedCount = { total?: number | null } | null; - -/** - * Lenient view over the AI SDK usage shapes we accept: the nested - * `LanguageModelV3Usage`, the flat `ai` `LanguageModelUsage` (with token details), and - * legacy `promptTokens`/`completionTokens` objects. - */ -type AnyUsage = (LanguageModelV3Usage | LanguageModelUsage) & { - promptTokens?: number | NestedCount; - completionTokens?: number | NestedCount; - inputTokenDetails?: { - noCacheTokens?: number | null; - cacheReadTokens?: number | null; - cacheWriteTokens?: number | null; - } | null; - outputTokenDetails?: { - textTokens?: number | null; - reasoningTokens?: number | null; - } | null; - cachedInputTokens?: number | null; - reasoningTokens?: number | null; +type TrackTokensParams = TokenPools & { + customerId: string; + modelId: string; + featureId?: string; + entityId?: string; + properties?: Record; }; -type ExclusivePools = { - inputTokens: number; - outputTokens: number; - cacheReadTokens: number; - cacheWriteTokens: number; - reasoningTokens: number; +/** Structural view of the autumn-js client; older versions may not ship balances.trackTokens. */ +export type AutumnClient = { + balances?: { + trackTokens?: (params: TrackTokensParams) => Promise; + }; }; -const flatCount = ( - value: number | NestedCount | undefined, -): number | undefined => { - if (typeof value === "number") return value; - return value?.total ?? undefined; +export type WithAutumnOptions = { + /** Autumn SDK client instance. */ + autumn: AutumnClient; + /** The AI SDK language model to wrap. */ + model: LanguageModelV3; + /** The Autumn customer ID to attribute usage to. */ + customerId: string; + /** Override the provider prefix used in the model name (e.g. "openrouter", "custom"). Falls back to `model.provider`. */ + providerId?: string; + /** Target a specific AI credit system feature. Auto-detected if omitted. */ + featureId?: string; + /** Entity ID for entity-scoped balance tracking. */ + entityId?: string; + /** Additional properties to attach to each usage event. */ + properties?: Record; }; export const withAutumn = ({ @@ -56,103 +44,21 @@ export const withAutumn = ({ featureId, entityId, properties, -}: { - /** Autumn SDK client instance. */ - autumn: Autumn; - /** The AI SDK language model to wrap. */ - model: LanguageModelV3; - /** The Autumn customer ID to attribute usage to. */ - customerId: string; - /** Override the provider prefix used in the model name. Falls back to `model.provider`. */ - providerId?: "custom" | string; - /** Target a specific AI credit system feature. Auto-detected if omitted. */ - featureId?: string; - /** Entity ID for entity-scoped balance tracking. */ - entityId?: string; - /** Additional properties to attach to the usage event. */ - properties?: Record; -}) => { - const provider = providerId ?? model.provider; - const modelName = `${provider}${PROVIDER_SEPARATOR}${model.modelId}`; +}: WithAutumnOptions): LanguageModelV3 => { + const modelName = `${providerId ?? model.provider}/${model.modelId}`; - const required = (value: number | undefined, label: string): number => { - if (value == null) { - throw new Error( - `[Autumn] ${label} token usage was not returned by the model provider (${modelName}). This provider may not support usage tracking.`, - ); - } - return value; - }; - - const normalizeUsage = (usage: AnyUsage): ExclusivePools => { - const input = usage.inputTokens; - const output = usage.outputTokens; - - if (input != null && typeof input === "object") { - const cacheReadTokens = input.cacheRead ?? 0; - const cacheWriteTokens = input.cacheWrite ?? 0; - const textInput = - input.noCache ?? - (input.total != null - ? input.total - cacheReadTokens - cacheWriteTokens - : undefined); - const out = typeof output === "object" ? output : null; - const reasoningTokens = out?.reasoning ?? 0; - const textOutput = - out?.text ?? - (out?.total != null ? out.total - reasoningTokens : undefined); - return { - inputTokens: required(textInput, "Input"), - outputTokens: required(textOutput, "Output"), - cacheReadTokens: Math.max(0, cacheReadTokens), - cacheWriteTokens: Math.max(0, cacheWriteTokens), - reasoningTokens: Math.max(0, reasoningTokens), - }; - } - - const inputDetails = usage.inputTokenDetails; - const outputDetails = usage.outputTokenDetails; - const cacheReadTokens = - inputDetails?.cacheReadTokens ?? usage.cachedInputTokens ?? 0; - const cacheWriteTokens = inputDetails?.cacheWriteTokens ?? 0; - const reasoningTokens = - outputDetails?.reasoningTokens ?? usage.reasoningTokens ?? 0; - - const rawInput = - typeof input === "number" ? input : flatCount(usage.promptTokens); - const textInput = - inputDetails?.noCacheTokens ?? - (rawInput != null - ? rawInput - cacheReadTokens - cacheWriteTokens - : undefined); - - const rawOutput = - typeof output === "number" ? output : flatCount(usage.completionTokens); - const textOutput = - outputDetails?.textTokens ?? - (rawOutput != null ? rawOutput - reasoningTokens : undefined); - - return { - inputTokens: Math.max(0, required(textInput, "Input")), - outputTokens: Math.max(0, required(textOutput, "Output")), - cacheReadTokens: Math.max(0, cacheReadTokens), - cacheWriteTokens: Math.max(0, cacheWriteTokens), - reasoningTokens: Math.max(0, reasoningTokens), - }; - }; - - const trackUsage = async (usage: AnyUsage) => { + const trackUsage = async (usage: UsageLike) => { try { - const pools = normalizeUsage(usage); - // @ts-ignore trackTokens is generated from OpenAPI; local autumn-js types may not include it yet. - await autumn.balances.trackTokens({ + const trackTokens = autumn.balances?.trackTokens; + if (!trackTokens) { + throw new Error( + "autumn-js client does not support balances.trackTokens — upgrade autumn-js.", + ); + } + await trackTokens({ + ...normalizeUsage(usage, modelName), customerId, modelId: modelName, - inputTokens: pools.inputTokens, - outputTokens: pools.outputTokens, - cacheReadTokens: pools.cacheReadTokens, - cacheWriteTokens: pools.cacheWriteTokens, - reasoningTokens: pools.reasoningTokens, featureId, entityId, properties, @@ -166,7 +72,7 @@ export const withAutumn = ({ specificationVersion: "v3", wrapGenerate: async ({ doGenerate }) => { const result = await doGenerate(); - await trackUsage(result.usage as AnyUsage); + await trackUsage(result.usage as UsageLike); return result; }, wrapStream: async ({ doStream }) => { @@ -181,14 +87,12 @@ export const withAutumn = ({ const transformStream = new TransformStream({ transform(chunk, controller) { if (chunk.type === "finish" && chunk.usage) { - trackingPromise = trackUsage(chunk.usage as AnyUsage); + trackingPromise = trackUsage(chunk.usage as UsageLike); } controller.enqueue(chunk); }, async flush() { - if (trackingPromise) { - await trackingPromise; - } + await trackingPromise; }, }); @@ -199,5 +103,5 @@ export const withAutumn = ({ }, }; - return wrapLanguageModel({ model: model, middleware }); + return wrapLanguageModel({ model, middleware }); }; diff --git a/packages/ai-sdk/src/usage.ts b/packages/ai-sdk/src/usage.ts new file mode 100644 index 000000000..eb73be51a --- /dev/null +++ b/packages/ai-sdk/src/usage.ts @@ -0,0 +1,117 @@ +type NestedTokens = { + total?: number | null; + noCache?: number | null; + cacheRead?: number | null; + cacheWrite?: number | null; + text?: number | null; + reasoning?: number | null; +}; + +type LegacyCount = number | { total?: number | null } | null; + +/** Lenient view over AI SDK usage shapes: nested V3 counts, flat counts with token details, and legacy prompt/completion counts. */ +export type UsageLike = { + inputTokens?: number | NestedTokens | null; + outputTokens?: number | NestedTokens | null; + promptTokens?: LegacyCount; + completionTokens?: LegacyCount; + inputTokenDetails?: { + noCacheTokens?: number | null; + cacheReadTokens?: number | null; + cacheWriteTokens?: number | null; + } | null; + outputTokenDetails?: { + textTokens?: number | null; + reasoningTokens?: number | null; + } | null; + cachedInputTokens?: number | null; + reasoningTokens?: number | null; +}; + +export type TokenPools = { + inputTokens: number; + outputTokens: number; + cacheReadTokens: number; + cacheWriteTokens: number; + reasoningTokens: number; +}; + +const flatCount = (value: LegacyCount | undefined): number | undefined => + typeof value === "number" ? value : (value?.total ?? undefined); + +const isNested = ( + value: number | NestedTokens | null | undefined, +): value is NestedTokens => value != null && typeof value === "object"; + +const toParts = (usage: UsageLike) => { + const input = usage.inputTokens; + const output = usage.outputTokens; + + if (isNested(input)) { + const out = isNested(output) ? output : undefined; + return { + cacheRead: input.cacheRead ?? 0, + cacheWrite: input.cacheWrite ?? 0, + reasoning: out?.reasoning ?? 0, + textInput: input.noCache, + totalInput: input.total, + textOutput: out?.text, + totalOutput: out?.total, + }; + } + + return { + cacheRead: + usage.inputTokenDetails?.cacheReadTokens ?? usage.cachedInputTokens ?? 0, + cacheWrite: usage.inputTokenDetails?.cacheWriteTokens ?? 0, + reasoning: + usage.outputTokenDetails?.reasoningTokens ?? usage.reasoningTokens ?? 0, + textInput: usage.inputTokenDetails?.noCacheTokens, + totalInput: + typeof input === "number" ? input : flatCount(usage.promptTokens), + textOutput: usage.outputTokenDetails?.textTokens, + totalOutput: + typeof output === "number" ? output : flatCount(usage.completionTokens), + }; +}; + +const clamp = (value: number) => Math.max(0, value); + +/** Splits provider usage into exclusive token pools; throws if the provider returned no usable counts. */ +export const normalizeUsage = ( + usage: UsageLike, + modelName: string, +): TokenPools => { + const parts = toParts(usage); + + const required = ( + value: number | null | undefined, + label: string, + ): number => { + if (value == null) { + throw new Error( + `[Autumn] ${label} token usage was not returned by the model provider (${modelName}). This provider may not support usage tracking.`, + ); + } + return value; + }; + + const textInput = + parts.textInput ?? + (parts.totalInput != null + ? parts.totalInput - parts.cacheRead - parts.cacheWrite + : undefined); + const textOutput = + parts.textOutput ?? + (parts.totalOutput != null + ? parts.totalOutput - parts.reasoning + : undefined); + + return { + inputTokens: clamp(required(textInput, "Input")), + outputTokens: clamp(required(textOutput, "Output")), + cacheReadTokens: clamp(parts.cacheRead), + cacheWriteTokens: clamp(parts.cacheWrite), + reasoningTokens: clamp(parts.reasoning), + }; +}; diff --git a/packages/ai-sdk/tests/unit/usage.test.ts b/packages/ai-sdk/tests/unit/usage.test.ts new file mode 100644 index 000000000..132e39107 --- /dev/null +++ b/packages/ai-sdk/tests/unit/usage.test.ts @@ -0,0 +1,93 @@ +import { describe, expect, test } from "bun:test"; +import { normalizeUsage } from "../../src/usage.js"; + +const MODEL = "openai/gpt-test"; + +describe("normalizeUsage", () => { + test("nested V3 counts split into exclusive pools", () => { + expect( + normalizeUsage( + { + inputTokens: { total: 13, noCache: 10, cacheRead: 2, cacheWrite: 1 }, + outputTokens: { total: 7, text: 5, reasoning: 2 }, + }, + MODEL, + ), + ).toEqual({ + inputTokens: 10, + outputTokens: 5, + cacheReadTokens: 2, + cacheWriteTokens: 1, + reasoningTokens: 2, + }); + }); + + test("nested totals without breakdowns derive text pools", () => { + expect( + normalizeUsage( + { + inputTokens: { total: 13, cacheRead: 2, cacheWrite: 1 }, + outputTokens: { total: 7, reasoning: 2 }, + }, + MODEL, + ), + ).toEqual({ + inputTokens: 10, + outputTokens: 5, + cacheReadTokens: 2, + cacheWriteTokens: 1, + reasoningTokens: 2, + }); + }); + + test("flat counts with token details", () => { + expect( + normalizeUsage( + { + inputTokens: 13, + outputTokens: 7, + inputTokenDetails: { cacheReadTokens: 2, cacheWriteTokens: 1 }, + outputTokenDetails: { reasoningTokens: 2 }, + }, + MODEL, + ), + ).toEqual({ + inputTokens: 10, + outputTokens: 5, + cacheReadTokens: 2, + cacheWriteTokens: 1, + reasoningTokens: 2, + }); + }); + + test("legacy prompt/completion counts", () => { + expect( + normalizeUsage( + { promptTokens: 100, completionTokens: { total: 50 } }, + MODEL, + ), + ).toEqual({ + inputTokens: 100, + outputTokens: 50, + cacheReadTokens: 0, + cacheWriteTokens: 0, + reasoningTokens: 0, + }); + }); + + test("inconsistent totals clamp to zero instead of going negative", () => { + const pools = normalizeUsage( + { + inputTokens: { total: 1, cacheRead: 5, cacheWrite: 0 }, + outputTokens: { total: 1, reasoning: 5 }, + }, + MODEL, + ); + expect(pools.inputTokens).toBe(0); + expect(pools.outputTokens).toBe(0); + }); + + test("missing usage throws with the model name", () => { + expect(() => normalizeUsage({}, MODEL)).toThrow(/gpt-test/); + }); +}); From f6068edc1b69dd606c58d7697761d0a5320c3775 Mon Sep 17 00:00:00 2001 From: Charlie Lamb Date: Wed, 10 Jun 2026 14:28:32 +0100 Subject: [PATCH 36/46] chore: add params to track metadata --- .../track/utils/getTokenTrackParams.ts | 18 ++- .../internal/features/aiCreditSystemUtils.ts | 135 ++++++++++++------ .../creditSystems/ai-model-resolution.test.ts | 28 +++- .../unit/features/get-credit-cost.test.ts | 98 ++++++++++++- 4 files changed, 229 insertions(+), 50 deletions(-) diff --git a/server/src/internal/balances/track/utils/getTokenTrackParams.ts b/server/src/internal/balances/track/utils/getTokenTrackParams.ts index 56d00e36b..e75950bbd 100644 --- a/server/src/internal/balances/track/utils/getTokenTrackParams.ts +++ b/server/src/internal/balances/track/utils/getTokenTrackParams.ts @@ -11,7 +11,7 @@ import { import type { AutumnContext } from "@/honoUtils/HonoEnv.js"; import { getOrSetCachedFullSubject } from "@/internal/customers/cache/fullSubject/actions/getOrSetCachedFullSubject.js"; import { getOrSetCachedFullCustomer } from "@/internal/customers/cusUtils/fullCustomerCacheUtils/getOrSetCachedFullCustomer.js"; -import { getModelCreditCost } from "@/internal/features/aiCreditSystemUtils.js"; +import { getModelCreditCostBreakdown } from "@/internal/features/aiCreditSystemUtils.js"; import { isFullSubjectRolloutEnabled } from "@/internal/misc/rollouts/fullSubjectRolloutUtils.js"; import type { FeatureDeduction } from "../../utils/types/featureDeduction.js"; @@ -118,7 +118,7 @@ export const getTokenTrackParams = async ({ entityId: input.entity_id, }); - const cost = await getModelCreditCost({ + const pricing = await getModelCreditCostBreakdown({ modelName: input.model_id, creditSystem: aiCreditFeature, input: input.input_tokens, @@ -129,6 +129,7 @@ export const getTokenTrackParams = async ({ audioOutput: input.audio_output_tokens, reasoning: input.reasoning_tokens, }); + const cost = pricing.cost; const featureDeductions: FeatureDeduction[] = [ { @@ -161,6 +162,19 @@ export const getTokenTrackParams = async ({ audio_output_tokens: input.audio_output_tokens, reasoning_tokens: input.reasoning_tokens, cost, + base_cost: pricing.baseCost, + markup: pricing.markup, + markup_source: pricing.markupSource, + tier_applied: pricing.tierApplied, + rates: { + input: pricing.rates.input, + output: pricing.rates.output, + cache_read: pricing.rates.cacheRead, + cache_write: pricing.rates.cacheWrite, + audio_input: pricing.rates.audioInput, + audio_output: pricing.rates.audioOutput, + reasoning: pricing.rates.reasoning, + }, }, idempotency_key: input.idempotency_key, overage_behavior: input.overage_behavior, diff --git a/server/src/internal/features/aiCreditSystemUtils.ts b/server/src/internal/features/aiCreditSystemUtils.ts index 77392e1f3..b0eb44de2 100644 --- a/server/src/internal/features/aiCreditSystemUtils.ts +++ b/server/src/internal/features/aiCreditSystemUtils.ts @@ -7,7 +7,6 @@ import { type ModelsDevModel, type ModelsDevProvider, RecaseError, - resolveInheritedMarkup, splitModelId, } from "@autumn/shared"; import { Decimal } from "decimal.js"; @@ -78,7 +77,7 @@ const resolveModel = ({ const getEffectiveCost = ( cost: ModelsDevCost, totalInputTokens: number, -): ModelsDevCost => { +): { effective: ModelsDevCost; tierApplied: boolean } => { if (cost.tiers?.length) { let chosen: ModelsDevCostTier | undefined; for (const tier of cost.tiers) { @@ -91,25 +90,51 @@ const getEffectiveCost = ( } if (chosen) { return { - ...cost, - input: chosen.input, - output: chosen.output, - cache_read: chosen.cache_read ?? cost.cache_read, - cache_write: chosen.cache_write ?? cost.cache_write, + effective: { + ...cost, + input: chosen.input, + output: chosen.output, + cache_read: chosen.cache_read ?? cost.cache_read, + cache_write: chosen.cache_write ?? cost.cache_write, + }, + tierApplied: true, }; } - return cost; + return { effective: cost, tierApplied: false }; } if (cost.context_over_200k && totalInputTokens > LARGE_CONTEXT_THRESHOLD) { return { - ...cost, - input: cost.context_over_200k.input, - output: cost.context_over_200k.output, - cache_read: cost.context_over_200k.cache_read ?? cost.cache_read, - cache_write: cost.context_over_200k.cache_write ?? cost.cache_write, + effective: { + ...cost, + input: cost.context_over_200k.input, + output: cost.context_over_200k.output, + cache_read: cost.context_over_200k.cache_read ?? cost.cache_read, + cache_write: cost.context_over_200k.cache_write ?? cost.cache_write, + }, + tierApplied: true, }; } - return cost; + return { effective: cost, tierApplied: false }; +}; + +/** Effective per-token rates ($/M) used for a charge, after tier overlays and fallbacks. */ +export type ModelCostRates = { + input: number; + output: number; + cacheRead: number; + cacheWrite: number; + audioInput: number; + audioOutput: number; + reasoning: number; +}; + +export type ModelCostBreakdown = { + cost: number; + baseCost: number; + markup: number; + markupSource: "model" | "provider" | "default" | "none"; + tierApplied: boolean; + rates: ModelCostRates; }; const computeCost = ({ @@ -120,7 +145,7 @@ const computeCost = ({ cost: ModelsDevCost; tokens: TokenInput; markup: number; -}): number => { +}): { cost: number; baseCost: number; tierApplied: boolean; rates: ModelCostRates } => { const cacheRead = tokens.cacheRead ?? 0; const cacheWrite = tokens.cacheWrite ?? 0; const audioInput = tokens.audioInput ?? 0; @@ -128,28 +153,37 @@ const computeCost = ({ const reasoning = tokens.reasoning ?? 0; const totalInput = tokens.input + cacheRead + cacheWrite; - const effective = getEffectiveCost(cost, totalInput); - const inputRate = effective.input; - const outputRate = effective.output; + const { effective, tierApplied } = getEffectiveCost(cost, totalInput); // Pools without a published rate fall back to the base text rate. - const cacheReadRate = effective.cache_read ?? inputRate; - const cacheWriteRate = effective.cache_write ?? inputRate; - const audioInputRate = effective.input_audio ?? inputRate; - const audioOutputRate = effective.output_audio ?? outputRate; - const reasoningRate = effective.reasoning ?? outputRate; + const rates: ModelCostRates = { + input: effective.input, + output: effective.output, + cacheRead: effective.cache_read ?? effective.input, + cacheWrite: effective.cache_write ?? effective.input, + audioInput: effective.input_audio ?? effective.input, + audioOutput: effective.output_audio ?? effective.output, + reasoning: effective.reasoning ?? effective.output, + }; - return new Decimal(inputRate) + const baseCost = new Decimal(rates.input) .mul(tokens.input) - .add(new Decimal(outputRate).mul(tokens.output)) - .add(new Decimal(cacheReadRate).mul(cacheRead)) - .add(new Decimal(cacheWriteRate).mul(cacheWrite)) - .add(new Decimal(audioInputRate).mul(audioInput)) - .add(new Decimal(audioOutputRate).mul(audioOutput)) - .add(new Decimal(reasoningRate).mul(reasoning)) - .div(1_000_000) - .mul(new Decimal(1).add(new Decimal(markup).div(100))) - .toNumber(); + .add(new Decimal(rates.output).mul(tokens.output)) + .add(new Decimal(rates.cacheRead).mul(cacheRead)) + .add(new Decimal(rates.cacheWrite).mul(cacheWrite)) + .add(new Decimal(rates.audioInput).mul(audioInput)) + .add(new Decimal(rates.audioOutput).mul(audioOutput)) + .add(new Decimal(rates.reasoning).mul(reasoning)) + .div(1_000_000); + + return { + cost: baseCost + .mul(new Decimal(1).add(new Decimal(markup).div(100))) + .toNumber(), + baseCost: baseCost.toNumber(), + tierApplied, + rates, + }; }; const resolveAiMarkup = ({ @@ -160,38 +194,41 @@ const resolveAiMarkup = ({ modelName: string; creditSystem: Feature; modelMarkup?: { markup?: number | null } | null; -}) => { +}): { markup: number; source: ModelCostBreakdown["markupSource"] } => { if (modelMarkup?.markup != null) { - return modelMarkup.markup; + return { markup: modelMarkup.markup, source: "model" }; } const { provider } = splitModelId(modelName); const providerMarkup = provider ? creditSystem.config?.provider_markups?.[provider]?.markup : undefined; + if (providerMarkup != null) { + return { markup: providerMarkup, source: "provider" }; + } - return ( - resolveInheritedMarkup({ - providerMarkup, - defaultMarkup: creditSystem.config?.default_markup, - }) ?? 0 - ); + const defaultMarkup = creditSystem.config?.default_markup; + if (defaultMarkup != null) { + return { markup: defaultMarkup, source: "default" }; + } + + return { markup: 0, source: "none" }; }; -export const getModelCreditCost = async ({ +export const getModelCreditCostBreakdown = async ({ modelName, creditSystem, ...tokens }: { modelName: string; creditSystem: Feature; -} & TokenInput): Promise => { +} & TokenInput): Promise => { const markups = creditSystem.model_markups || {}; const pricingData = await getModelsDevPricing(); const resolved = resolveModel({ modelName, pricingData }); const markupEntry = markups[modelName]; - const markup = resolveAiMarkup({ + const { markup, source } = resolveAiMarkup({ modelName, creditSystem, modelMarkup: markupEntry, @@ -207,16 +244,22 @@ export const getModelCreditCost = async ({ data: { modelName }, }); } - return computeCost({ + const computed = computeCost({ cost: { input: markupEntry.input_cost, output: markupEntry.output_cost }, tokens: { input: tokens.input, output: tokens.output }, markup, }); + return { ...computed, markup, markupSource: source }; } - return computeCost({ + const computed = computeCost({ cost: resolved.model.cost, tokens, markup, }); + return { ...computed, markup, markupSource: source }; }; + +export const getModelCreditCost = async ( + args: { modelName: string; creditSystem: Feature } & TokenInput, +): Promise => (await getModelCreditCostBreakdown(args)).cost; diff --git a/server/tests/advanced/creditSystems/ai-model-resolution.test.ts b/server/tests/advanced/creditSystems/ai-model-resolution.test.ts index d63689846..bb6830d2d 100644 --- a/server/tests/advanced/creditSystems/ai-model-resolution.test.ts +++ b/server/tests/advanced/creditSystems/ai-model-resolution.test.ts @@ -89,7 +89,7 @@ mock.module("@/internal/features/utils/getModelPricing.js", () => ({ getModelsDevPricing: async () => pricingData, })); -const { getModelCreditCost } = await import( +const { getModelCreditCost, getModelCreditCostBreakdown } = await import( "@/internal/features/aiCreditSystemUtils.js" ); @@ -247,4 +247,30 @@ describe("computeCost — token pools", () => { }); expect(cost).toBeCloseTo(((5 * 1000 + 25 * 500) / PER_MILLION) * 1.5, 10); }); + + test("breakdown reports tier_applied and the tier rates actually used", async () => { + const above = await getModelCreditCostBreakdown({ + modelName: "openai/gpt-5", + creditSystem: makeFeature(), + input: 300_000, + output: 1000, + }); + expect(above.tierApplied).toBe(true); + expect(above.rates.input).toBe(2); + expect(above.rates.cacheRead).toBe(1); + expect(above.baseCost).toBeCloseTo( + (2 * 300_000 + 4 * 1000) / PER_MILLION, + 10, + ); + expect(above.cost).toBe(above.baseCost); + + const below = await getModelCreditCostBreakdown({ + modelName: "openai/gpt-5", + creditSystem: makeFeature(), + input: 1000, + output: 1000, + }); + expect(below.tierApplied).toBe(false); + expect(below.rates.input).toBe(1); + }); }); diff --git a/server/tests/unit/features/get-credit-cost.test.ts b/server/tests/unit/features/get-credit-cost.test.ts index 7fb7c5553..5f64b31fa 100644 --- a/server/tests/unit/features/get-credit-cost.test.ts +++ b/server/tests/unit/features/get-credit-cost.test.ts @@ -5,7 +5,10 @@ import { FeatureType, FeatureUsageType, } from "@autumn/shared"; -import { getModelCreditCost } from "@/internal/features/aiCreditSystemUtils.js"; +import { + getModelCreditCost, + getModelCreditCostBreakdown, +} from "@/internal/features/aiCreditSystemUtils.js"; import { getCreditCost } from "@/internal/features/creditSystemUtils.js"; // Uses custom/* models so pricing resolves offline (no models.dev fetch). @@ -81,3 +84,96 @@ describe("getModelCreditCost — token pricing", () => { }); }); }); + +describe("getModelCreditCostBreakdown — pricing audit trail", () => { + test("records base cost, markup source, and effective rates", async () => { + const withMarkup: Feature = { + ...aiCreditFeature, + model_markups: { + [CUSTOM_MODEL]: { markup: 50, input_cost: 1000, output_cost: 2000 }, + }, + }; + const breakdown = await getModelCreditCostBreakdown({ + modelName: CUSTOM_MODEL, + creditSystem: withMarkup, + input: 1000, + output: 500, + }); + + expect(breakdown.baseCost).toBeCloseTo(2.0, 10); + expect(breakdown.cost).toBeCloseTo(3.0, 10); + expect(breakdown.markup).toBe(50); + expect(breakdown.markupSource).toBe("model"); + expect(breakdown.tierApplied).toBe(false); + expect(breakdown.rates.input).toBe(1000); + expect(breakdown.rates.output).toBe(2000); + // Unpublished pools fall back to the text rates. + expect(breakdown.rates.cacheRead).toBe(1000); + expect(breakdown.rates.reasoning).toBe(2000); + }); + + test("explicit markup 0 reports source model; no markup anywhere reports none", async () => { + const explicitZero = await getModelCreditCostBreakdown({ + modelName: CUSTOM_MODEL, + creditSystem: aiCreditFeature, + input: 1000, + output: 500, + }); + expect(explicitZero.markup).toBe(0); + expect(explicitZero.markupSource).toBe("model"); + + const unconfigured = await getModelCreditCostBreakdown({ + modelName: CUSTOM_MODEL, + creditSystem: { + ...aiCreditFeature, + model_markups: { + [CUSTOM_MODEL]: { input_cost: 1000, output_cost: 2000 }, + }, + }, + input: 1000, + output: 500, + }); + expect(unconfigured.markup).toBe(0); + expect(unconfigured.markupSource).toBe("none"); + }); + + test("reports provider and default markup sources", async () => { + const noModelMarkup: Feature = { + ...aiCreditFeature, + config: { + schema: [], + usage_type: FeatureUsageType.Single, + default_markup: 10, + provider_markups: { custom: { markup: 20 } }, + }, + model_markups: { + [CUSTOM_MODEL]: { input_cost: 1000, output_cost: 2000 }, + }, + }; + + const provider = await getModelCreditCostBreakdown({ + modelName: CUSTOM_MODEL, + creditSystem: noModelMarkup, + input: 1000, + output: 500, + }); + expect(provider.markup).toBe(20); + expect(provider.markupSource).toBe("provider"); + + const defaultOnly = await getModelCreditCostBreakdown({ + modelName: CUSTOM_MODEL, + creditSystem: { + ...noModelMarkup, + config: { + schema: [], + usage_type: FeatureUsageType.Single, + default_markup: 10, + }, + }, + input: 1000, + output: 500, + }); + expect(defaultOnly.markup).toBe(10); + expect(defaultOnly.markupSource).toBe("default"); + }); +}); From c443b3ae39a73dfe63fa2b054e57a6d3a9921e04 Mon Sep 17 00:00:00 2001 From: Charlie Lamb Date: Wed, 10 Jun 2026 14:57:28 +0100 Subject: [PATCH 37/46] chore: ui cleanup --- .../credit-systems/components/AiCreditSchemaTable.tsx | 3 +-- .../products/features/feature-list/CreditListColumns.tsx | 6 +++--- 2 files changed, 4 insertions(+), 5 deletions(-) diff --git a/vite/src/views/products/features/credit-systems/components/AiCreditSchemaTable.tsx b/vite/src/views/products/features/credit-systems/components/AiCreditSchemaTable.tsx index d4544cf27..6ad09edc1 100644 --- a/vite/src/views/products/features/credit-systems/components/AiCreditSchemaTable.tsx +++ b/vite/src/views/products/features/credit-systems/components/AiCreditSchemaTable.tsx @@ -255,7 +255,6 @@ export function AiCreditSchemaTable({
Markup %
)} diff --git a/vite/src/views/products/features/feature-list/CreditListColumns.tsx b/vite/src/views/products/features/feature-list/CreditListColumns.tsx index bbf35d616..0a7abf6a1 100644 --- a/vite/src/views/products/features/feature-list/CreditListColumns.tsx +++ b/vite/src/views/products/features/feature-list/CreditListColumns.tsx @@ -56,7 +56,7 @@ export const createCreditListColumns = ( }, { header: "Type", - size: 120, + size: 160, accessorKey: "type", cell: ({ row }: { row: Row }) => { const isAi = isAiCreditSystem(row.original.type); @@ -65,12 +65,12 @@ export const createCreditListColumns = ( {isAi ? ( <> - AI + AI Credit System ) : ( <> - Standard + Credit System )}
From 36311a71cf4323ccca2aaeb796996292d8d47dea Mon Sep 17 00:00:00 2001 From: Ridhwan Hussain <73362400+TheUntraceable@users.noreply.github.com.> Date: Wed, 10 Jun 2026 15:10:03 +0100 Subject: [PATCH 38/46] fix: allow explicit setting of a models markup to 0 --- .../features/credit-systems/components/EditableNumberCell.tsx | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/vite/src/views/products/features/credit-systems/components/EditableNumberCell.tsx b/vite/src/views/products/features/credit-systems/components/EditableNumberCell.tsx index 39dee2cbc..d949b212e 100644 --- a/vite/src/views/products/features/credit-systems/components/EditableNumberCell.tsx +++ b/vite/src/views/products/features/credit-systems/components/EditableNumberCell.tsx @@ -31,9 +31,7 @@ export function EditableNumberCell({ const [local, setLocal] = useState(""); const [focused, setFocused] = useState(false); - const hasValue = allowUndefined - ? currentValue != null && currentValue !== 0 - : currentValue != null; + const hasValue = currentValue != null; const displayed = focused ? local : hasValue ? String(currentValue) : ""; return ( From d5e5cfc5ff0697e7e8a17a10d3157489d4af0686 Mon Sep 17 00:00:00 2001 From: Ridhwan Hussain <73362400+TheUntraceable@users.noreply.github.com.> Date: Wed, 10 Jun 2026 15:46:14 +0100 Subject: [PATCH 39/46] fix: let people do negative markups for free models --- server/src/internal/features/featureUtils.ts | 10 ++++++---- shared/api/features/apiFeatureV1.ts | 5 +++-- shared/api/features/crud/common/baseFeatureParamsV1.ts | 4 ++-- shared/drizzle/0002_shocking_wong.sql | 1 + .../models/featureModels/featureConfig/creditConfig.ts | 6 +++--- 5 files changed, 15 insertions(+), 11 deletions(-) create mode 100644 shared/drizzle/0002_shocking_wong.sql diff --git a/server/src/internal/features/featureUtils.ts b/server/src/internal/features/featureUtils.ts index cf8eb5311..35e66c579 100644 --- a/server/src/internal/features/featureUtils.ts +++ b/server/src/internal/features/featureUtils.ts @@ -80,9 +80,10 @@ export const validateCreditSystem = ( const defaultMarkup = newConfig.default_markup; if (defaultMarkup != null) { const parsedDefaultMarkup = Number(defaultMarkup); - if (Number.isNaN(parsedDefaultMarkup) || parsedDefaultMarkup < 0) { + if (Number.isNaN(parsedDefaultMarkup) || parsedDefaultMarkup < -100) { throw new RecaseError({ - message: "Default markup should be a non-negative number", + message: + "Default markup must be -100 or greater (-100 makes usage free)", code: ErrCode.InvalidFeature, statusCode: 400, }); @@ -94,9 +95,10 @@ export const validateCreditSystem = ( if (providerMarkups != null) { for (const [provider, entry] of Object.entries(providerMarkups)) { const markup = Number(entry?.markup); - if (!provider || Number.isNaN(markup) || markup < 0) { + if (!provider || Number.isNaN(markup) || markup < -100) { throw new RecaseError({ - message: "Provider markups must be non-negative numbers", + message: + "Provider markups must be -100 or greater (-100 makes usage free)", code: ErrCode.InvalidFeature, statusCode: 400, }); diff --git a/shared/api/features/apiFeatureV1.ts b/shared/api/features/apiFeatureV1.ts index 9ff81d6f4..0f63a62c4 100644 --- a/shared/api/features/apiFeatureV1.ts +++ b/shared/api/features/apiFeatureV1.ts @@ -50,8 +50,9 @@ export const ApiFeatureV1Schema = z.object({ description: "Per-model markup overrides for AI credit systems.", }), - default_markup: z.number().min(0).optional().meta({ - description: "Default percentage markup for AI credit systems.", + default_markup: z.number().min(-100).optional().meta({ + description: + "Default percentage markup for AI credit systems. Use -100 to make usage free.", }), provider_markups: ProviderMarkupsSchema.optional().meta({ diff --git a/shared/api/features/crud/common/baseFeatureParamsV1.ts b/shared/api/features/crud/common/baseFeatureParamsV1.ts index e9a103ed7..7394be049 100644 --- a/shared/api/features/crud/common/baseFeatureParamsV1.ts +++ b/shared/api/features/crud/common/baseFeatureParamsV1.ts @@ -55,9 +55,9 @@ export const BaseFeatureV1ParamsSchema = z.object({ "Per-model markup overrides for AI credit systems. Maps model IDs to their markup configuration.", }), - default_markup: z.number().min(0).optional().meta({ + default_markup: z.number().min(-100).optional().meta({ description: - "Default percentage markup for this AI credit system. Used when no model or provider markup applies.", + "Default percentage markup for this AI credit system. Used when no model or provider markup applies. Use -100 to make usage free.", }), provider_markups: ProviderMarkupsSchema.optional().meta({ diff --git a/shared/drizzle/0002_shocking_wong.sql b/shared/drizzle/0002_shocking_wong.sql new file mode 100644 index 000000000..be0692e60 --- /dev/null +++ b/shared/drizzle/0002_shocking_wong.sql @@ -0,0 +1 @@ +ALTER TABLE "features" ADD COLUMN "model_markups" jsonb DEFAULT null; \ No newline at end of file diff --git a/shared/models/featureModels/featureConfig/creditConfig.ts b/shared/models/featureModels/featureConfig/creditConfig.ts index 69f51666a..63d524d42 100644 --- a/shared/models/featureModels/featureConfig/creditConfig.ts +++ b/shared/models/featureModels/featureConfig/creditConfig.ts @@ -8,7 +8,7 @@ export const CreditSchemaItemSchema = z.object({ }); const MarkupEntrySchema = z.object({ - markup: z.number().min(0), // percentage markup, e.g. 20 for 20% + markup: z.number().min(-100), // percentage markup, e.g. 20 for 20%, -100 for free }); export const ProviderMarkupsSchema = z @@ -26,7 +26,7 @@ export const CreditSystemConfigSchema = z.object({ }), ), usage_type: z.nativeEnum(FeatureUsageType), - default_markup: z.number().min(0).optional(), + default_markup: z.number().min(-100).optional(), provider_markups: ProviderMarkupsSchema, }); @@ -34,7 +34,7 @@ export const ModelMarkupsSchema = z .record( z.string(), // Represents the model name in "provider/model" format, e.g. "anthropic/claude-2" MarkupEntrySchema.extend({ - markup: z.number().min(0).optional(), // Omit to inherit provider/global markup + markup: z.number().min(-100).optional(), // Omit to inherit provider/global markup input_cost: z.number().min(0).optional(), // $/M tokens, required for custom/ models output_cost: z.number().min(0).optional(), // $/M tokens, required for custom/ models }), From ac26721eea4398d6c2a652f8c7b722d795ae5ea1 Mon Sep 17 00:00:00 2001 From: Charlie Lamb Date: Wed, 10 Jun 2026 16:08:04 +0100 Subject: [PATCH 40/46] fix: zero credit cost deducts nothing instead of defaulting to 1 A -100% markup model computes cost 0, but the deduction scripts coerced credit_cost 0 -> 1, silently charging 1 credit unit per free call. Keep the nil -> 1 default; treat an explicit 0 as free (consume the usage, leave main and rollover balances untouched). --- .../fullSubjectDeduction/deductFromRolloversV2.lua | 8 +++++++- .../fullSubjectDeduction/runDeductionOnContextV2.lua | 7 ++++++- 2 files changed, 13 insertions(+), 2 deletions(-) diff --git a/server/src/_luaScriptsV2/fullSubjectDeduction/deductFromRolloversV2.lua b/server/src/_luaScriptsV2/fullSubjectDeduction/deductFromRolloversV2.lua index 7ee5f7f47..60b8acdc8 100644 --- a/server/src/_luaScriptsV2/fullSubjectDeduction/deductFromRolloversV2.lua +++ b/server/src/_luaScriptsV2/fullSubjectDeduction/deductFromRolloversV2.lua @@ -74,9 +74,15 @@ local function deduct_from_rollovers(params) local rollover_id = rollover_obj.id local credit_cost = rollover_obj.credit_cost - if is_nil(credit_cost) or credit_cost == 0 then + if is_nil(credit_cost) then credit_cost = 1 end + if credit_cost == 0 then + -- Zero credit cost (e.g. -100% markup AI model): usage is free, leave rollovers untouched. + logger.log(" Rollover %s credit_cost=0 - free deduction, skipping", rollover_id) + remaining = 0 + break + end local rollover_data = context.rollovers[rollover_id] if not rollover_data then diff --git a/server/src/_luaScriptsV2/fullSubjectDeduction/runDeductionOnContextV2.lua b/server/src/_luaScriptsV2/fullSubjectDeduction/runDeductionOnContextV2.lua index e092ee3a3..1c590a479 100644 --- a/server/src/_luaScriptsV2/fullSubjectDeduction/runDeductionOnContextV2.lua +++ b/server/src/_luaScriptsV2/fullSubjectDeduction/runDeductionOnContextV2.lua @@ -50,7 +50,7 @@ local function process_deduction_pass(params) local ent_id = ent_obj.customer_entitlement_id local credit_cost = ent_obj.credit_cost local ent_feature_id = ent_obj.feature_id - if credit_cost == cjson.null or credit_cost == nil or credit_cost == 0 then + if credit_cost == cjson.null or credit_cost == nil then credit_cost = 1 end @@ -91,6 +91,11 @@ local function process_deduction_pass(params) if not should_process then logger.log("%s skipping %s - usage_allowed=false or not in context", pass_name, ent_id) + elseif credit_cost == 0 then + -- Zero credit cost (e.g. -100% markup AI model): the usage is free. + -- Consume the requested amount without touching any balance. + logger.log("%s ent %s credit_cost=0 - free deduction, no balance change", pass_name, ent_id) + remaining_amount = 0 else local deducted = deduct_from_main_balance({ context = context, From a14f6e43960c121f4ed6b4bba49318fdb34fa8e6 Mon Sep 17 00:00:00 2001 From: Charlie Lamb Date: Wed, 10 Jun 2026 16:19:31 +0100 Subject: [PATCH 41/46] chore: add beta label --- vite/src/components/v2/buttons/GroupedTabButton.tsx | 2 +- .../credit-systems/components/CreditSystemSchema.tsx | 11 ++++++++++- 2 files changed, 11 insertions(+), 2 deletions(-) diff --git a/vite/src/components/v2/buttons/GroupedTabButton.tsx b/vite/src/components/v2/buttons/GroupedTabButton.tsx index 2c8a6bac8..611b2b418 100644 --- a/vite/src/components/v2/buttons/GroupedTabButton.tsx +++ b/vite/src/components/v2/buttons/GroupedTabButton.tsx @@ -6,7 +6,7 @@ interface GroupedTabButtonProps { onValueChange: (value: string) => void; options: Array<{ value: string; - label: string; + label: React.ReactNode; icon?: React.ReactNode; }>; className?: string; diff --git a/vite/src/views/products/features/credit-systems/components/CreditSystemSchema.tsx b/vite/src/views/products/features/credit-systems/components/CreditSystemSchema.tsx index d1abcb963..f67bbad95 100644 --- a/vite/src/views/products/features/credit-systems/components/CreditSystemSchema.tsx +++ b/vite/src/views/products/features/credit-systems/components/CreditSystemSchema.tsx @@ -1,6 +1,7 @@ import { FeatureType, isAiCreditSystem } from "@autumn/shared"; import { useStore } from "@tanstack/react-form"; import { useMemo } from "react"; +import { BetaBadge } from "@/components/v2/badges/BetaBadge"; import { GroupedTabButton } from "@/components/v2/buttons/GroupedTabButton"; import { SheetSection } from "@/components/v2/sheets/SharedSheetComponents"; import type { CreditSystemFormInstance } from "../hooks/useCreditSystemForm"; @@ -44,7 +45,15 @@ export function CreditSystemSchema({ const modeOptions = useMemo( () => [ { value: "classic", label: "Classic" }, - { value: "ai", label: "AI" }, + { + value: "ai", + label: ( + + AI + + + ), + }, ], [], ); From 8f0218050d1cfe3e128e0723111e09dd7ec18637 Mon Sep 17 00:00:00 2001 From: Charlie Lamb Date: Wed, 10 Jun 2026 17:03:45 +0100 Subject: [PATCH 42/46] docs: update ai credit system docs --- .../balances/trackTokens.mdx | 41 +++++++---- .../api-reference/balances/trackTokens.mdx | 69 ++++++++++++++----- .../customers/tracking-usage.mdx | 13 +++- .../external-providers/ai-sdk.mdx | 4 ++ .../modelling-pricing/credit-systems.mdx | 35 ++++++++-- 5 files changed, 126 insertions(+), 36 deletions(-) diff --git a/apps/docs/api-reference-generator/balances/trackTokens.mdx b/apps/docs/api-reference-generator/balances/trackTokens.mdx index f523f8d11..fa947801c 100644 --- a/apps/docs/api-reference-generator/balances/trackTokens.mdx +++ b/apps/docs/api-reference-generator/balances/trackTokens.mdx @@ -3,15 +3,15 @@ title: "Track Token Usage" openapi: "openapi POST /v1/balances.track_tokens" --- -import { DynamicParamField } from "/components/dynamic-param-field.jsx"; -import { DynamicResponseExample } from "/components/dynamic-response-example.jsx"; -import { DynamicResponseField } from "/components/dynamic-response-field.jsx"; +import { DynamicParamField } from "/snippets/dynamic-param-field.jsx"; +import { DynamicResponseField } from "/snippets/dynamic-response-field.jsx"; +import { DynamicResponseExample } from "/snippets/dynamic-response-example.jsx"; Track AI token usage against a customer's AI credit system balance. Converts token counts to a dollar cost using [Models.dev](https://models.dev) pricing and your configured markup, then deducts from the customer's credit balance. -The `model_id` must use `provider/model` format, matching the provider and model keys from [Models.dev](https://models.dev). For providers with nested model paths (like OpenRouter), include the full path: `openrouter/anthropic/claude-opus-4.6`. +The `model_id` must use `provider/model` format, matching the provider and model keys from [Models.dev](https://models.dev). For providers with nested model paths (like OpenRouter), include the full path: `openrouter/anthropic/claude-opus-4.6`. The first path segment is the provider key used for `providerMarkups` lookup. ### Common Use Cases @@ -26,16 +26,15 @@ await autumn.balances.trackTokens({ }); ``` -```typescript OpenAI +```typescript With cache + reasoning await autumn.balances.trackTokens({ customerId: "cus_123", - modelId: "openai/gpt-4o", - inputTokens: 500, - outputTokens: 200, - properties: { - conversation_id: "conv_abc123", - prompt_type: "summarization" - } + modelId: "anthropic/claude-opus-4-6", + inputTokens: 800, // excludes the cached tokens below + outputTokens: 350, // excludes the reasoning tokens below + cacheReadTokens: 1000, + cacheWriteTokens: 200, + reasoningTokens: 150 }); ``` @@ -59,3 +58,21 @@ await autumn.balances.trackTokens({ ``` + +### Token Pools + +Each token parameter is an exclusive pool — no token should be counted in more than one. `input_tokens` is non-cached text input only (cached tokens go in `cache_read_tokens` / `cache_write_tokens`), and `output_tokens` is text output only (reasoning tokens go in `reasoning_tokens`, audio in `audio_input_tokens` / `audio_output_tokens`). Each pool is billed at the model's published rate for that pool, falling back to the text input/output rate when the model has none. + + + If you pass a provider's raw totals (e.g. OpenAI's `prompt_tokens` and `completion_tokens`), subtract the cache and reasoning counts first — otherwise those tokens are billed twice. The [`@useautumn/ai-sdk` wrapper](/documentation/external-providers/ai-sdk) does this normalization for you. + + +### Markup Resolution + +Markups are optional — the credit system's default markup applies unless overridden per provider or per model. With no markups set, the Models.dev base cost is charged as-is. A markup of `-100` makes the model free — the usage event is still recorded, but nothing is deducted. See [AI Credit Systems](/documentation/modelling-pricing/credit-systems#ai-credit-systems) for configuration. + +The recorded event's `properties` include the full pricing breakdown: `cost`, `base_cost`, `markup`, `markup_source` (`model`, `provider`, or `default`), `tier_applied` (whether large-context tier pricing applied), and the per-pool `rates` used. + + + `feature_id` is auto-detected when the customer has exactly one AI credit system. The request fails if the customer has none, or has more than one and `feature_id` is omitted. + diff --git a/apps/docs/mintlify/api-reference/balances/trackTokens.mdx b/apps/docs/mintlify/api-reference/balances/trackTokens.mdx index a258f5df3..28c892213 100644 --- a/apps/docs/mintlify/api-reference/balances/trackTokens.mdx +++ b/apps/docs/mintlify/api-reference/balances/trackTokens.mdx @@ -3,15 +3,15 @@ title: "Track Token Usage" openapi: "openapi POST /v1/balances.track_tokens" --- -import { DynamicParamField } from "/components/dynamic-param-field.jsx"; -import { DynamicResponseExample } from "/components/dynamic-response-example.jsx"; -import { DynamicResponseField } from "/components/dynamic-response-field.jsx"; +import { DynamicParamField } from "/snippets/dynamic-param-field.jsx"; +import { DynamicResponseField } from "/snippets/dynamic-response-field.jsx"; +import { DynamicResponseExample } from "/snippets/dynamic-response-example.jsx"; Track AI token usage against a customer's AI credit system balance. Converts token counts to a dollar cost using [Models.dev](https://models.dev) pricing and your configured markup, then deducts from the customer's credit balance. -The `model_id` must use `provider/model` format, matching the provider and model keys from [Models.dev](https://models.dev). For providers with nested model paths (like OpenRouter), include the full path: `openrouter/anthropic/claude-opus-4.6`. +The `model_id` must use `provider/model` format, matching the provider and model keys from [Models.dev](https://models.dev). For providers with nested model paths (like OpenRouter), include the full path: `openrouter/anthropic/claude-opus-4.6`. The first path segment is the provider key used for provider-level markup lookup. ### Common Use Cases @@ -26,23 +26,22 @@ await autumn.balances.trackTokens({ }); ``` -```typescript OpenAI +```typescript With cache + reasoning await autumn.balances.trackTokens({ customerId: "cus_123", - modelId: "openai/gpt-4o", - inputTokens: 500, - outputTokens: 200, - properties: { - conversation_id: "conv_abc123", - prompt_type: "summarization" - } + modelId: "anthropic/claude-opus-4-6", + inputTokens: 800, // excludes the cached tokens below + outputTokens: 350, // excludes the reasoning tokens below + cacheReadTokens: 1000, + cacheWriteTokens: 200, + reasoningTokens: 150 }); ``` ```typescript OpenRouter (nested path) await autumn.balances.trackTokens({ customerId: "cus_123", - modelId: "openrouter/anthropic/claude-opus-4.6", + modelId: "openrouter/anthropic/claude-opus-4-6", inputTokens: 2000, outputTokens: 1000 }); @@ -60,6 +59,22 @@ await autumn.balances.trackTokens({ +### Token Pools + +Each token parameter is an exclusive pool — no token should be counted in more than one. Each pool is billed at the model's published rate for that pool, falling back to the text input/output rate when the model has none. + + + If you pass a provider's raw totals (e.g. OpenAI's `prompt_tokens` and `completion_tokens`), subtract the cache and reasoning counts first — otherwise those tokens are billed twice. The [`@useautumn/ai-sdk` wrapper](/documentation/external-providers/ai-sdk) does this normalization for you. + + +### Markup Resolution + +Markups are optional — the credit system's default markup applies unless overridden per provider or per model. With no markups set, the Models.dev base cost is charged as-is. A markup of `-100` makes the model free — the usage event is still recorded, but nothing is deducted. See [AI Credit Systems](/documentation/modelling-pricing/credit-systems#ai-credit-systems) for configuration. + + + `feature_id` is auto-detected when the customer has exactly one AI credit system. The request fails if the customer has none, or has more than one and `feature_id` is omitted. + + ### Body Parameters @@ -71,15 +86,35 @@ await autumn.balances.trackTokens({ - Number of input tokens consumed. + Number of non-cached text input tokens consumed. Exclusive of the cache and audio token pools. - Number of output tokens consumed. + Number of text output tokens consumed. Exclusive of the reasoning and audio output pools. + + + + Number of cached input tokens read, billed at the model's cache read rate. + + + + Number of input tokens written to the cache, billed at the model's cache write rate. + + + + Number of reasoning tokens generated, billed at the model's reasoning rate (falls back to the output rate). + + + + Number of audio input tokens consumed, billed at the model's audio input rate (falls back to the input rate). + + + + Number of audio output tokens generated, billed at the model's audio output rate (falls back to the output rate). - The ID of the AI credit system feature. If omitted, automatically detects the organization's AI credit system feature. + The ID of the AI credit system feature. If omitted, automatically detects the customer's AI credit system feature. Required when the customer has more than one. @@ -87,7 +122,7 @@ await autumn.balances.trackTokens({ - Additional properties to attach to this usage event. The `model`, `input_tokens`, and `output_tokens` values are automatically included. + Additional properties to attach to this usage event. The token counts and a pricing breakdown (`cost`, `base_cost`, `markup`, `markup_source`, `tier_applied`, `rates`) are automatically included. ### Response diff --git a/apps/docs/mintlify/documentation/customers/tracking-usage.mdx b/apps/docs/mintlify/documentation/customers/tracking-usage.mdx index 0df65ea15..4c5bff968 100644 --- a/apps/docs/mintlify/documentation/customers/tracking-usage.mdx +++ b/apps/docs/mintlify/documentation/customers/tracking-usage.mdx @@ -142,6 +142,14 @@ The `modelId` must be in `provider/model` format, matching the provider and mode For providers with nested model paths (like OpenRouter), include the full path after the provider: `openrouter/anthropic/claude-opus-4.6`. +Token counts are **exclusive pools**: `inputTokens` should exclude cached tokens (pass those as `cacheReadTokens` / `cacheWriteTokens`) and `outputTokens` should exclude reasoning tokens (pass those as `reasoningTokens`). Audio tokens go in `audioInputTokens` / `audioOutputTokens`. See the [API reference](/api-reference/balances/trackTokens) for the full parameter list. + + + `autumn.balances.trackTokens` requires an autumn-js release that includes + the method. On older versions, call the REST endpoint directly — see the + cURL tab below. + + ```typescript TypeScript @@ -185,7 +193,10 @@ curl -X POST "https://api.useautumn.com/v1/balances.track_tokens" \ - If your organization has only one AI credit system feature, you can omit the `featureId` parameter — it will be auto-detected. + If the customer has exactly one AI credit system feature, you can omit the + `featureId` parameter — it will be auto-detected. The request fails with an + error if the customer has no AI credit system, or has more than one and no + `featureId` is provided. ### Vercel AI SDK integration diff --git a/apps/docs/mintlify/documentation/external-providers/ai-sdk.mdx b/apps/docs/mintlify/documentation/external-providers/ai-sdk.mdx index e9bb2b92a..74a99dac2 100644 --- a/apps/docs/mintlify/documentation/external-providers/ai-sdk.mdx +++ b/apps/docs/mintlify/documentation/external-providers/ai-sdk.mdx @@ -73,6 +73,10 @@ for await (const chunk of result.textStream) { } ``` +## Token pools + +The wrapper normalizes the AI SDK's usage object into the exclusive token pools that [trackTokens](/api-reference/balances/trackTokens) expects: text input (excluding cached tokens), text output (excluding reasoning tokens), cache reads, cache writes, and reasoning tokens. Each pool is billed at the model's published rate, so cached and reasoning-heavy requests are priced correctly without any extra work. + ## Model ID format The wrapped model constructs the `modelId` sent to Autumn using `provider/model` format, derived from the AI SDK model's `provider` and `modelId` fields. This must match a valid provider and model key from [Models.dev](https://models.dev). diff --git a/apps/docs/mintlify/documentation/modelling-pricing/credit-systems.mdx b/apps/docs/mintlify/documentation/modelling-pricing/credit-systems.mdx index f0fa1d9f9..661ccb98e 100644 --- a/apps/docs/mintlify/documentation/modelling-pricing/credit-systems.mdx +++ b/apps/docs/mintlify/documentation/modelling-pricing/credit-systems.mdx @@ -279,7 +279,20 @@ For AI applications that need to track token usage with per-model pricing, you c -Define an AI credit system with `modelMarkups` that maps model IDs to pricing configuration: +Markups are optional. `defaultMarkup` applies to every model unless overridden — by `providerMarkups` (keyed by the first segment of the model ID, e.g. `openrouter`), or by `modelMarkups` for a specific model, which takes highest priority. With no markups set, models are billed at their Models.dev base cost. + +A markup of `-100` makes the model free: usage events are still recorded, but nothing is deducted from the balance. + +```ts Simplest setup — one markup for everything +export const aiCredits = feature({ + id: 'ai_credits', + name: 'AI Credits', + type: 'ai_credit_system', + defaultMarkup: 30, // every model billed at models.dev cost + 30% +}); +``` + +Or mix the levels for finer control: ```ts autumn.config.ts import { feature, item, plan } from 'atmn'; @@ -288,10 +301,17 @@ export const aiCredits = feature({ id: 'ai_credits', name: 'AI Credits', type: 'ai_credit_system', + // Global fallback markup + defaultMarkup: 30, + // Per-provider defaults + providerMarkups: { + openrouter: { markup: 25 }, + }, + // Per-model overrides (highest priority) modelMarkups: { 'anthropic/claude-opus-4-5': { markup: 20 }, 'anthropic/claude-sonnet-4-5': { markup: 15 }, - 'openai/gpt-4o': { markup: 15 }, + 'openai/gpt-4o-mini': { markup: -100 }, // free for customers // For custom/self-hosted models, specify input/output costs in $/M tokens 'custom/my-model': { markup: 25, inputCost: 0.01, outputCost: 0.03 }, }, @@ -319,9 +339,10 @@ Push changes with `atmn push`. 1. Navigate to the features page, under Plans. 2. Click "Create Credit System" 3. Toggle "AI Credit System" to enable model-based pricing -4. Add the models you want to support with their markup percentages -5. For custom models, also specify input/output costs per million tokens -6. Click "Create" +4. Set a default markup %, and optionally add providers with their own default markups +5. Add the models you want to support, overriding the markup per model where needed +6. For custom models, also specify input/output costs per million tokens +7. Click "Create" @@ -333,7 +354,9 @@ Model IDs follow the `provider/model` format: - OpenRouter models: `openrouter/anthropic/claude-opus-4.6` - Custom models: `custom/my-model-name` -For standard models, pricing is automatically fetched from models.dev. For custom models, you must specify `inputCost` and `outputCost` in dollars per million tokens. +For standard models, pricing is automatically fetched from models.dev, including separate rates for cache reads/writes, reasoning, and audio tokens where the model publishes them, plus large-context tier pricing (e.g. above 200k input tokens) when applicable. + +For custom models, you must specify both `inputCost` and `outputCost` in dollars per million tokens — tracking fails if either is missing. Custom models bill input and output tokens only; cache, reasoning, and audio pools are ignored. ### Tracking Token Usage From b0beb3ed59615373ea74e128bf2a619bc2eea995 Mon Sep 17 00:00:00 2001 From: Charlie Lamb Date: Wed, 10 Jun 2026 17:10:46 +0100 Subject: [PATCH 43/46] chore: improve credit system type tabs --- .../features/credit-systems/components/CreditSystemSchema.tsx | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vite/src/views/products/features/credit-systems/components/CreditSystemSchema.tsx b/vite/src/views/products/features/credit-systems/components/CreditSystemSchema.tsx index f67bbad95..14cffbab1 100644 --- a/vite/src/views/products/features/credit-systems/components/CreditSystemSchema.tsx +++ b/vite/src/views/products/features/credit-systems/components/CreditSystemSchema.tsx @@ -74,7 +74,7 @@ export function CreditSystemSchema({ value={mode} onValueChange={handleModeChange} options={modeOptions} - className="w-fit" + className="w-full" /> )} From a9dbb595100b8c0c5a9bce01c505b7949ffc6a64 Mon Sep 17 00:00:00 2001 From: Ridhwan Hussain <73362400+TheUntraceable@users.noreply.github.com.> Date: Wed, 10 Jun 2026 17:36:28 +0100 Subject: [PATCH 44/46] docs: use correct openrouter anthropic model ID --- apps/docs/mintlify/api-reference/balances/trackTokens.mdx | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/apps/docs/mintlify/api-reference/balances/trackTokens.mdx b/apps/docs/mintlify/api-reference/balances/trackTokens.mdx index 28c892213..41894bbdc 100644 --- a/apps/docs/mintlify/api-reference/balances/trackTokens.mdx +++ b/apps/docs/mintlify/api-reference/balances/trackTokens.mdx @@ -41,7 +41,7 @@ await autumn.balances.trackTokens({ ```typescript OpenRouter (nested path) await autumn.balances.trackTokens({ customerId: "cus_123", - modelId: "openrouter/anthropic/claude-opus-4-6", + modelId: "openrouter/anthropic/claude-opus-4.6", inputTokens: 2000, outputTokens: 1000 }); From 8169b82aaeda29910ff77f1c9611d8782821e46c Mon Sep 17 00:00:00 2001 From: Ridhwan Hussain <73362400+TheUntraceable@users.noreply.github.com.> Date: Wed, 10 Jun 2026 17:38:17 +0100 Subject: [PATCH 45/46] docs: use correct openrouter anthropic model ID --- apps/docs/mintlify/documentation/external-providers/ai-sdk.mdx | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/apps/docs/mintlify/documentation/external-providers/ai-sdk.mdx b/apps/docs/mintlify/documentation/external-providers/ai-sdk.mdx index 74a99dac2..2c44f2670 100644 --- a/apps/docs/mintlify/documentation/external-providers/ai-sdk.mdx +++ b/apps/docs/mintlify/documentation/external-providers/ai-sdk.mdx @@ -99,7 +99,7 @@ const model = withAutumn({ customerId: "user_123", providerId: "openrouter", // Override provider prefix }); -// Sends modelId as "openrouter/anthropic/claude-opus-4-6" +// Sends modelId as "openrouter/anthropic/claude-opus-4.6" ``` ## Options From 5ccd4f0fa89d3ddf751be7ca6f4293929d4a3c01 Mon Sep 17 00:00:00 2001 From: Ridhwan Hussain <73362400+TheUntraceable@users.noreply.github.com.> Date: Wed, 10 Jun 2026 17:39:54 +0100 Subject: [PATCH 46/46] docs: use correct openrouter anthropic model ID --- apps/docs/mintlify/documentation/external-providers/ai-sdk.mdx | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/apps/docs/mintlify/documentation/external-providers/ai-sdk.mdx b/apps/docs/mintlify/documentation/external-providers/ai-sdk.mdx index 2c44f2670..bf526f96e 100644 --- a/apps/docs/mintlify/documentation/external-providers/ai-sdk.mdx +++ b/apps/docs/mintlify/documentation/external-providers/ai-sdk.mdx @@ -95,7 +95,7 @@ const openrouter = createOpenRouter(); const model = withAutumn({ autumn, - model: openrouter("anthropic/claude-opus-4-6"), + model: openrouter("anthropic/claude-opus-4.6"), customerId: "user_123", providerId: "openrouter", // Override provider prefix });