From 7b8b3752e467c98567d575b48fd91b30348d2376 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Filip=20Kujund=C5=BEi=C4=87?= Date: Fri, 2 Oct 2026 19:05:41 +0200 Subject: [PATCH] fix(ai): price the chat credit gate at the factored cost The chat gate read `model.costs` directly to decide affordability, cap `max_tokens` and size the credit hold, while usage is recorded at cost times the AI cost factor, so all three ran low by the factor on every request. The gate now resolves the factor and applies it to both rates. The factor is looked up by the key the provider records usage under. Providers expose that key through `meteringModelKey` and use it when recording, so the gate and the charge can't drift apart; the unreported-stream backstop records under the same key. Co-Authored-By: Claude Opus 5.5 --- .../ChatCompletionDriver.metering.test.ts | 113 ++++++++++++++++++ .../drivers/ai-chat/ChatCompletionDriver.ts | 27 ++++- .../providers/alibaba/AlibabaProvider.ts | 7 +- .../providers/azure/AzureChatProvider.ts | 7 +- .../providers/azure/AzureResponsesProvider.ts | 7 +- .../providers/byteplus/BytePlusProvider.ts | 7 +- .../providers/claude/ClaudeProvider.ts | 9 +- .../providers/deepseek/DeepSeekProvider.ts | 7 +- .../providers/gemini/GeminiChatProvider.ts | 7 +- .../ai-chat/providers/groq/GroqAIProvider.ts | 7 +- .../providers/hoonify/HoonifyProvider.ts | 7 +- .../providers/infron/InfronProvider.ts | 9 +- .../ai-chat/providers/meta/MetaProvider.ts | 7 +- .../providers/minimax/MiniMaxProvider.ts | 7 +- .../providers/mistral/MistralAiProvider.ts | 7 +- .../providers/moonshot/MoonshotProvider.ts | 7 +- .../neuralwatt/NeuralwattProvider.ts | 9 +- .../providers/ollama/OllamaProvider.ts | 5 + .../openai/OpenAiChatCompletionsProvider.ts | 7 +- .../openai/OpenAiChatResponsesProvider.ts | 7 +- .../openrouter/OpenRouterProvider.ts | 9 +- .../providers/together/TogetherAIProvider.ts | 7 +- .../ai-chat/providers/xai/XAIProvider.ts | 7 +- .../ai-chat/providers/zai/ZAIProvider.ts | 7 +- src/backend/drivers/ai-chat/types.ts | 6 + 25 files changed, 277 insertions(+), 29 deletions(-) diff --git a/src/backend/drivers/ai-chat/ChatCompletionDriver.metering.test.ts b/src/backend/drivers/ai-chat/ChatCompletionDriver.metering.test.ts index 9609f08ea..b579b8576 100644 --- a/src/backend/drivers/ai-chat/ChatCompletionDriver.metering.test.ts +++ b/src/backend/drivers/ai-chat/ChatCompletionDriver.metering.test.ts @@ -434,6 +434,119 @@ describe('ChatCompletionDriver credit gate against real metering', () => { }); }); +describe('ChatCompletionDriver credit gate with an AI cost factor', () => { + const FACTOR = 1.3; + let seenKeys: string[]; + const listener = (key: string, event: { factor: number }) => { + seenKeys.push(key); + event.factor = FACTOR; + }; + + beforeEach(() => { + seenKeys = []; + server.clients.event.on('ai.cost.factor.*', listener as never); + }); + + afterEach(() => { + server.clients.event.off('ai.cost.factor.*', listener as never); + }); + + const okResult = { + message: { role: 'assistant', content: [{ type: 'text', text: 'ok' }] }, + usage: { input_tokens: 1, output_tokens: 1 }, + finish_reason: 'stop', + } as never; + + // Usage is recorded at cost × factor, so a cap sized at the bare cost + // lets a completion spend past what's left. + it('caps output at the factored price', async () => { + const actor = freeUser() as never; + const metering = server.services.metering; + const allowance = (await metering.getActorSubscription(actor)) + .monthUsageAllowance; + await metering.incrementUsage( + actor, + 'test:prior-spend', + 1, + Math.floor(allowance * 0.9), + ); + const remaining = await metering.getRemainingUsage(actor); + const completeSpy = vi + .spyOn(FakeChatProvider.prototype, 'complete') + .mockResolvedValueOnce(okResult); + + await withTestActor( + () => + driver.complete({ + model: 'priced', + messages: [{ role: 'user', content: 'go' }], + max_tokens: 100_000, + }), + actor, + ); + + const passed = completeSpy.mock.calls[0]![0] as { max_tokens: number }; + expect(passed.max_tokens * 2000 * FACTOR).toBeLessThanOrEqual( + remaining, + ); + expect(passed.max_tokens).toBeGreaterThan( + remaining / (2000 * FACTOR) - 10, + ); + }); + + it('reserves the factored worst case', async () => { + const actor = freeUser() as never; + const reserve = vi.spyOn(server.services.metering, 'reserveCredits'); + const completeSpy = vi + .spyOn(FakeChatProvider.prototype, 'complete') + .mockResolvedValueOnce(okResult); + + await withTestActor( + () => + driver.complete({ + model: 'priced', + messages: [{ role: 'user', content: 'go' }], + max_tokens: 10, + }), + actor, + ); + + const passed = completeSpy.mock.calls[0]![0] as { max_tokens: number }; + expect(passed.max_tokens).toBe(10); + const [, held] = reserve.mock.calls[0]!; + // At least the 10 capped output tokens at 2000 × factor. + expect(held).toBeGreaterThanOrEqual(10 * 2000 * FACTOR); + }); + + // The factor is looked up by the key the provider records under, or a + // per-model factor would price the gate and the charge differently. + it('asks for the factor under the metering key the provider records with', async () => { + const proto = FakeChatProvider.prototype as unknown as { + meteringModelKey?: (id: string) => string; + }; + proto.meteringModelKey = (id) => `fake-recorded:${id}`; + try { + vi.spyOn( + FakeChatProvider.prototype, + 'complete', + ).mockResolvedValueOnce(okResult); + await withTestActor( + () => + driver.complete({ + model: 'priced', + messages: [{ role: 'user', content: 'go' }], + }), + freeUser() as never, + ); + expect(seenKeys).toContain( + 'ai.cost.factor.ai-chat.fake-recorded:priced', + ); + } finally { + delete proto.meteringModelKey; + } + }); +}); + describe('ChatCompletionDriver credit gate on multimodal prompts', () => { it('prices attachments into the pre-flight affordability check', async () => { const actor = freeUser() as never; diff --git a/src/backend/drivers/ai-chat/ChatCompletionDriver.ts b/src/backend/drivers/ai-chat/ChatCompletionDriver.ts index 65f939b2b..005f8c4f6 100644 --- a/src/backend/drivers/ai-chat/ChatCompletionDriver.ts +++ b/src/backend/drivers/ai-chat/ChatCompletionDriver.ts @@ -948,11 +948,21 @@ export class ChatCompletionDriver extends PuterDriver { // A prompt estimated past a long-context threshold pays the raised // rates on input and output alike. const multipliers = longContextMultipliers(model, promptTokenEstimate); + // Recorded costs pass through the AI cost factor; the gate prices at + // the same rate or it under-reserves on every request. + const costFactor = await this.#aiMetering.costFactor( + actor, + this.#meteringModelKey(model), + ); // `|| 0` also catches NaN from a malformed cost table. const inputTokenCost = - (Number(model.costs?.[inputKey] ?? 0) || 0) * multipliers.input; + (Number(model.costs?.[inputKey] ?? 0) || 0) * + multipliers.input * + costFactor; const outputTokenCost = - (Number(model.costs?.[outputKey] ?? 0) || 0) * multipliers.output; + (Number(model.costs?.[outputKey] ?? 0) || 0) * + multipliers.output * + costFactor; const approximateInputCost = promptTokenEstimate * inputTokenCost; const minimumCredits = Number(model.minimumCredits || 1); @@ -1021,7 +1031,8 @@ export class ChatCompletionDriver extends PuterDriver { // What this attempt can cost at worst: the prompt, plus output run to // the cap just set. Capped output is what makes the number finite — // for a model with no output price the output term is zero and the - // prompt estimate stands alone. + // prompt estimate stands alone. Already factored, so it's held through + // the plain service rather than scaled a second time. const worstCaseCost = approximateInputCost + (args.max_tokens ?? 0) * outputTokenCost; return this.services.metering.reserveCredits( @@ -1030,6 +1041,14 @@ export class ChatCompletionDriver extends PuterDriver { ); } + /** The key `model`'s usage is recorded and cost-factored under. */ + #meteringModelKey(model: IChatModel): string { + return ( + this.#providers[model.provider!]?.meteringModelKey?.(model.id) ?? + `${model.provider}:${model.id}` + ); + } + /** * Charge a stream that produced output but never reported usage. * @@ -1085,7 +1104,7 @@ export class ChatCompletionDriver extends PuterDriver { [`estimated_${outputKey}`]: outputTokens, }, actor, - `${model.provider}:${model.id}`, + this.#meteringModelKey(model), { // Undefined when the model has no cost table: the entry is // recorded unpriced rather than free. diff --git a/src/backend/drivers/ai-chat/providers/alibaba/AlibabaProvider.ts b/src/backend/drivers/ai-chat/providers/alibaba/AlibabaProvider.ts index 0ab886afb..61b0999f8 100644 --- a/src/backend/drivers/ai-chat/providers/alibaba/AlibabaProvider.ts +++ b/src/backend/drivers/ai-chat/providers/alibaba/AlibabaProvider.ts @@ -58,6 +58,11 @@ export class AlibabaProvider implements IChatProvider { return modelLookupNames(this.models()); } + /** The model key this provider records usage under. */ + meteringModelKey(modelId: string): string { + return `alibaba:${modelId}`; + } + async complete({ messages, stream, @@ -100,7 +105,7 @@ export class AlibabaProvider implements IChatProvider { this.#meteringService.utilRecordUsageObject( trackedUsage, actor!, - `alibaba:${modelUsed.id}`, + this.meteringModelKey(modelUsed.id), costsOverride, ); return trackedUsage; diff --git a/src/backend/drivers/ai-chat/providers/azure/AzureChatProvider.ts b/src/backend/drivers/ai-chat/providers/azure/AzureChatProvider.ts index bbc92c4af..9e55e0345 100644 --- a/src/backend/drivers/ai-chat/providers/azure/AzureChatProvider.ts +++ b/src/backend/drivers/ai-chat/providers/azure/AzureChatProvider.ts @@ -112,6 +112,11 @@ export class AzureChatProvider implements IChatProvider { return this.#defaultModel; } + /** The model key this provider records usage under. */ + meteringModelKey(modelId: string | undefined): string { + return `azure-openai:${modelId}`; + } + async complete( params: ICompleteArguments, ): ReturnType { @@ -275,7 +280,7 @@ export class AzureChatProvider implements IChatProvider { this.#meteringService.utilRecordUsageObject( trackedUsage, actor!, - `azure-openai:${modelUsed?.id}`, + this.meteringModelKey(modelUsed?.id), costsOverrideFromModel, ); return trackedUsage; diff --git a/src/backend/drivers/ai-chat/providers/azure/AzureResponsesProvider.ts b/src/backend/drivers/ai-chat/providers/azure/AzureResponsesProvider.ts index 2d24a4c92..ee97209fd 100644 --- a/src/backend/drivers/ai-chat/providers/azure/AzureResponsesProvider.ts +++ b/src/backend/drivers/ai-chat/providers/azure/AzureResponsesProvider.ts @@ -94,6 +94,11 @@ export class AzureResponsesProvider implements IChatProvider { return this.#defaultModel; } + /** The model key this provider records usage under. */ + meteringModelKey(modelId: string | undefined): string { + return `azure-openai:${modelId}`; + } + async complete({ messages, model, @@ -262,7 +267,7 @@ export class AzureResponsesProvider implements IChatProvider { this.#meteringService.utilRecordUsageObject( trackedUsage, actor, - `azure-openai:${modelUsed?.id}`, + this.meteringModelKey(modelUsed?.id), costsOverrideFromModel, ); return trackedUsage; diff --git a/src/backend/drivers/ai-chat/providers/byteplus/BytePlusProvider.ts b/src/backend/drivers/ai-chat/providers/byteplus/BytePlusProvider.ts index b6f2a9921..31bd1f2d1 100644 --- a/src/backend/drivers/ai-chat/providers/byteplus/BytePlusProvider.ts +++ b/src/backend/drivers/ai-chat/providers/byteplus/BytePlusProvider.ts @@ -80,6 +80,11 @@ export class BytePlusProvider implements IChatProvider { return modelLookupNames(this.models()); } + /** The model key this provider records usage under. */ + meteringModelKey(modelId: string): string { + return `byteplus:${modelId}`; + } + async complete( params: ICompleteArguments, ): ReturnType { @@ -151,7 +156,7 @@ export class BytePlusProvider implements IChatProvider { this.#meteringService.utilRecordUsageObject( trackedUsage, actor!, - `byteplus:${modelUsed.id}`, + this.meteringModelKey(modelUsed.id), costsOverride, ); return trackedUsage; diff --git a/src/backend/drivers/ai-chat/providers/claude/ClaudeProvider.ts b/src/backend/drivers/ai-chat/providers/claude/ClaudeProvider.ts index ff88f533c..8dd181278 100644 --- a/src/backend/drivers/ai-chat/providers/claude/ClaudeProvider.ts +++ b/src/backend/drivers/ai-chat/providers/claude/ClaudeProvider.ts @@ -135,6 +135,11 @@ export class ClaudeProvider implements IChatProvider { return modelLookupNames(this.models()); } + /** The model key this provider records usage under. */ + meteringModelKey(modelId: string): string { + return `claude:${modelId}`; + } + async complete({ messages, stream, @@ -669,7 +674,7 @@ export class ClaudeProvider implements IChatProvider { this.#meteringService.utilRecordUsageObject( usageSum, actor, - `claude:${modelUsed.id}`, + this.meteringModelKey(modelUsed.id), costsOverrideFromModel, ); chatStream.end(usageSum); @@ -696,7 +701,7 @@ export class ClaudeProvider implements IChatProvider { this.#meteringService.utilRecordUsageObject( usage, actor, - `claude:${modelUsed.id}`, + this.meteringModelKey(modelUsed.id), costsOverrideFromModel, ); diff --git a/src/backend/drivers/ai-chat/providers/deepseek/DeepSeekProvider.ts b/src/backend/drivers/ai-chat/providers/deepseek/DeepSeekProvider.ts index c750b2321..e188cca7d 100644 --- a/src/backend/drivers/ai-chat/providers/deepseek/DeepSeekProvider.ts +++ b/src/backend/drivers/ai-chat/providers/deepseek/DeepSeekProvider.ts @@ -52,6 +52,11 @@ export class DeepSeekProvider implements IChatProvider { return modelLookupNames(this.models()); } + /** The model key this provider records usage under. */ + meteringModelKey(modelId: string): string { + return `deepseek:${modelId}`; + } + async complete({ messages, stream, @@ -131,7 +136,7 @@ export class DeepSeekProvider implements IChatProvider { this.#meteringService.utilRecordUsageObject( trackedUsage, actor!, - `deepseek:${modelUsed.id}`, + this.meteringModelKey(modelUsed.id), costsOverrideFromModel, ); return trackedUsage; diff --git a/src/backend/drivers/ai-chat/providers/gemini/GeminiChatProvider.ts b/src/backend/drivers/ai-chat/providers/gemini/GeminiChatProvider.ts index a70a5c813..572890fe6 100644 --- a/src/backend/drivers/ai-chat/providers/gemini/GeminiChatProvider.ts +++ b/src/backend/drivers/ai-chat/providers/gemini/GeminiChatProvider.ts @@ -58,6 +58,11 @@ export class GeminiChatProvider implements IChatProvider { return modelLookupNames(await this.models()); } + /** The model key this provider records usage under. */ + meteringModelKey(modelId: string | undefined): string { + return `gemini:${modelId}`; + } + async complete({ messages, stream, @@ -156,7 +161,7 @@ export class GeminiChatProvider implements IChatProvider { this.meteringService.utilRecordUsageObject( trackedUsage, actor!, - `gemini:${modelUsed?.id}`, + this.meteringModelKey(modelUsed?.id), costsOverrideFromModel, ); diff --git a/src/backend/drivers/ai-chat/providers/groq/GroqAIProvider.ts b/src/backend/drivers/ai-chat/providers/groq/GroqAIProvider.ts index 690ca2e82..8f1fd8e7b 100644 --- a/src/backend/drivers/ai-chat/providers/groq/GroqAIProvider.ts +++ b/src/backend/drivers/ai-chat/providers/groq/GroqAIProvider.ts @@ -51,6 +51,11 @@ export class GroqAIProvider implements IChatProvider { return modelLookupNames(this.models()); } + /** The model key this provider records usage under. */ + meteringModelKey(modelId: string): string { + return `groq:${modelId}`; + } + async complete({ messages, model, @@ -99,7 +104,7 @@ export class GroqAIProvider implements IChatProvider { this.#meteringService.utilRecordUsageObject( trackedUsage, actor, - `groq:${modelUsed.id}`, + this.meteringModelKey(modelUsed.id), costsOverride, ); return trackedUsage; diff --git a/src/backend/drivers/ai-chat/providers/hoonify/HoonifyProvider.ts b/src/backend/drivers/ai-chat/providers/hoonify/HoonifyProvider.ts index e07d799b6..9325b76ff 100644 --- a/src/backend/drivers/ai-chat/providers/hoonify/HoonifyProvider.ts +++ b/src/backend/drivers/ai-chat/providers/hoonify/HoonifyProvider.ts @@ -77,6 +77,11 @@ export class HoonifyProvider implements IChatProvider { return modelIds; } + /** The model key this provider records usage under. */ + meteringModelKey(modelId: string): string { + return modelId; + } + async complete( params: ICompleteArguments, ): ReturnType { @@ -149,7 +154,7 @@ export class HoonifyProvider implements IChatProvider { this.#meteringService.utilRecordUsageObject( trackedUsage, actor, - modelUsed.id, + this.meteringModelKey(modelUsed.id), costsOverrideFromModel, ); return trackedUsage; diff --git a/src/backend/drivers/ai-chat/providers/infron/InfronProvider.ts b/src/backend/drivers/ai-chat/providers/infron/InfronProvider.ts index 90f43123f..ab6360074 100644 --- a/src/backend/drivers/ai-chat/providers/infron/InfronProvider.ts +++ b/src/backend/drivers/ai-chat/providers/infron/InfronProvider.ts @@ -223,6 +223,11 @@ export class InfronProvider implements IChatProvider { } /** AI Chat completion method. See AIChatService for more details. */ + /** The model key this provider records usage under. */ + meteringModelKey(modelId: string): string { + return modelId; + } + async complete({ messages, stream, @@ -311,7 +316,7 @@ export class InfronProvider implements IChatProvider { this.#meteringService.utilRecordUsageObject( billedTrackedUsage, actor, - modelUsed.id, + this.meteringModelKey(modelUsed.id), costOverwrites, ); (billedTrackedUsage as Record).usd_cents = @@ -331,7 +336,7 @@ export class InfronProvider implements IChatProvider { this.#meteringService.utilRecordUsageObject( trackedUsage, actor, - modelUsed.id, + this.meteringModelKey(modelUsed.id), costOverwrites, ); return trackedUsage; diff --git a/src/backend/drivers/ai-chat/providers/meta/MetaProvider.ts b/src/backend/drivers/ai-chat/providers/meta/MetaProvider.ts index c5d0e36ae..b4a21825e 100644 --- a/src/backend/drivers/ai-chat/providers/meta/MetaProvider.ts +++ b/src/backend/drivers/ai-chat/providers/meta/MetaProvider.ts @@ -99,6 +99,11 @@ export class MetaProvider implements IChatProvider { return modelLookupNames(this.models()); } + /** The model key this provider records usage under. */ + meteringModelKey(modelId: string): string { + return `meta:${modelId}`; + } + async complete( params: ICompleteArguments, ): ReturnType { @@ -227,7 +232,7 @@ export class MetaProvider implements IChatProvider { this.#meteringService.utilRecordUsageObject( trackedUsage, actor!, - `meta:${modelUsed.id}`, + this.meteringModelKey(modelUsed.id), costsOverride, ); return trackedUsage; diff --git a/src/backend/drivers/ai-chat/providers/minimax/MiniMaxProvider.ts b/src/backend/drivers/ai-chat/providers/minimax/MiniMaxProvider.ts index bdf0871c8..53939f43d 100644 --- a/src/backend/drivers/ai-chat/providers/minimax/MiniMaxProvider.ts +++ b/src/backend/drivers/ai-chat/providers/minimax/MiniMaxProvider.ts @@ -58,6 +58,11 @@ export class MiniMaxProvider implements IChatProvider { return modelLookupNames(this.models()); } + /** The model key this provider records usage under. */ + meteringModelKey(modelId: string): string { + return `minimax:${modelId}`; + } + async complete({ messages, stream, @@ -111,7 +116,7 @@ export class MiniMaxProvider implements IChatProvider { this.#meteringService.utilRecordUsageObject( trackedUsage, actor!, - `minimax:${modelUsed.id}`, + this.meteringModelKey(modelUsed.id), costsOverride, ); return trackedUsage; diff --git a/src/backend/drivers/ai-chat/providers/mistral/MistralAiProvider.ts b/src/backend/drivers/ai-chat/providers/mistral/MistralAiProvider.ts index 83d6256cb..c9addfce8 100644 --- a/src/backend/drivers/ai-chat/providers/mistral/MistralAiProvider.ts +++ b/src/backend/drivers/ai-chat/providers/mistral/MistralAiProvider.ts @@ -165,6 +165,11 @@ export class MistralAIProvider implements IChatProvider { }); } + /** The model key this provider records usage under. */ + meteringModelKey(modelId: string): string { + return `mistral:${modelId}`; + } + async complete({ messages, stream, @@ -396,7 +401,7 @@ export class MistralAIProvider implements IChatProvider { this.#meteringService.utilRecordUsageObject( trackedUsage, actor, - `mistral:${selectedModel.id}`, + this.meteringModelKey(selectedModel.id), costsOverrideFromModel, ); return trackedUsage; diff --git a/src/backend/drivers/ai-chat/providers/moonshot/MoonshotProvider.ts b/src/backend/drivers/ai-chat/providers/moonshot/MoonshotProvider.ts index ebb26de62..1d656bdb7 100644 --- a/src/backend/drivers/ai-chat/providers/moonshot/MoonshotProvider.ts +++ b/src/backend/drivers/ai-chat/providers/moonshot/MoonshotProvider.ts @@ -56,6 +56,11 @@ export class MoonshotProvider implements IChatProvider { return modelLookupNames(this.models()); } + /** The model key this provider records usage under. */ + meteringModelKey(modelId: string): string { + return `moonshotai:${modelId}`; + } + async complete({ messages, stream, @@ -107,7 +112,7 @@ export class MoonshotProvider implements IChatProvider { this.#meteringService.utilRecordUsageObject( trackedUsage, actor, - `moonshotai:${modelUsed.id}`, + this.meteringModelKey(modelUsed.id), costsOverride, ); return trackedUsage; diff --git a/src/backend/drivers/ai-chat/providers/neuralwatt/NeuralwattProvider.ts b/src/backend/drivers/ai-chat/providers/neuralwatt/NeuralwattProvider.ts index f72dedde3..a548e994b 100644 --- a/src/backend/drivers/ai-chat/providers/neuralwatt/NeuralwattProvider.ts +++ b/src/backend/drivers/ai-chat/providers/neuralwatt/NeuralwattProvider.ts @@ -157,6 +157,11 @@ export class NeuralwattProvider implements IChatProvider { return undefined; } + /** The model key this provider records usage under. */ + meteringModelKey(modelId: string): string { + return modelId; + } + async complete({ messages, stream, @@ -305,7 +310,7 @@ export class NeuralwattProvider implements IChatProvider { this.#meteringService.utilRecordUsageObject( billedTrackedUsage, actor!, - modelUsed.id, + this.meteringModelKey(modelUsed.id), costOverwrites, ); const result = annotate(billedTrackedUsage); @@ -327,7 +332,7 @@ export class NeuralwattProvider implements IChatProvider { this.#meteringService.utilRecordUsageObject( trackedUsage, actor!, - modelUsed.id, + this.meteringModelKey(modelUsed.id), costOverwrites, ); return annotate(trackedUsage); diff --git a/src/backend/drivers/ai-chat/providers/ollama/OllamaProvider.ts b/src/backend/drivers/ai-chat/providers/ollama/OllamaProvider.ts index 4a2affda0..3e6a060b5 100644 --- a/src/backend/drivers/ai-chat/providers/ollama/OllamaProvider.ts +++ b/src/backend/drivers/ai-chat/providers/ollama/OllamaProvider.ts @@ -108,6 +108,11 @@ export class OllamaChatProvider implements IChatProvider { } return model_names; } + /** The model key this provider records usage under. */ + meteringModelKey(modelId: string): string { + return modelId; + } + async complete({ messages, stream, diff --git a/src/backend/drivers/ai-chat/providers/openai/OpenAiChatCompletionsProvider.ts b/src/backend/drivers/ai-chat/providers/openai/OpenAiChatCompletionsProvider.ts index 9237df7db..de9da5349 100644 --- a/src/backend/drivers/ai-chat/providers/openai/OpenAiChatCompletionsProvider.ts +++ b/src/backend/drivers/ai-chat/providers/openai/OpenAiChatCompletionsProvider.ts @@ -96,6 +96,11 @@ export class OpenAiChatProvider implements IChatProvider { return this.#defaultModel; } + /** The model key this provider records usage under. */ + meteringModelKey(modelId: string | undefined): string { + return `openai:${modelId}`; + } + async complete( params: ICompleteArguments, ): ReturnType { @@ -247,7 +252,7 @@ export class OpenAiChatProvider implements IChatProvider { this.#meteringService.utilRecordUsageObject( trackedUsage, actor, - `openai:${modelUsed?.id}`, + this.meteringModelKey(modelUsed?.id), costsOverrideFromModel, ); return trackedUsage; diff --git a/src/backend/drivers/ai-chat/providers/openai/OpenAiChatResponsesProvider.ts b/src/backend/drivers/ai-chat/providers/openai/OpenAiChatResponsesProvider.ts index 7a4f444c7..57396f41c 100644 --- a/src/backend/drivers/ai-chat/providers/openai/OpenAiChatResponsesProvider.ts +++ b/src/backend/drivers/ai-chat/providers/openai/OpenAiChatResponsesProvider.ts @@ -88,6 +88,11 @@ export class OpenAiResponsesChatProvider implements IChatProvider { return this.#defaultModel; } + /** The model key this provider records usage under. */ + meteringModelKey(modelId: string | undefined): string { + return `openai:${modelId}`; + } + async complete({ messages, model, @@ -280,7 +285,7 @@ export class OpenAiResponsesChatProvider implements IChatProvider { this.#meteringService.utilRecordUsageObject( trackedUsage, actor!, - `openai:${modelUsed?.id}`, + this.meteringModelKey(modelUsed?.id), costsOverrideFromModel, ); return trackedUsage; diff --git a/src/backend/drivers/ai-chat/providers/openrouter/OpenRouterProvider.ts b/src/backend/drivers/ai-chat/providers/openrouter/OpenRouterProvider.ts index be20cb48c..ce2299ce8 100644 --- a/src/backend/drivers/ai-chat/providers/openrouter/OpenRouterProvider.ts +++ b/src/backend/drivers/ai-chat/providers/openrouter/OpenRouterProvider.ts @@ -93,6 +93,11 @@ export class OpenRouterProvider implements IChatProvider { } /** AI Chat completion method. See AIChatService for more details. */ + /** The model key this provider records usage under. */ + meteringModelKey(modelId: string): string { + return modelId; + } + async complete({ messages, stream, @@ -194,7 +199,7 @@ export class OpenRouterProvider implements IChatProvider { this.#meteringService.utilRecordUsageObject( trackedUsage, actor, - modelUsed.id, + this.meteringModelKey(modelUsed.id), costOverwrites, ); (trackedUsage as Record).usd_cents = @@ -221,7 +226,7 @@ export class OpenRouterProvider implements IChatProvider { this.#meteringService.utilRecordUsageObject( trackedUsage, actor, - modelUsed.id, + this.meteringModelKey(modelUsed.id), costOverwrites, ); return trackedUsage; diff --git a/src/backend/drivers/ai-chat/providers/together/TogetherAIProvider.ts b/src/backend/drivers/ai-chat/providers/together/TogetherAIProvider.ts index 5278b4376..617edcf2f 100644 --- a/src/backend/drivers/ai-chat/providers/together/TogetherAIProvider.ts +++ b/src/backend/drivers/ai-chat/providers/together/TogetherAIProvider.ts @@ -110,6 +110,11 @@ export class TogetherAIProvider implements IChatProvider { return modelLookupNames(await this.models()); } + /** The model key this provider records usage under. */ + meteringModelKey(modelId: string): string { + return `togetherai:${modelId.replace(/^togetherai:/, '')}`; + } + async complete({ messages, stream, @@ -174,7 +179,7 @@ export class TogetherAIProvider implements IChatProvider { this.#meteringService.utilRecordUsageObject( trackedUsage, actor, - `togetherai:${modelIdForParams}`, + this.meteringModelKey(modelUsed.id), costsOverride, ); return trackedUsage; diff --git a/src/backend/drivers/ai-chat/providers/xai/XAIProvider.ts b/src/backend/drivers/ai-chat/providers/xai/XAIProvider.ts index e09b2ffe1..2a39b6d5c 100644 --- a/src/backend/drivers/ai-chat/providers/xai/XAIProvider.ts +++ b/src/backend/drivers/ai-chat/providers/xai/XAIProvider.ts @@ -56,6 +56,11 @@ export class XAIProvider implements IChatProvider { return modelLookupNames(this.models()); } + /** The model key this provider records usage under. */ + meteringModelKey(modelId: string): string { + return `xai:${modelId}`; + } + async complete({ messages, stream, @@ -105,7 +110,7 @@ export class XAIProvider implements IChatProvider { this.#meteringService.utilRecordUsageObject( trackedUsage, actor, - `xai:${modelUsed.id}`, + this.meteringModelKey(modelUsed.id), costsOverride, ); return metered; diff --git a/src/backend/drivers/ai-chat/providers/zai/ZAIProvider.ts b/src/backend/drivers/ai-chat/providers/zai/ZAIProvider.ts index c1eac5006..eea785d06 100644 --- a/src/backend/drivers/ai-chat/providers/zai/ZAIProvider.ts +++ b/src/backend/drivers/ai-chat/providers/zai/ZAIProvider.ts @@ -79,6 +79,11 @@ export class ZAIProvider implements IChatProvider { return modelLookupNames(this.models()); } + /** The model key this provider records usage under. */ + meteringModelKey(modelId: string): string { + return `zai:${modelId}`; + } + async complete( params: ICompleteArguments, ): ReturnType { @@ -162,7 +167,7 @@ export class ZAIProvider implements IChatProvider { this.#meteringService.utilRecordUsageObject( trackedUsage, actor!, - `zai:${modelUsed.id}`, + this.meteringModelKey(modelUsed.id), costsOverrideFromModel, ); return trackedUsage; diff --git a/src/backend/drivers/ai-chat/types.ts b/src/backend/drivers/ai-chat/types.ts index 0d8bf5c3f..937abc256 100644 --- a/src/backend/drivers/ai-chat/types.ts +++ b/src/backend/drivers/ai-chat/types.ts @@ -173,4 +173,10 @@ export interface IChatProvider { * attempt. */ readonly resolvesPuterPaths?: boolean; + /** + * The model key usage is recorded under, which is also what the AI cost + * factor is looked up by. The driver assumes `:` when + * absent. + */ + meteringModelKey?(modelId: string): string; }