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 <noreply@anthropic.com>
This commit is contained in:
Filip KujundžićandClaude Opus 5.5 authored and 404oops committed 2026-10-02 21:21:18 +02:00
1 parent 521812be02
commit 7b8b3752e4
25 files changed
+277 -29

No files matched your search

@@ -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;
@@ -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.
@@ -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;
@@ -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<IChatProvider['complete']> {
@@ -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;
@@ -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;
@@ -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<IChatProvider['complete']> {
@@ -151,7 +156,7 @@ export class BytePlusProvider implements IChatProvider {
this.#meteringService.utilRecordUsageObject(
trackedUsage,
actor!,
`byteplus:${modelUsed.id}`,
this.meteringModelKey(modelUsed.id),
costsOverride,
);
return trackedUsage;
@@ -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,
);
@@ -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;
@@ -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,
);
@@ -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;
@@ -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<IChatProvider['complete']> {
@@ -149,7 +154,7 @@ export class HoonifyProvider implements IChatProvider {
this.#meteringService.utilRecordUsageObject(
trackedUsage,
actor,
modelUsed.id,
this.meteringModelKey(modelUsed.id),
costsOverrideFromModel,
);
return trackedUsage;
@@ -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<string, number>).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;
@@ -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<IChatProvider['complete']> {
@@ -227,7 +232,7 @@ export class MetaProvider implements IChatProvider {
this.#meteringService.utilRecordUsageObject(
trackedUsage,
actor!,
`meta:${modelUsed.id}`,
this.meteringModelKey(modelUsed.id),
costsOverride,
);
return trackedUsage;
@@ -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;
@@ -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;
@@ -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;
@@ -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);
@@ -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,
@@ -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<IChatProvider['complete']> {
@@ -247,7 +252,7 @@ export class OpenAiChatProvider implements IChatProvider {
this.#meteringService.utilRecordUsageObject(
trackedUsage,
actor,
`openai:${modelUsed?.id}`,
this.meteringModelKey(modelUsed?.id),
costsOverrideFromModel,
);
return trackedUsage;
@@ -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;
@@ -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<string, number>).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;
@@ -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;
@@ -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;
@@ -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<IChatProvider['complete']> {
@@ -162,7 +167,7 @@ export class ZAIProvider implements IChatProvider {
this.#meteringService.utilRecordUsageObject(
trackedUsage,
actor!,
`zai:${modelUsed.id}`,
this.meteringModelKey(modelUsed.id),
costsOverrideFromModel,
);
return trackedUsage;
+6
View File
@@ -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 `<provider>:<model id>` when
* absent.
*/
meteringModelKey?(modelId: string): string;
}