From fb7968a1c732d0b7c9cf4ac3154ef360811a429a Mon Sep 17 00:00:00 2001 From: Daniel Salazar Date: Sat, 15 Aug 2026 10:52:29 -0700 Subject: [PATCH] fix: metering hardening; handle burst of unfinished ai requests (#3585) --- doc/alarms.md | 10 +- src/backend/core/http/HttpError.ts | 52 +- .../ChatCompletionDriver.edges.test.ts | 39 +- .../ChatCompletionDriver.metering.test.ts | 508 ++++++++++++++++++ .../ChatCompletionDriver.routing.test.ts | 67 +++ .../ai-chat/ChatCompletionDriver.test.ts | 39 +- .../drivers/ai-chat/ChatCompletionDriver.ts | 430 +++++++++++---- .../providers/claude/ClaudeProvider.ts | 6 +- .../drivers/ai-chat/utils/OpenAIUtil.js | 125 +++-- .../drivers/ai-chat/utils/Streaming.js | 117 +++- .../drivers/ai-chat/utils/pricing.test.ts | 28 +- src/backend/drivers/ai-chat/utils/pricing.ts | 32 +- .../ai-chat/utils/usageEstimate.test.ts | 187 +++++++ .../drivers/ai-chat/utils/usageEstimate.ts | 194 +++++++ .../providers/gemini/GeminiImageProvider.ts | 8 +- .../providers/openai/OpenAiImageProvider.ts | 14 +- src/backend/drivers/util/tokenEstimate.ts | 37 ++ src/backend/server.ts | 6 + .../services/metering/MeteringService.test.ts | 88 +++ .../services/metering/MeteringService.ts | 79 ++- src/backend/services/metering/types.ts | 24 + src/backend/stores/index.ts | 3 + .../stores/metering/CreditHoldStore.test.ts | 144 +++++ .../stores/metering/CreditHoldStore.ts | 233 ++++++++ src/docs/src/rate-limits-and-quotas.md | 4 + 25 files changed, 2227 insertions(+), 247 deletions(-) create mode 100644 src/backend/drivers/ai-chat/ChatCompletionDriver.metering.test.ts create mode 100644 src/backend/drivers/ai-chat/utils/usageEstimate.test.ts create mode 100644 src/backend/drivers/ai-chat/utils/usageEstimate.ts create mode 100644 src/backend/drivers/util/tokenEstimate.ts create mode 100644 src/backend/stores/metering/CreditHoldStore.test.ts create mode 100644 src/backend/stores/metering/CreditHoldStore.ts diff --git a/doc/alarms.md b/doc/alarms.md index fa01b9ec2..8a9df3540 100644 --- a/doc/alarms.md +++ b/doc/alarms.md @@ -44,7 +44,15 @@ pass one explicitly unless you really mean "page someone". event; the 429 is the whole signal, and alarming on it only produces noise proportional to traffic. An *upstream provider* rate-limiting us is the opposite case and still alarms (`upstream_rate_limited`, `info`) — that one - is not something we chose. + is not something we chose. The exception is a free model: nothing is billed, + nothing is actionable, and the throttling is the price of the model, so the + AI chat driver marks those `noAlarm` and only paid models still record. + +`noAlarm` is the code-level mute: an `HttpError` carrying it skips the +terminal gate entirely, no matter its status or code. It is for failures the +call site *already knows* are expected and traffic-proportional — reach for it +there, not as a way to quiet an alarm you haven't diagnosed +(`severityOverrides` below is the knob for that). An extension whose signals are all one tier can default its own local `raiseAlarm` helper to that tier instead of repeating it at every call site — diff --git a/src/backend/core/http/HttpError.ts b/src/backend/core/http/HttpError.ts index 3e26321c1..2a7046295 100644 --- a/src/backend/core/http/HttpError.ts +++ b/src/backend/core/http/HttpError.ts @@ -61,23 +61,22 @@ export type LegacyErrorCodes = * * This file is part of Puter. * - * Puter is free software: you can redistribute it and/or modify - * it under the terms of the GNU Affero General Public License as published - * by the Free Software Foundation, either version 3 of the License, or - * (at your option) any later version. + * Puter is free software: you can redistribute it and/or modify it under the + * terms of the GNU Affero General Public License as published by the Free + * Software Foundation, either version 3 of the License, or (at your option) any + * later version. * - * This program is distributed in the hope that it will be useful, - * but WITHOUT ANY WARRANTY; without even the implied warranty of - * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the - * GNU Affero General Public License for more details. + * This program is distributed in the hope that it will be useful, but WITHOUT + * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS + * FOR A PARTICULAR PURPOSE. See the GNU Affero General Public License for more + * details. * * You should have received a copy of the GNU Affero General Public License - * along with this program. If not, see . + * along with this program. If not, see + * [https://www.gnu.org/licenses/](https://www.gnu.org/licenses/). */ -/** - * Options accepted by `HttpError`. All optional. - */ +/** Options accepted by `HttpError`. All optional. */ export interface HttpErrorOptions { /** Underlying error. Set as the standard `Error.cause`. */ cause?: unknown; @@ -89,25 +88,38 @@ export interface HttpErrorOptions { legacyCode?: LegacyErrorCodes | (string & {}); /** * Modern, structured error code. If both `legacyCode` and `code` are set, - * the legacy one takes the `code` slot in the response body and `code` - * is emitted as `errorCode`, so clients keying on either field find - * what they expect. + * the legacy one takes the `code` slot in the response body and `code` is + * emitted as `errorCode`, so clients keying on either field find what they + * expect. */ code?: string; /** Additional fields merged into the response body. */ fields?: Record; + /** + * Skip the terminal alarm gate for this error. Set it only where the call + * site already knows the failure is expected and arrives in proportion to + * traffic — an upstream rate limit on a zero-cost model, say. Never + * serialized to the client. + */ + noAlarm?: boolean; } /** * The single error type controllers and services throw to surface an HTTP - * failure. The terminal `errorHandler` middleware catches it, serializes a - * JSON body, and sets the response status. + * failure. The terminal `errorHandler` middleware catches it, serializes a JSON + * body, and sets the response status. * * Usage: + * * ```ts * throw new HttpError(404, 'Item not found'); - * throw new HttpError(409, 'Cannot overwrite directory', { legacyCode: 'is_directory' }); - * throw new HttpError(403, 'Forbidden', { legacyCode: 'forbidden', fields: { target } }); + * throw new HttpError(409, 'Cannot overwrite directory', { + * legacyCode: 'is_directory', + * }); + * throw new HttpError(403, 'Forbidden', { + * legacyCode: 'forbidden', + * fields: { target }, + * }); * ``` * * Express 5 forwards thrown errors (sync and async) to error-handling @@ -118,6 +130,7 @@ export class HttpError extends Error { readonly legacyCode?: LegacyErrorCodes | (string & {}); readonly code?: string; readonly fields?: Record; + readonly noAlarm?: boolean; constructor( statusCode: number, @@ -133,6 +146,7 @@ export class HttpError extends Error { this.legacyCode = options.legacyCode; this.code = options.code; this.fields = options.fields; + this.noAlarm = options.noAlarm; } } diff --git a/src/backend/drivers/ai-chat/ChatCompletionDriver.edges.test.ts b/src/backend/drivers/ai-chat/ChatCompletionDriver.edges.test.ts index 02790c1c6..7096a683d 100644 --- a/src/backend/drivers/ai-chat/ChatCompletionDriver.edges.test.ts +++ b/src/backend/drivers/ai-chat/ChatCompletionDriver.edges.test.ts @@ -148,10 +148,13 @@ const completeFake = (args: Record) => } as never), ); -const errorFor = async (thrown: unknown): Promise => { +const errorFor = async ( + thrown: unknown, + args: Record = {}, +): Promise => { vi.spyOn(FakeChatProvider.prototype, 'complete').mockRejectedValue(thrown); try { - await completeFake({}); + await completeFake(args); } catch (e) { return e as HttpError; } @@ -257,6 +260,32 @@ describe('ChatCompletionDriver exhausted-chain classification', () => { expect(err.message).toBe('AI provider rate limit exceeded'); }); + it('mutes the alarm when every rate-limited attempt was on a free model', async () => { + // `fake` is priced at zero throughout: an upstream throttle there is + // expected and costs nobody anything, so the caller still gets the + // 429 but nothing is recorded. + const err = await errorFor( + Object.assign(new Error('slow down'), { status: 429 }), + ); + expect(err.noAlarm).toBe(true); + }); + + it('keeps the alarm when a paid model is the one being rate limited', async () => { + const err = await errorFor( + Object.assign(new Error('slow down'), { status: 429 }), + { model: 'costly' }, + ); + expect(err).toMatchObject({ legacyCode: 'upstream_rate_limited' }); + expect(err.noAlarm).toBe(false); + }); + + it('leaves non-rate-limit failures on a free model alarming as usual', async () => { + const err = await errorFor( + Object.assign(new Error('invalid api key'), { status: 401 }), + ); + expect(err.noAlarm).toBeFalsy(); + }); + it('classifies a rate limit reported only in the message text', async () => { const err = await errorFor(new Error('Quota exceeded for this key')); expect(err.statusCode).toBe(429); @@ -411,9 +440,9 @@ describe('ChatCompletionDriver cross-provider fallback', () => { Object.assign(new Error('azure down'), { status: 503 }), ); const openai = vi.spyOn(OpenAiChatProvider.prototype, 'complete'); - vi.spyOn(server.services.metering, 'hasEnoughCredits') - .mockResolvedValueOnce(true) // pre-flight - .mockResolvedValue(false); // drained by a parallel request + vi.spyOn(server.services.metering, 'getRemainingUsage') + .mockResolvedValueOnce(1_000_000) // pre-flight + .mockResolvedValue(0); // drained by a parallel request await expect(completeShared()).rejects.toMatchObject({ statusCode: 402, diff --git a/src/backend/drivers/ai-chat/ChatCompletionDriver.metering.test.ts b/src/backend/drivers/ai-chat/ChatCompletionDriver.metering.test.ts new file mode 100644 index 000000000..9609f08ea --- /dev/null +++ b/src/backend/drivers/ai-chat/ChatCompletionDriver.metering.test.ts @@ -0,0 +1,508 @@ +/* + * Copyright (C) 2024-present Puter Technologies Inc. + * + * This file is part of Puter. + * + * Puter is free software: you can redistribute it and/or modify + * it under the terms of the GNU Affero General Public License as published + * by the Free Software Foundation, either version 3 of the License, or + * (at your option) any later version. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * GNU Affero General Public License for more details. + * + * You should have received a copy of the GNU Affero General Public License + * along with this program. If not, see . + */ + +/** + * What a chat completion is charged when the usual path doesn't complete. + * + * Providers meter from the usage they hand to `chatStream.end`, which is the + * last thing a stream does — so every way a stream can stop early is a way a + * completion the upstream billed us for reaches the account as free. These + * tests pin the driver's backstop: output that was produced gets charged, and + * output that was never produced doesn't. + */ +import type { Readable } from 'node:stream'; +import { + afterAll, + afterEach, + beforeAll, + beforeEach, + describe, + expect, + it, + vi, +} from 'vitest'; + +import type { UsageInput } from '../../services/metering/types.js'; +import { PuterServer } from '../../server.js'; +import { setupTestServer } from '../../testUtil.js'; +import { withTestActor } from '../integrationTestUtil.js'; +import { ChatCompletionDriver } from './ChatCompletionDriver.js'; +import { FakeChatProvider } from './providers/FakeChatProvider.js'; +import type { IChatCompleteResult } from './types.js'; +import type { AIChatStream } from './utils/Streaming.js'; + +let server: PuterServer; + +const PRICED_MODEL = { + id: 'priced', + aliases: [], + costs_currency: 'usd-cents', + costs: { input_tokens: 1000, output_tokens: 2000 }, + max_tokens: 8192, +}; + +const makeDriver = async () => { + const d = new ChatCompletionDriver( + { providers: { ollama: { enabled: false } } } as never, + server.clients, + server.stores, + server.services, + ); + d.onServerStart(); + for (let i = 0; i < 200; i++) { + const m = await d.models(); + if (m.length > 0) return d; + await new Promise((r) => setTimeout(r, 5)); + } + throw new Error('ChatCompletionDriver model map never populated in test'); +}; + +const drain = async (stream: Readable): Promise => { + for await (const _chunk of stream as AsyncIterable) { + void _chunk; + } + // The driver meters in the `finally` of the fire-and-forget stream pump, + // which can land a tick after the last byte. + await new Promise((r) => setTimeout(r, 20)); +}; + +/** Every usage entry the driver recorded, flattened across batches. */ +const meteredUsages = ( + spy: ReturnType, +): UsageInput[] => + spy.mock.calls.flatMap((call) => (call[1] as UsageInput[]) ?? []); + +beforeAll(async () => { + server = await setupTestServer(); +}); + +afterAll(async () => { + await server?.shutdown(); +}); + +let driver: ChatCompletionDriver; + +beforeEach(async () => { + vi.spyOn(FakeChatProvider.prototype, 'models').mockResolvedValue([ + PRICED_MODEL, + ] as never); + driver = await makeDriver(); +}); + +afterEach(() => { + vi.restoreAllMocks(); +}); + +const streamOf = ( + init: (args: { chatStream: AIChatStream }) => Promise, +): IChatCompleteResult => + ({ + init_chat_stream: init, + stream: true, + }) as unknown as IChatCompleteResult; + +const startStream = async (driverUnderTest: ChatCompletionDriver) => + (await withTestActor(() => + driverUnderTest.complete({ + model: 'priced', + messages: [{ role: 'user', content: 'write me something' }], + stream: true, + }), + )) as unknown as { stream: Readable }; + +describe('ChatCompletionDriver streaming metering backstop', () => { + it('charges an estimate when a stream dies after producing output', async () => { + const metered = vi.spyOn(server.services.metering, 'batchIncrementUsages'); + + vi.spyOn(FakeChatProvider.prototype, 'complete').mockResolvedValueOnce( + streamOf(async ({ chatStream }) => { + const message = chatStream.message(); + const block = message.contentBlock({ type: 'text' }); + block.addText('x'.repeat(4000)); + throw new Error('upstream died mid-response'); + }) as never, + ); + + const result = await startStream(driver); + await drain(result.stream); + + const estimated = meteredUsages(metered).filter((u) => + u.usageType.includes('estimated_'), + ); + expect(estimated.length).toBe(2); + const output = estimated.find((u) => + u.usageType.endsWith('estimated_output_tokens'), + )!; + // 4000 characters at 4 chars/token, priced at 2000 ucents/token. + expect(output.usageAmount).toBe(1000); + expect(output.costOverride).toBe(2_000_000); + }); + + it('charges an estimate when a stream completes without a usage report', async () => { + const metered = vi.spyOn(server.services.metering, 'batchIncrementUsages'); + + vi.spyOn(FakeChatProvider.prototype, 'complete').mockResolvedValueOnce( + streamOf(async ({ chatStream }) => { + const message = chatStream.message(); + const block = message.contentBlock({ type: 'text' }); + block.addText('y'.repeat(2000)); + // Provider sent no usage chunk — `end` with nothing to report. + chatStream.end(undefined as never); + }) as never, + ); + + const result = await startStream(driver); + await drain(result.stream); + + const estimated = meteredUsages(metered).filter((u) => + u.usageType.includes('estimated_'), + ); + expect(estimated.length).toBe(2); + }); + + it('does not charge when the stream failed before producing anything', async () => { + const metered = vi.spyOn(server.services.metering, 'batchIncrementUsages'); + + vi.spyOn(FakeChatProvider.prototype, 'complete').mockResolvedValueOnce( + streamOf(async () => { + throw new Error('upstream refused the request'); + }) as never, + ); + + const result = await startStream(driver); + await drain(result.stream); + + expect( + meteredUsages(metered).filter((u) => + u.usageType.includes('estimated_'), + ), + ).toHaveLength(0); + }); + + // The provider metered, then died before `chatStream.end` — the one + // path where the backstop used to charge a second, estimated time on + // top of the real usage. + it('does not double-charge a stream whose provider metered before it threw', async () => { + const metered = vi.spyOn(server.services.metering, 'batchIncrementUsages'); + + vi.spyOn(FakeChatProvider.prototype, 'complete').mockResolvedValueOnce( + streamOf(async ({ chatStream }) => { + const message = chatStream.message(); + const block = message.contentBlock({ type: 'text' }); + block.addText('w'.repeat(4000)); + // Real usage was metered here (providers report the moment + // they meter)... + chatStream.reportUsage({ input_tokens: 10, output_tokens: 1000 }); + // ...then the stream died before reaching chatStream.end — + // e.g. a malformed tool-call payload failing to parse. + throw new Error('malformed tool-call payload'); + }) as never, + ); + + const result = await startStream(driver); + await drain(result.stream); + + expect( + meteredUsages(metered).filter((u) => + u.usageType.includes('estimated_'), + ), + ).toHaveLength(0); + }); + + // `end({})` meters nothing — an empty usage object must not pass for a + // usage report, or a stream full of output goes out billed at zero. + it('charges the estimate when the usage report is empty', async () => { + const metered = vi.spyOn(server.services.metering, 'batchIncrementUsages'); + + vi.spyOn(FakeChatProvider.prototype, 'complete').mockResolvedValueOnce( + streamOf(async ({ chatStream }) => { + const message = chatStream.message(); + const block = message.contentBlock({ type: 'text' }); + block.addText('v'.repeat(2000)); + chatStream.end({} as never); + }) as never, + ); + + const result = await startStream(driver); + await drain(result.stream); + + expect( + meteredUsages(metered).filter((u) => + u.usageType.includes('estimated_'), + ), + ).toHaveLength(2); + }); + + it('leaves a provider-reported stream alone', async () => { + const metered = vi.spyOn(server.services.metering, 'batchIncrementUsages'); + + vi.spyOn(FakeChatProvider.prototype, 'complete').mockResolvedValueOnce( + streamOf(async ({ chatStream }) => { + const message = chatStream.message(); + const block = message.contentBlock({ type: 'text' }); + block.addText('z'.repeat(2000)); + chatStream.end({ input_tokens: 10, output_tokens: 500 }); + }) as never, + ); + + const result = await startStream(driver); + await drain(result.stream); + + expect( + meteredUsages(metered).filter((u) => + u.usageType.includes('estimated_'), + ), + ).toHaveLength(0); + }); +}); + +const freeUser = () => ({ + user: { + uuid: `chat-gate-${Math.random().toString(36).slice(2)}`, + username: 'chat-gate-user', + email: 'chat-gate@test.com', + }, +}); + +// The gate against the real MeteringService, no mocks: a free account whose +// month has already outrun its allowance must not reach a provider again. +describe('ChatCompletionDriver credit gate against real metering', () => { + + it('rejects a free account that has already spent its allowance', 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, + allowance + 1, + ); + expect(await metering.getRemainingUsage(actor)).toBe(0); + + const completeSpy = vi.spyOn(FakeChatProvider.prototype, 'complete'); + + await expect( + withTestActor( + () => + driver.complete({ + model: 'priced', + messages: [{ role: 'user', content: 'one more' }], + }), + actor, + ), + ).rejects.toMatchObject({ + statusCode: 402, + legacyCode: 'insufficient_funds', + }); + expect(completeSpy).not.toHaveBeenCalled(); + }); + + // The incident this exists for: usage is recorded when a completion + // finishes, so a second request that starts while the first is still + // running used to read a balance that had nothing in flight subtracted + // from it, and was told it could spend the whole thing too. Concurrency, + // not budget, was what bounded the spend. + it('does not let a second request spend a balance the first already has in flight', async () => { + const actor = freeUser() as never; + const metering = server.services.metering; + // Spent down to where one completion's worst case is the whole of + // what's left — the shape of an expensive model against a small + // allowance, which is when parallel requests overshoot. + const allowance = (await metering.getActorSubscription(actor)) + .monthUsageAllowance; + await metering.incrementUsage( + actor, + 'test:prior-spend', + 1, + Math.floor(allowance * 0.9), + ); + + let releaseFirst: (v: unknown) => void = () => {}; + const firstInFlight = new Promise((r) => { + releaseFirst = r; + }); + const completeSpy = vi + .spyOn(FakeChatProvider.prototype, 'complete') + .mockImplementationOnce( + async () => + firstInFlight.then(() => ({ + message: { + role: 'assistant', + content: [{ type: 'text', text: 'ok' }], + }, + usage: { input_tokens: 1, output_tokens: 1 }, + finish_reason: 'stop', + })) as never, + ); + + const first = withTestActor( + () => + driver.complete({ + model: 'priced', + messages: [{ role: 'user', content: 'first request' }], + }), + actor, + ); + // Let the first request clear the gate and take its hold. + await vi.waitFor(() => expect(completeSpy).toHaveBeenCalledTimes(1)); + + await expect( + withTestActor( + () => + driver.complete({ + model: 'priced', + messages: [{ role: 'user', content: 'second request' }], + }), + actor, + ), + ).rejects.toMatchObject({ + statusCode: 402, + legacyCode: 'insufficient_funds', + }); + // The second request never reached a provider. + expect(completeSpy).toHaveBeenCalledTimes(1); + + releaseFirst(undefined); + await first; + + // And the hold is given back, so the account can spend again. + expect( + await server.services.metering.getRemainingUsage(actor), + ).toBeGreaterThan(0); + }); + + it('caps output to the balance left, so one call cannot run away with the month', async () => { + const actor = freeUser() as never; + const metering = server.services.metering; + const allowance = (await metering.getActorSubscription(actor)) + .monthUsageAllowance; + + // Nine tenths spent — a tenth of the allowance is left to bound the + // next completion's output. + 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({ + message: { + role: 'assistant', + content: [{ type: 'text', text: 'ok' }], + }, + usage: { input_tokens: 1, output_tokens: 1 }, + finish_reason: 'stop', + } as never); + + 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 }; + // 2000 ucents per output token — the cap has to fit what's left. + expect(passed.max_tokens).toBeDefined(); + expect(passed.max_tokens! * 2000).toBeLessThanOrEqual(remaining); + }); +}); + +describe('ChatCompletionDriver credit gate on multimodal prompts', () => { + it('prices attachments into the pre-flight affordability check', async () => { + const actor = freeUser() as never; + const metering = server.services.metering; + const allowance = (await metering.getActorSubscription(actor)) + .monthUsageAllowance; + + // Leave half of what the frame below estimates to: a plain-text + // prompt is still affordable, the same prompt carrying the frame is + // not. (~150KB of base64 payload ≈ 1000 tokens at 1000 ucents each.) + await metering.incrementUsage( + actor, + 'test:prior-spend', + 1, + allowance - 500_000, + ); + + const completeSpy = vi + .spyOn(FakeChatProvider.prototype, 'complete') + .mockResolvedValue({ + message: { + role: 'assistant', + content: [{ type: 'text', text: 'ok' }], + }, + usage: { input_tokens: 1, output_tokens: 1 }, + finish_reason: 'stop', + } as never); + + // Text alone clears the gate... + await withTestActor( + () => + driver.complete({ + model: 'priced', + messages: [{ role: 'user', content: 'describe' }], + max_tokens: 10, + }), + actor, + ); + expect(completeSpy).toHaveBeenCalledTimes(1); + + // ...the same prompt carrying a frame that used to price as ~nothing + // does not. + await expect( + withTestActor( + () => + driver.complete({ + model: 'priced', + messages: [ + { + role: 'user', + content: [ + { type: 'text', text: 'describe' }, + { + type: 'image_url', + image_url: { + url: `data:image/jpeg;base64,${'A'.repeat(200_000)}`, + }, + }, + ], + }, + ], + max_tokens: 10, + }), + actor, + ), + ).rejects.toMatchObject({ + statusCode: 402, + legacyCode: 'insufficient_funds', + }); + expect(completeSpy).toHaveBeenCalledTimes(1); + }); +}); diff --git a/src/backend/drivers/ai-chat/ChatCompletionDriver.routing.test.ts b/src/backend/drivers/ai-chat/ChatCompletionDriver.routing.test.ts index 25330bc3f..13da0ea45 100644 --- a/src/backend/drivers/ai-chat/ChatCompletionDriver.routing.test.ts +++ b/src/backend/drivers/ai-chat/ChatCompletionDriver.routing.test.ts @@ -293,6 +293,73 @@ describe('ChatCompletionDriver duplicate-model fallback', () => { }); }); +// Each attempt in a fallback chain is a whole completion at that model's +// prices, so each one has to clear the credit gate on its own. The chain used +// to re-check with a nominal 1-microcent amount, which any account with a +// fraction of a credit left passed — three attempts could cost three times +// what the balance allowed. +describe('ChatCompletionDriver credit gate across the fallback chain', () => { + it('runs the full gate once per attempt, not a nominal re-check', async () => { + // Observed, not stubbed — the counter that shows each attempt did a + // real balance read of its own. + const remaining = vi.spyOn( + server.services.metering, + 'getRemainingUsage', + ); + + const attempts = await attemptsFor('deepseek-v4-pro'); + + expect(attempts).toHaveLength(3); + // The balance is read for each attempt because each one's + // affordability and output cap are decided at that model's prices + // against what's actually left. + expect(remaining.mock.calls.length).toBeGreaterThanOrEqual( + attempts.length, + ); + remaining.mockRestore(); + }); + + it('aborts the chain when the balance runs out mid-fallback', async () => { + const actor = { + user: { + uuid: `routing-gate-${Math.random().toString(36).slice(2)}`, + username: 'routing-gate-user', + email: 'routing-gate@test.com', + }, + } as never; + const metering = server.services.metering; + + // The first attempt fails upstream, and while it does, a "parallel + // request" spends the rest of the month's allowance — real usage + // rows, not a stubbed balance — so the next attempt's gate reads an + // account with nothing left. + createMock.mockImplementation(async () => { + const { remaining } = await metering.getAllowedUsage(actor); + await metering.incrementUsage( + actor, + 'test:parallel-spend', + 1, + remaining, + ); + throw new Error('upstream down'); + }); + + await expect( + withTestActor( + () => + driver.complete({ + model: 'deepseek-v4-pro', + messages: [{ role: 'user', content: 'hi' }], + }), + actor, + ), + ).rejects.toMatchObject({ + statusCode: 402, + legacyCode: 'insufficient_funds', + }); + }); +}); + describe('ChatCompletionDriver unhealthy-route skipping', () => { it('skips a route marked by an earlier failure and serves the next one', async () => { markRouteUnhealthy('deepseek', 'deepseek-v4-pro'); diff --git a/src/backend/drivers/ai-chat/ChatCompletionDriver.test.ts b/src/backend/drivers/ai-chat/ChatCompletionDriver.test.ts index b4d1a65b0..133abdd62 100644 --- a/src/backend/drivers/ai-chat/ChatCompletionDriver.test.ts +++ b/src/backend/drivers/ai-chat/ChatCompletionDriver.test.ts @@ -481,8 +481,8 @@ describe('ChatCompletionDriver.complete validation event routing', () => { describe('ChatCompletionDriver.complete credit gate and max_tokens cap', () => { it('throws 402 `insufficient_funds` when the actor has no remaining credits', async () => { - vi.spyOn(server.services.metering, 'hasEnoughCredits').mockResolvedValue( - false, + vi.spyOn(server.services.metering, 'getRemainingUsage').mockResolvedValue( + 0, ); await expect( @@ -562,11 +562,8 @@ describe('ChatCompletionDriver.complete credit gate and max_tokens cap', () => { ]); const d = await makeDriver(); - // Pass the cheap pre-flight gate but leave a balance too small to + // Pass the cheap pre-flight check but leave a balance too small to // afford a single 2000-microcent output token. - vi.spyOn(server.services.metering, 'hasEnoughCredits').mockResolvedValue( - true, - ); vi.spyOn(server.services.metering, 'getRemainingUsage').mockResolvedValue( 100, ); @@ -610,10 +607,6 @@ describe('ChatCompletionDriver.complete credit gate and max_tokens cap', () => { ] as never); const d = await makeDriver(); - vi.spyOn( - server.services.metering, - 'hasEnoughCredits', - ).mockResolvedValue(true); vi.spyOn( server.services.metering, 'getRemainingUsage', @@ -658,9 +651,6 @@ describe('ChatCompletionDriver.complete credit gate and max_tokens cap', () => { ]); const d = await makeDriver(); // Plenty of credits so the credit gate doesn't intercept first. - vi.spyOn(server.services.metering, 'hasEnoughCredits').mockResolvedValue( - true, - ); vi.spyOn(server.services.metering, 'getRemainingUsage').mockResolvedValue( 1_000_000, ); @@ -768,19 +758,20 @@ describe('ChatCompletionDriver.complete fallback and error envelope', () => { }); }); - it('re-checks `hasEnoughCredits` between fallback attempts so a parallel request that drains the wallet aborts the chain', async () => { - // The primary provider throws; the fallback loop checks credits - // before its next upstream hit. We force `false` on the second - // check to verify the 402 short-circuit, even though no actual - // fallback model is wired (the loop bails on the credit gate - // before `#findFallback` decides there's nowhere to go). + it('re-reads the balance between fallback attempts so a parallel request that drains the wallet aborts the chain', async () => { + // The primary provider throws; the fallback loop runs the full gate + // (one balance read per attempt) before its next upstream hit. We + // force an empty balance on the second read to verify the 402 + // short-circuit, even though no actual fallback model is wired (the + // loop bails on the credit gate before `#findFallback` decides + // there's nowhere to go). vi.spyOn(FakeChatProvider.prototype, 'complete').mockRejectedValueOnce( new Error('boom'), ); - const credits = vi - .spyOn(server.services.metering, 'hasEnoughCredits') - .mockResolvedValueOnce(true) // pre-flight - .mockResolvedValueOnce(false); // mid-fallback re-check + const remaining = vi + .spyOn(server.services.metering, 'getRemainingUsage') + .mockResolvedValueOnce(1_000_000) // pre-flight + .mockResolvedValueOnce(0); // drained mid-fallback // No second provider serves `fake`, so `#findFallback` returns // null and the loop exits before reaching the credit re-check. @@ -795,7 +786,7 @@ describe('ChatCompletionDriver.complete fallback and error envelope', () => { }), ), ).rejects.toMatchObject({ statusCode: 500 }); - expect(credits.mock.calls.length).toBeGreaterThanOrEqual(1); + expect(remaining.mock.calls.length).toBeGreaterThanOrEqual(1); }); }); diff --git a/src/backend/drivers/ai-chat/ChatCompletionDriver.ts b/src/backend/drivers/ai-chat/ChatCompletionDriver.ts index 233636523..8ca8347a5 100644 --- a/src/backend/drivers/ai-chat/ChatCompletionDriver.ts +++ b/src/backend/drivers/ai-chat/ChatCompletionDriver.ts @@ -20,12 +20,15 @@ import crypto from 'node:crypto'; import { PassThrough } from 'node:stream'; import { EventMap } from '../../clients/event/types.js'; +import type { Actor } from '../../core/actor.js'; import { Context } from '../../core/context.js'; -import { HttpError } from '../../core/http/HttpError.js'; +import { HttpError, isHttpError } from '../../core/http/HttpError.js'; import { DEFAULT_FREE_SUBSCRIPTION, DEFAULT_TEMP_SUBSCRIPTION, } from '../../services/metering/consts.js'; +import type { CreditHold } from '../../services/metering/types.js'; +import { NO_CREDIT_HOLD } from '../../services/metering/types.js'; import type { DriverStreamResult } from '../meta.js'; import { PuterDriver } from '../types.js'; import { AI_CONCURRENT, AI_RATE_LIMIT } from '../util/aiLimits.js'; @@ -57,7 +60,6 @@ import type { } from './types.js'; import { normalize_tools_object } from './utils/FunctionCalling.js'; import { - extract_text, normalize_messages, normalize_single_message, } from './utils/Messages.js'; @@ -66,14 +68,35 @@ import { isIdentityKey, normalizeModelKey, } from './utils/modelRouting.js'; +import { costKeys, isFreeModel } from './utils/pricing.js'; import { isRouteUnhealthy, markRouteUnhealthy, } from './utils/providerHealth.js'; import { AIChatStream } from './utils/Streaming.js'; +import { + estimateOutputTokens, + estimatePromptTokens, +} from './utils/usageEstimate.js'; const MAX_ATTEMPTS = 3; // the first attempt plus two fallbacks +/** + * How often a streaming completion renews its credit hold. Holds default to a + * 10-minute TTL; a long generation (a reasoning model with a large + * `max_tokens`) can stream past that, and a hold that expires mid-stream + * reopens the overspend window it exists to close. + */ +const HOLD_RENEW_INTERVAL_MS = 5 * 60 * 1000; + +/** + * A moderation refusal is a completion that was produced and charged, then + * withheld — not a route failure. Retrying it on a fallback provider would bill + * the account again for another completion the user will never see. + */ +const isModerationRefusal = (e: unknown): boolean => + isHttpError(e) && e.code === 'moderation_flagged'; + type ProviderAttempt = { model: string; provider: string; @@ -157,13 +180,17 @@ const routeId = (provider: string, modelId: string) => `${provider}:${modelId}`; * * Per-class rules (see also alarm gate in server.ts): * - * - All rate-limited → 429 `upstream_rate_limited` (paged: forced alert) + * - All rate-limited → 429 `upstream_rate_limited` (alerted, unless every attempt + * was on a free model — see `allModelsFree`) * - All auth failures → 500 `upstream_auth_failed` (paged: our config) * - All upstream 5xx → 400 `upstream_provider_unavailable` (no page) * - All upstream 4xx (other) → 400 `upstream_bad_request` (no page) * - Mixed → 400 `upstream_failed` (no page) */ -const classifyAttempts = (attempts: ProviderAttempt[]): HttpError => { +const classifyAttempts = ( + attempts: ProviderAttempt[], + { allModelsFree = false } = {}, +): HttpError => { const fields = { attempts }; if (attempts.length === 0) { return new HttpError(500, 'No providers attempted', { @@ -176,6 +203,11 @@ const classifyAttempts = (attempts: ProviderAttempt[]): HttpError => { return new HttpError(429, 'AI provider rate limit exceeded', { legacyCode: 'upstream_rate_limited', fields, + // A free model getting throttled upstream is the deal we took + // when we picked it up for nothing: there's no billing at stake + // and nothing to act on, and the volume tracks traffic. The + // caller still gets the 429; we just don't record it. + noAlarm: allModelsFree, }); } if (attempts.every(isAuthFailure)) { @@ -334,6 +366,15 @@ export class ChatCompletionDriver extends PuterDriver { normalize_tools_object(args.tools); } + // Both estimated once, before any attempt: providers rewrite + // `args.messages` in place (tool_use blocks move out of `content`), so + // an estimate taken after a failed attempt would undercount the same + // prompt — and the gate writes each attempt's output cap into + // `args.max_tokens`, so the user's requested value has to be kept + // apart from what the previous attempt was capped to. + const promptTokenEstimate = estimatePromptTokens(args.messages ?? []); + const requestedMaxTokens = args.max_tokens; + const completionId = crypto .randomUUID() .replaceAll('-', '') @@ -376,89 +417,20 @@ export class ChatCompletionDriver extends PuterDriver { } } - // -- Credit / subscription gates (metering) -------------------- - // Cheap pre-flight: reject when the user can't afford even the - // approximate input cost, keep subscriber-only models gated, and - // cap `max_tokens` so output can't exceed remaining credits. - // Skipped for blocked requests since fake-chat is free and the - // user shouldn't see a billing error in place of the abuse page. + // Skipped for blocked requests since fake-chat is free and the user + // shouldn't see a billing error in place of the abuse page. + // + // The gate hands back a hold on what this attempt could cost, which + // stands in for its usage until the real numbers land. It is released + // on every way out of this method — including the streaming path, + // where "done" is the stream draining rather than this method + // returning. + let hold: CreditHold = NO_CREDIT_HOLD; if (!blocked) { - const metering = this.services.metering; - const inputCostKey = - (model.input_cost_key as string | undefined) ?? 'input_tokens'; - const outputCostKey = - (model.output_cost_key as string | undefined) ?? - 'output_tokens'; - const inputTokenCost = Number(model.costs?.[inputCostKey] ?? 0); - const outputTokenCost = Number(model.costs?.[outputCostKey] ?? 0); - const text = extract_text(args.messages ?? []); - // Rough estimator from v1 — avg of char/4 and word*(4/3), halved. - // See https://help.openai.com/en/articles/4936856 - const approximateTokenCount = Math.floor( - (text.length / 4 + text.split(/\s+/).length * (4 / 3)) / 2, - ); - const approximateInputCost = approximateTokenCount * inputTokenCost; - const minimumCredits = Number(model.minimumCredits || 1); - - const usageAllowed = await metering.hasEnoughCredits( - actor, - Math.max(approximateInputCost, minimumCredits), - ); - if (!usageAllowed) { - throw new HttpError(402, 'No usage left for request.', { - legacyCode: 'insufficient_funds', - }); - } - - if (model.subscriberOnly) { - const subscription = await metering.getActorSubscription(actor); - const isDefaultPolicy = - subscription.id === DEFAULT_FREE_SUBSCRIPTION || - subscription.id === DEFAULT_TEMP_SUBSCRIPTION; - if (isDefaultPolicy) { - throw new HttpError( - 403, - `The model ${model.id} is only available to subscribers. Please subscribe to access this model.`, - { legacyCode: 'permission_denied' }, - ); - } - } - - if (outputTokenCost > 0) { - const remainingCredits = - await metering.getRemainingUsage(actor); - const maxAllowedOutputUcents = - remainingCredits - approximateInputCost; - const maxAllowedOutputTokens = - maxAllowedOutputUcents / outputTokenCost; - // A provider may not know a model's output ceiling. Drop the - // term rather than let a missing value drive the cap: `null` - // coerces to 0, so the subtraction goes negative instead of - // NaN and the user is told they're out of credits. - const modelOutputCeiling = - Number.isFinite(model.max_tokens) && model.max_tokens > 0 - ? model.max_tokens - approximateTokenCount - : Number.POSITIVE_INFINITY; - const cap = Math.floor( - Math.min( - args.max_tokens ?? Number.POSITIVE_INFINITY, - maxAllowedOutputTokens, - modelOutputCeiling, - ), - ); - // `cap` is the credit-bounded ceiling on output tokens. When - // it drops below 1 the user can't afford even a single output - // token, so reject the request. Crucially we must NOT leave - // `max_tokens` unset here: an undefined max_tokens lets the - // provider run to the model's full output limit (e.g. 128k for - // Claude), billing far past the user's remaining balance. - if (cap < 1) { - throw new HttpError(402, 'No usage left for request.', { - legacyCode: 'insufficient_funds', - }); - } - args.max_tokens = cap; - } + hold = await this.#applyCreditGate(actor, model, args, { + promptTokenEstimate, + requestedMaxTokens, + }); } // First attempt @@ -473,18 +445,18 @@ export class ChatCompletionDriver extends PuterDriver { const attempts: ProviderAttempt[] = []; let res: IChatCompleteResult | undefined; + // Tracked across the chain so the classifier can tell a chain that + // only ever touched free models from one that cost the user something. + let allModelsFree = true; // A failed route is remembered briefly so the next request skips it // rather than paying its timeout again. - const recordFailure = ( - modelId: string, - providerId: string, - err: unknown, - ) => { - const attempt = toAttempt(modelId, providerId, err); + const recordFailure = (failed: IChatModel, err: unknown) => { + const attempt = toAttempt(failed.id, failed.provider!, err); attempts.push(attempt); + if (!isFreeModel(failed)) allModelsFree = false; if (isRouteLevelFailure(attempt)) { - markRouteUnhealthy(providerId, modelId); + markRouteUnhealthy(failed.provider!, failed.id); } }; @@ -495,7 +467,15 @@ export class ChatCompletionDriver extends PuterDriver { provider: model.provider, }); } catch (e) { - recordFailure(model.id, model.provider!, e); + // This attempt is over and cost whatever it cost; the next one + // takes a hold of its own. + await hold.release(); + hold = NO_CREDIT_HOLD; + + // A withheld completion was still a completion — charged, final, + // not a route failure worth another (billed) attempt elsewhere. + if (isModerationRefusal(e)) throw e; + recordFailure(model, e); // Fallback loop — the bucket holds every provider that serves this // model, ranked by `compareModelPreference`, so each miss walks one @@ -511,14 +491,18 @@ export class ChatCompletionDriver extends PuterDriver { const fbProvider = this.#providers[fallback.provider!]; if (!fbProvider) break; - // Credits can be exhausted mid-fallback by parallel requests; - // re-check before another upstream hit. Same bail as the - // pre-flight above. - const fallbackUsageAllowed = - await this.services.metering.hasEnoughCredits(actor, 1); - if (!fallbackUsageAllowed) { - throw new HttpError(402, 'No usage left for request.', { - legacyCode: 'insufficient_funds', + // Every attempt is a whole completion the account pays for, so + // each one goes through the full gate again rather than a + // token "do they have anything left" check: the balance may + // have been spent by a parallel request, and the fallback is + // a different model at a different price, whose output has to + // be capped against what is actually left. + // The previous attempt released its hold when it failed, so + // this one starts from nothing held. + if (!blocked) { + hold = await this.#applyCreditGate(actor, fallback, args, { + promptTokenEstimate, + requestedMaxTokens, }); } @@ -533,14 +517,18 @@ export class ChatCompletionDriver extends PuterDriver { model = fallback; lastError = null; } catch (fbErr) { + await hold.release(); + hold = NO_CREDIT_HOLD; + if (isModerationRefusal(fbErr)) throw fbErr; lastError = fbErr as Error; - recordFailure(fallback.id, fallback.provider!, fbErr); + recordFailure(fallback, fbErr); } } } if (!res) { - throw classifyAttempts(attempts); + await hold.release(); + throw classifyAttempts(attempts, { allModelsFree }); } const username = actor.user?.username; @@ -588,6 +576,13 @@ export class ChatCompletionDriver extends PuterDriver { return originalEnd(enrichedUsage!); }; + // The hold lives for the whole stream, which can outlast its TTL — + // keep pushing the deadline out until the pump is done. + const renewHold = setInterval(() => { + void hold.extend?.(); + }, HOLD_RENEW_INTERVAL_MS); + renewHold.unref?.(); + // Fire-and-forget — the stream writes happen async while the // response is being piped to the client. (async () => { @@ -602,6 +597,26 @@ export class ChatCompletionDriver extends PuterDriver { ); passthrough.end(); } finally { + clearInterval(renewHold); + // Providers report usage the moment they meter it (see + // `AIChatStream.reportUsage`); a stream that never got + // there was never charged for. + if (!blocked && !chatStream.reportedUsage) { + this.#meterUnreportedStream({ + actor, + chatStream, + model, + promptTokenEstimate, + completionId, + username, + intendedProvider, + }); + } + // Held until the generation is actually over: for a + // stream, the provider returns as soon as it has a + // populator, and everything the account pays for happens + // after that. + await hold.release(); if (cleanup) await cleanup(); } })(); @@ -615,6 +630,10 @@ export class ChatCompletionDriver extends PuterDriver { return streamResult as unknown as IChatCompleteResult; } + // The provider recorded this completion's usage before returning it, + // so the hold has served its purpose. + await hold.release(); + // -- Post-completion audit event ------------------------------ // Only for non-streaming results (streaming emits from the // `chatStream.end` wrapper above). Extensions like prompt_block / @@ -676,10 +695,7 @@ export class ChatCompletionDriver extends PuterDriver { outputMicroCents: number; totalMicroCents: number; } | null { - const inputKey = - (model.input_cost_key as string | undefined) ?? 'input_tokens'; - const outputKey = - (model.output_cost_key as string | undefined) ?? 'output_tokens'; + const { inputKey, outputKey } = costKeys(model); const costs = model.costs; if (!costs) return null; @@ -763,6 +779,199 @@ export class ChatCompletionDriver extends PuterDriver { }; } + /** + * The credit and subscription gate for one upstream attempt. + * + * Runs before every attempt, not once per request: each attempt is a whole + * completion the account pays for, at that model's prices, against whatever + * balance is left by the time it starts. + * + * Rejects when the account can't afford the approximate input cost, keeps + * subscriber-only models gated, and tightens `args.max_tokens` so the + * output this attempt can produce is bounded by the remaining balance. + * + * Returns a hold on what the attempt can cost at worst, so requests this + * account is running in parallel see the spend before it is recorded. The + * caller releases it once the attempt is done. + */ + async #applyCreditGate( + actor: Actor, + model: IChatModel, + args: ICompleteArguments, + estimates: { + /** + * Prompt tokens, estimated once before any attempt — counts + * attachments as well as text, and predates any in-place message + * rewriting a previous attempt's provider did. + */ + promptTokenEstimate: number; + /** + * What the user asked for, kept apart from `args.max_tokens`, which + * carries the previous attempt's cap: a cheap fallback must not + * inherit the ceiling computed at an expensive model's price. + */ + requestedMaxTokens: number | undefined; + }, + ): Promise { + const metering = this.services.metering; + const { promptTokenEstimate, requestedMaxTokens } = estimates; + const { inputKey, outputKey } = costKeys(model); + // `|| 0` also catches NaN from a malformed cost table. + const inputTokenCost = Number(model.costs?.[inputKey] ?? 0) || 0; + const outputTokenCost = Number(model.costs?.[outputKey] ?? 0) || 0; + const approximateInputCost = promptTokenEstimate * inputTokenCost; + const minimumCredits = Number(model.minimumCredits || 1); + + // One balance read serves the whole gate: the affordability check + // here and the output cap below. + const remainingCredits = await metering.getRemainingUsage(actor); + if (remainingCredits < Math.max(approximateInputCost, minimumCredits)) { + throw new HttpError(402, 'No usage left for request.', { + legacyCode: 'insufficient_funds', + }); + } + + if (model.subscriberOnly) { + const subscription = await metering.getActorSubscription(actor); + const isDefaultPolicy = + subscription.id === DEFAULT_FREE_SUBSCRIPTION || + subscription.id === DEFAULT_TEMP_SUBSCRIPTION; + if (isDefaultPolicy) { + throw new HttpError( + 403, + `The model ${model.id} is only available to subscribers. Please subscribe to access this model.`, + { legacyCode: 'permission_denied' }, + ); + } + } + + if (outputTokenCost > 0) { + const maxAllowedOutputUcents = + remainingCredits - approximateInputCost; + const maxAllowedOutputTokens = + maxAllowedOutputUcents / outputTokenCost; + // A provider may not know a model's output ceiling. Drop the term + // rather than let a missing value drive the cap: `null` coerces to + // 0, so the subtraction goes negative instead of NaN and the user + // is told they're out of credits. + const modelOutputCeiling = + Number.isFinite(model.max_tokens) && model.max_tokens > 0 + ? model.max_tokens - promptTokenEstimate + : Number.POSITIVE_INFINITY; + const cap = Math.floor( + Math.min( + requestedMaxTokens ?? Number.POSITIVE_INFINITY, + maxAllowedOutputTokens, + modelOutputCeiling, + ), + ); + // `cap` is the credit-bounded ceiling on output tokens. When it + // drops below 1 the user can't afford even a single output token, + // so reject the request. Crucially we must NOT leave `max_tokens` + // unset here: an undefined max_tokens lets the provider run to the + // model's full output limit (e.g. 128k for Claude), billing far + // past the user's remaining balance. + if (cap < 1) { + throw new HttpError(402, 'No usage left for request.', { + legacyCode: 'insufficient_funds', + }); + } + args.max_tokens = cap; + } else { + // No output price, nothing to bound — but a previous attempt may + // have written its cap here; give this one the user's own value. + args.max_tokens = requestedMaxTokens; + } + + // 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. + const worstCaseCost = + approximateInputCost + (args.max_tokens ?? 0) * outputTokenCost; + return this.services.metering.reserveCredits( + actor, + Math.max(worstCaseCost, minimumCredits), + ); + } + + /** + * Charge a stream that produced output but never reported usage. + * + * Providers meter from the usage they hand to `chatStream.end`, at the very + * end of the stream — so anything that stops the stream short of that point + * (an upstream error mid-response, a malformed tool-call payload, a + * provider that never sends a usage chunk) leaves a completion the upstream + * has already billed us for and the account has paid nothing for. This is + * the backstop: what the stream actually emitted, priced off the model's + * own cost table. + * + * Only when there was output. A stream that failed before producing + * anything cost the user nothing, and charging an estimated prompt to + * someone whose request we failed to serve is worse than the leak. + * + * Recorded under `estimated_*` usage keys so the numbers stay separable + * from provider-reported ones in the usage breakdown. + */ + #meterUnreportedStream(params: { + actor: Actor; + chatStream: AIChatStream; + model: IChatModel; + /** Estimated before any provider rewrote `args.messages` in place. */ + promptTokenEstimate: number; + completionId: string; + username?: string; + intendedProvider: string; + }): void { + const { + actor, + chatStream, + model, + promptTokenEstimate, + completionId, + username, + intendedProvider, + } = params; + + const outputTokens = estimateOutputTokens(chatStream.outputChars ?? 0); + if (outputTokens <= 0) return; + + const { inputKey, outputKey } = costKeys(model); + const inputTokens = promptTokenEstimate; + const usage = { + [inputKey]: inputTokens, + [outputKey]: outputTokens, + }; + + const cost = this.#computeCost(usage, model); + this.services.metering.utilRecordUsageObject( + { + [`estimated_${inputKey}`]: inputTokens, + [`estimated_${outputKey}`]: outputTokens, + }, + actor, + `${model.provider}:${model.id}`, + { + // Undefined when the model has no cost table: the entry is + // recorded unpriced rather than free. + [`estimated_${inputKey}`]: cost?.inputMicroCents, + [`estimated_${outputKey}`]: cost?.outputMicroCents, + }, + ); + + console.warn( + `[ai-chat] stream ended without usage; charged an estimate (${completionId}, ${model.provider}:${model.id}, ~${inputTokens} in / ~${outputTokens} out)`, + ); + + this.#emitCostCalculated({ + completionId, + username, + usage, + model, + intendedProvider, + }); + } + // Add `usd_cents` to the usage object. Skips if the provider already // set an authoritative value (e.g. OpenRouter's `usage.cost`). // Sets `null` when cost data is unavailable for the model. @@ -795,10 +1004,7 @@ export class ChatCompletionDriver extends PuterDriver { params; const cost = this.#computeCost(usage, model); - const inputKey = - (model.input_cost_key as string | undefined) ?? 'input_tokens'; - const outputKey = - (model.output_cost_key as string | undefined) ?? 'output_tokens'; + const { inputKey, outputKey } = costKeys(model); const inputTokens = cost?.inputTokens ?? 0; const outputTokens = cost?.outputTokens ?? 0; const inputMicroCents = cost?.inputMicroCents ?? 0; diff --git a/src/backend/drivers/ai-chat/providers/claude/ClaudeProvider.ts b/src/backend/drivers/ai-chat/providers/claude/ClaudeProvider.ts index f67970e26..bbba0873e 100644 --- a/src/backend/drivers/ai-chat/providers/claude/ClaudeProvider.ts +++ b/src/backend/drivers/ai-chat/providers/claude/ClaudeProvider.ts @@ -536,7 +536,10 @@ export class ClaudeProvider implements IChatProvider { } } } - chatStream.end(usageSum); + // Metered before `end`: handing usage to `end` is what tells + // the driver this completion has been charged for, so anything + // that throws between the two would leave it charged to + // nobody. const costsOverrideFromModel = this.#buildCostsOverrideFromModel(usageSum, modelUsed); this.#meteringService.utilRecordUsageObject( @@ -545,6 +548,7 @@ export class ClaudeProvider implements IChatProvider { `claude:${modelUsed.id}`, costsOverrideFromModel, ); + chatStream.end(usageSum); }; return { diff --git a/src/backend/drivers/ai-chat/utils/OpenAIUtil.js b/src/backend/drivers/ai-chat/utils/OpenAIUtil.js index b67dee361..00e7c3569 100644 --- a/src/backend/drivers/ai-chat/utils/OpenAIUtil.js +++ b/src/backend/drivers/ai-chat/utils/OpenAIUtil.js @@ -5,26 +5,27 @@ import { HttpError } from '@heyputer/backend/src/core/http'; * * This file is part of Puter. * - * Puter is free software: you can redistribute it and/or modify - * it under the terms of the GNU Affero General Public License as published - * by the Free Software Foundation, either version 3 of the License, or - * (at your option) any later version. + * Puter is free software: you can redistribute it and/or modify it under the + * terms of the GNU Affero General Public License as published by the Free + * Software Foundation, either version 3 of the License, or (at your option) any + * later version. * - * This program is distributed in the hope that it will be useful, - * but WITHOUT ANY WARRANTY; without even the implied warranty of - * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the - * GNU Affero General Public License for more details. + * This program is distributed in the hope that it will be useful, but WITHOUT + * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS + * FOR A PARTICULAR PURPOSE. See the GNU Affero General Public License for more + * details. * * You should have received a copy of the GNU Affero General Public License - * along with this program. If not, see . + * along with this program. If not, see + * [https://www.gnu.org/licenses/](https://www.gnu.org/licenses/). */ /** - * Process input messages from Puter's normalized format to OpenAI's format - * May make changes in-place. + * Process input messages from Puter's normalized format to OpenAI's format May + * make changes in-place. * - * @param {Array} messages - array of normalized messages - * @returns {Array} - array of messages in OpenAI format + * @param {Message[]} messages - Array of normalized messages + * @returns {Message[]} - Array of messages in OpenAI format */ export const process_input_messages = async (messages) => { for (const msg of messages) { @@ -357,10 +358,22 @@ export const create_chat_stream_handler = } // TODO DS: this is a bit too abstracted... this is basically just doing the metering now - const usage = usage_calculator({ - usage: last_usage, - extra_content: last_extra_content, - }); + // No usage chunk means there is nothing to meter from — reaching into + // a null usage object here used to throw, which took down a response + // the upstream had already produced *and* skipped its billing. Leave + // the usage undefined instead; the driver charges an estimate for a + // stream that produced output nobody reported. + const usage = last_usage + ? usage_calculator({ + usage: last_usage, + extra_content: last_extra_content, + }) + : undefined; + // The calculator just metered. Reported here, not only via `end`: + // a throw in the block flushes below (a malformed tool-call payload, + // say) must not leave a metered stream looking unmetered — the driver + // would charge its estimate on top. + chatStream.reportUsage(usage); if (mode === 'text') textblock.end(); if (mode === 'tool') toolblock.end(); @@ -433,7 +446,13 @@ export const create_chat_stream_handler_responses_api = } // TODO DS: this is a bit too abstracted... this is basically just doing the metering now - const usage = usage_calculator({ usage: last_usage }); + // Missing usage is left undefined rather than fed to the calculator — + // see the sibling handler above, including why usage is reported + // before the block flushes. + const usage = last_usage + ? usage_calculator({ usage: last_usage }) + : undefined; + chatStream.reportUsage(usage); if (mode === 'text') textblock.end(); if (mode === 'tool') toolblock.end(); @@ -443,7 +462,13 @@ export const create_chat_stream_handler_responses_api = }; export const handle_completion_output = async ( - /** @type {Record & {usage_calculator:(args: {usage: import("openai/resources/completions.mjs").CompletionUsage})=> unknown }}*/ + /** + * @type {Record & { + * usage_calculator: (args: { + * usage: import('openai/resources/completions.mjs').CompletionUsage; + * }) => unknown; + * }} + */ { deviations, stream, completion, moderate, usage_calculator, finally_fn }, ) => { deviations = Object.assign( @@ -470,17 +495,10 @@ export const handle_completion_output = async ( if (finally_fn) await finally_fn(); - // We need to moderate the completion too - const mod_text = completion.choices[0].message.content; - if (moderate && mod_text !== null) { - const moderation_result = await moderate(mod_text); - if (moderation_result.flagged) { - throw new HttpError(400, 'message is not allowed', { - legacyCode: 'bad_request', - }); - } - } - + // Metered before moderation: the completion exists and the upstream has + // billed us for it whether or not we go on to withhold it, and running + // the moderation gate first meant a flagged completion was served to + // nobody and charged to nobody. const ret = completion.choices[0]; const completion_usage = deviations.coerce_completion_usage(completion); ret.usage = usage_calculator @@ -492,13 +510,30 @@ export const handle_completion_output = async ( input_tokens: completion_usage.prompt_tokens, output_tokens: completion_usage.completion_tokens, }; + + const mod_text = completion.choices[0].message.content; + if (moderate && mod_text !== null) { + const moderation_result = await moderate(mod_text); + if (moderation_result.flagged) { + // `code` tells the driver this is a refusal of a completion that + // was produced and charged, not a route failure — retrying it on + // a fallback provider would bill the account again for another + // completion the user will never see. + throw new HttpError(400, 'message is not allowed', { + legacyCode: 'bad_request', + code: 'moderation_flagged', + }); + } + } + return ret; }; /** - * * @param {object} params - * @param {(args: {usage: import("openai/resources/completions.mjs").CompletionUsage})=> unknown } params.usage_calculator + * @param {(args: { + * usage: import('openai/resources/completions.mjs').CompletionUsage; + * }) => unknown} params.usage_calculator * @returns */ export const handle_completion_output_responses_api = async ({ @@ -559,17 +594,6 @@ export const handle_completion_output_responses_api = async ({ }); } - // We need to moderate the completion too - const mod_text = completion.output_text; - if (moderate && mod_text !== null) { - const moderation_result = await moderate(mod_text); - if (moderation_result.flagged) { - throw new HttpError(400, 'message is not allowed', { - legacyCode: 'bad_request', - }); - } - } - const ret = { finish_reason: 'stop', index: 0, @@ -599,6 +623,9 @@ export const handle_completion_output_responses_api = async ({ delete ret.type; + // Metered before moderation, same as the sibling handler above: the + // completion exists and the upstream has billed us for it whether or not + // we go on to withhold it. ret.usage = usage_calculator ? usage_calculator({ ...completion, @@ -608,5 +635,17 @@ export const handle_completion_output_responses_api = async ({ input_tokens: completion.usage.input_tokens, output_tokens: completion.usage.output_tokens, }; + + const mod_text = completion.output_text; + if (moderate && mod_text !== null) { + const moderation_result = await moderate(mod_text); + if (moderation_result.flagged) { + throw new HttpError(400, 'message is not allowed', { + legacyCode: 'bad_request', + code: 'moderation_flagged', + }); + } + } + return ret; }; diff --git a/src/backend/drivers/ai-chat/utils/Streaming.js b/src/backend/drivers/ai-chat/utils/Streaming.js index e82f373d2..7178bc2f0 100644 --- a/src/backend/drivers/ai-chat/utils/Streaming.js +++ b/src/backend/drivers/ai-chat/utils/Streaming.js @@ -27,28 +27,25 @@ export class AIChatConstructStream { export class AIChatTextStream extends AIChatConstructStream { addText(text, extra_content) { - const json = JSON.stringify({ + this.chatStream.writeChunk({ type: 'text', text, ...(extra_content ? { extra_content } : {}), }); - this.chatStream.stream.write(`${json}\n`); } addReasoning(reasoning) { - const json = JSON.stringify({ + this.chatStream.writeChunk({ type: 'reasoning', reasoning, }); - this.chatStream.stream.write(`${json}\n`); } addExtraContent(extra_content) { - const json = JSON.stringify({ + this.chatStream.writeChunk({ type: 'extra_content', extra_content, }); - this.chatStream.stream.write(`${json}\n`); } } @@ -58,6 +55,9 @@ export class AIChatToolUseStream extends AIChatConstructStream { this.buffer = ''; } addPartialJSON(partial_json) { + // Counted as it accumulates, not when the block is written out: a + // stream that dies mid-tool-call still produced these characters. + this.chatStream.countOutput(partial_json); this.buffer += partial_json; } end() { @@ -65,13 +65,15 @@ export class AIChatToolUseStream extends AIChatConstructStream { this.buffer = '{}'; } if (process.env.DEBUG) console.log('BUFFER BEING PARSED', this.buffer); - const str = JSON.stringify({ - type: 'tool_use', - ...this.contentBlock, - input: JSON.parse(this.buffer), - ...(!this.contentBlock.text ? { text: '' } : {}), - }); - this.chatStream.stream.write(`${str}\n`); + this.chatStream.writeChunk( + { + type: 'tool_use', + ...this.contentBlock, + input: JSON.parse(this.buffer), + ...(!this.contentBlock.text ? { text: '' } : {}), + }, + { alreadyCounted: true }, + ); } } @@ -87,13 +89,90 @@ export class AIChatMessageStream extends AIChatConstructStream { } } +// Chunk fields that are envelope rather than model output — everything else +// string-valued in an outgoing chunk was generated by the model and counts +// toward the billing estimate for a stream whose usage is never reported. +const ENVELOPE_KEYS = new Set(['type', 'id', 'usage']); + +/** Characters of model-authored payload in one outgoing chunk value. */ +const payloadChars = (value) => { + if (typeof value === 'string') return value.length; + if (Array.isArray(value)) { + return value.reduce((n, v) => n + payloadChars(v), 0); + } + if (value && typeof value === 'object') { + let n = 0; + for (const [key, v] of Object.entries(value)) { + if (!ENVELOPE_KEYS.has(key)) n += payloadChars(v); + } + return n; + } + return 0; +}; + export class AIChatStream { stream; + /** + * Characters of model output that have gone through this stream. + * + * The only measure of what a completion produced that survives a stream + * ending badly. Providers report usage at the end and meter from it; when + * that end is never reached, this is what the driver falls back to so a + * completion that was generated is a completion that gets charged. + */ + outputChars = 0; + + /** + * Usage the provider has metered, recorded the moment it is known rather + * than when `end` is reached. A throw between metering and `end` must not + * disguise a charged stream as an uncharged one — the driver would then + * charge its estimate on top of the real usage. + * + * @type {Record | null} + */ + reportedUsage = null; + constructor({ stream }) { this.stream = stream; } + /** @param {string} text */ + countOutput(text) { + if (typeof text === 'string') this.outputChars += text.length; + } + + /** + * Mark this stream's usage as metered. A usage object with no finite + * numbers in it is not a report — nothing was recorded from it (a Claude + * stream that died before its usage events ends with `{}`, and the driver's + * `end` wrapper decorates that with `usd_cents: null`), so the driver's + * estimate is still the only charge the stream would ever get. + * + * @param {Record | undefined | null} usage + */ + reportUsage(usage) { + if (!usage) return; + if (Object.values(usage).some((v) => Number.isFinite(v))) { + this.reportedUsage = usage; + } + } + + /** + * Write one NDJSON chunk, counting its model-authored payload toward + * `outputChars` — counting lives at this choke point so a new chunk type is + * counted by default instead of arriving for free. Pass `alreadyCounted` + * when the payload was counted as it accumulated. + * + * @param {Record} chunk + * @param {{ alreadyCounted?: boolean }} [opts] + */ + writeChunk(chunk, { alreadyCounted = false } = {}) { + if (!alreadyCounted) this.outputChars += payloadChars(chunk); + this.stream.write(`${JSON.stringify(chunk)}\n`); + } + end(/** @type {Record} */ usage) { + this.reportUsage(usage); this.stream.write( `${JSON.stringify({ type: 'usage', @@ -113,13 +192,11 @@ export class AIChatStream { * @param {{ id?: string; encrypted_content: string }} compaction */ compaction({ id, encrypted_content }) { - this.stream.write( - `${JSON.stringify({ - type: 'compaction', - ...(id !== undefined ? { id } : {}), - encrypted_content, - })}\n`, - ); + this.writeChunk({ + type: 'compaction', + ...(id !== undefined ? { id } : {}), + encrypted_content, + }); } message() { diff --git a/src/backend/drivers/ai-chat/utils/pricing.test.ts b/src/backend/drivers/ai-chat/utils/pricing.test.ts index 4f1f01807..4acd8e51c 100644 --- a/src/backend/drivers/ai-chat/utils/pricing.test.ts +++ b/src/backend/drivers/ai-chat/utils/pricing.test.ts @@ -19,7 +19,7 @@ import { describe, expect, it } from 'vitest'; import type { IChatModel } from '../types.js'; -import { buildCostsOverride, usdPerMToken } from './pricing.js'; +import { buildCostsOverride, isFreeModel, usdPerMToken } from './pricing.js'; const model = (costs: Record): IChatModel => ({ @@ -43,6 +43,32 @@ describe('usdPerMToken', () => { }); }); +describe('isFreeModel', () => { + it('treats a table of nothing but zeroes as free', () => { + expect( + isFreeModel( + model({ + tokens: 1_000_000, + prompt_tokens: 0, + completion_tokens: 0, + }), + ), + ).toBe(true); + }); + + it('treats a model priced on any axis as paid', () => { + expect( + isFreeModel(model({ prompt_tokens: 0, completion_tokens: 200 })), + ).toBe(false); + expect(isFreeModel(model({ request: 1 }))).toBe(false); + }); + + it('treats a model with no cost data as unknown, not free', () => { + expect(isFreeModel(model({}))).toBe(false); + expect(isFreeModel(model({ tokens: 1_000_000 }))).toBe(false); + }); +}); + describe('buildCostsOverride', () => { it('multiplies each usage key by its own declared rate', () => { const overrides = buildCostsOverride( diff --git a/src/backend/drivers/ai-chat/utils/pricing.ts b/src/backend/drivers/ai-chat/utils/pricing.ts index 1d9f3fe3a..accf08ed0 100644 --- a/src/backend/drivers/ai-chat/utils/pricing.ts +++ b/src/backend/drivers/ai-chat/utils/pricing.ts @@ -42,6 +42,33 @@ export const usdPerMToken = ( const isRate = (value: unknown): value is number => typeof value === 'number' && Number.isFinite(value); +/** + * The usage keys a model's input and output are priced under. Most models use + * the defaults; a model whose provider reports usage under other names carries + * them in `input_cost_key`/`output_cost_key`. + */ +export const costKeys = ( + model: IChatModel, +): { inputKey: string; outputKey: string } => ({ + inputKey: (model.input_cost_key as string | undefined) ?? 'input_tokens', + outputKey: (model.output_cost_key as string | undefined) ?? 'output_tokens', +}); + +/** + * Whether a model costs the user nothing to run. + * + * Every rate in the cost table has to be zero — a model priced on one axis and + * free on another is a paid model. `tokens` is skipped: it's the scale the + * other numbers are expressed in, not a rate. A model with no cost table at all + * is _unknown_, not free, so it doesn't qualify. + */ +export const isFreeModel = (model: IChatModel): boolean => { + const rates = Object.entries(model.costs ?? {}).filter( + ([key]) => key !== 'tokens', + ); + return rates.length > 0 && rates.every(([, rate]) => Number(rate) === 0); +}; + /** * Prices a tracked-usage object against a model's cost table. * @@ -57,10 +84,7 @@ export const buildCostsOverride = ( trackedUsage: Record, model: IChatModel, ): Record => { - const inputKey = - (model.input_cost_key as string | undefined) ?? 'input_tokens'; - const outputKey = - (model.output_cost_key as string | undefined) ?? 'output_tokens'; + const { inputKey, outputKey } = costKeys(model); const costs = model.costs ?? {}; const inputRate = isRate(costs[inputKey]) ? costs[inputKey] : undefined; diff --git a/src/backend/drivers/ai-chat/utils/usageEstimate.test.ts b/src/backend/drivers/ai-chat/utils/usageEstimate.test.ts new file mode 100644 index 000000000..87b6bd348 --- /dev/null +++ b/src/backend/drivers/ai-chat/utils/usageEstimate.test.ts @@ -0,0 +1,187 @@ +/* + * Copyright (C) 2024-present Puter Technologies Inc. + * + * This file is part of Puter. + * + * Puter is free software: you can redistribute it and/or modify + * it under the terms of the GNU Affero General Public License as published + * by the Free Software Foundation, either version 3 of the License, or + * (at your option) any later version. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * GNU Affero General Public License for more details. + * + * You should have received a copy of the GNU Affero General Public License + * along with this program. If not, see . + */ + +import { describe, expect, it } from 'vitest'; +import { + estimateOutputTokens, + estimatePromptTokens, + estimateTextTokens, +} from './usageEstimate.js'; + +// A ~150KB base64 payload — the shape a client sends video frames in. +const base64Image = (bytes: number) => + `data:image/jpeg;base64,${'A'.repeat(Math.ceil(bytes / (3 / 4)))}`; + +describe('estimateTextTokens', () => { + it('is zero for empty text', () => { + expect(estimateTextTokens('')).toBe(0); + }); + + it('scales with length', () => { + const short = estimateTextTokens('hello world'); + const long = estimateTextTokens('hello world '.repeat(100)); + expect(long).toBeGreaterThan(short); + }); +}); + +describe('estimatePromptTokens', () => { + it('counts plain string content', () => { + const tokens = estimatePromptTokens([ + { role: 'user', content: 'a fairly ordinary sentence to price' }, + ]); + expect(tokens).toBeGreaterThan(0); + }); + + it('counts normalized text parts', () => { + const tokens = estimatePromptTokens([ + { + role: 'user', + content: [{ type: 'text', text: 'describe this frame' }], + }, + ]); + expect(tokens).toBeGreaterThan(0); + }); + + // The leak this estimator exists for: a prompt made almost entirely of + // image data used to price as an empty one, so an account with nothing + // left could still send it. + it('prices an inline image on its payload size, not as free', () => { + const tokens = estimatePromptTokens([ + { + role: 'user', + content: [ + { type: 'text', text: 'what is happening here?' }, + { + type: 'image_url', + image_url: { url: base64Image(150_000) }, + }, + ], + }, + ]); + expect(tokens).toBeGreaterThan(900); + }); + + it('scales with the number of frames attached', () => { + const frame = { + type: 'image_url', + image_url: { url: base64Image(150_000) }, + }; + const one = estimatePromptTokens([ + { role: 'user', content: [frame] }, + ]); + const ten = estimatePromptTokens([ + { role: 'user', content: Array.from({ length: 10 }, () => frame) }, + ]); + expect(ten).toBeGreaterThan(one * 9); + }); + + it('prices Anthropic-style base64 sources', () => { + const tokens = estimatePromptTokens([ + { + role: 'user', + content: [ + { + type: 'image', + source: { + type: 'base64', + media_type: 'image/png', + data: 'A'.repeat(200_000), + }, + }, + ], + }, + ]); + expect(tokens).toBeGreaterThan(900); + }); + + it('charges a flat estimate for attachments it cannot measure', () => { + const remote = estimatePromptTokens([ + { + role: 'user', + content: [ + { + type: 'image_url', + image_url: { url: 'https://example.com/cat.png' }, + }, + ], + }, + ]); + const fsRef = estimatePromptTokens([ + { + role: 'user', + content: [{ puter_path: '/user/Desktop/scan.pdf' }], + }, + ]); + expect(remote).toBeGreaterThan(0); + expect(fsRef).toBeGreaterThan(0); + }); + + it('counts tool traffic as the text it serializes to', () => { + const tokens = estimatePromptTokens([ + { + role: 'assistant', + content: [ + { + type: 'tool_use', + id: 't1', + name: 'search', + input: { query: 'a'.repeat(4000) }, + }, + ], + }, + ]); + expect(tokens).toBeGreaterThan(100); + }); + + // Model-output shapes come back around as history in multi-turn + // conversations. They are text; falling through to the attachment + // default would price a few sentences of reasoning like a full frame + // and 402 accounts that could easily afford their prompt. + it('prices replayed thinking/reasoning/refusal blocks as their text', () => { + const sentence = 'a short run of replayed reasoning text'; + const asText = estimatePromptTokens([ + { role: 'assistant', content: [{ type: 'text', text: sentence }] }, + ]); + for (const part of [ + { type: 'thinking', thinking: sentence }, + { type: 'reasoning', reasoning: sentence }, + { type: 'refusal', refusal: sentence }, + ]) { + expect( + estimatePromptTokens([{ role: 'assistant', content: [part] }]), + ).toBe(asText); + } + }); + + it('is zero for nothing at all', () => { + expect(estimatePromptTokens([])).toBe(0); + expect(estimatePromptTokens(undefined)).toBe(0); + }); +}); + +describe('estimateOutputTokens', () => { + it('is zero when nothing was emitted', () => { + expect(estimateOutputTokens(0)).toBe(0); + expect(estimateOutputTokens(Number.NaN)).toBe(0); + }); + + it('converts characters to tokens', () => { + expect(estimateOutputTokens(400)).toBe(100); + }); +}); diff --git a/src/backend/drivers/ai-chat/utils/usageEstimate.ts b/src/backend/drivers/ai-chat/utils/usageEstimate.ts new file mode 100644 index 000000000..1883b9106 --- /dev/null +++ b/src/backend/drivers/ai-chat/utils/usageEstimate.ts @@ -0,0 +1,194 @@ +/* + * Copyright (C) 2024-present Puter Technologies Inc. + * + * This file is part of Puter. + * + * Puter is free software: you can redistribute it and/or modify + * it under the terms of the GNU Affero General Public License as published + * by the Free Software Foundation, either version 3 of the License, or + * (at your option) any later version. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * GNU Affero General Public License for more details. + * + * You should have received a copy of the GNU Affero General Public License + * along with this program. If not, see . + */ + +/** + * Token estimates for prompts we haven't sent yet and completions the upstream + * never told us the size of. + * + * Two things use these and neither can be exact: the credit gate, which has to + * decide whether a request is affordable before anyone has counted its tokens, + * and the fallback charge for a stream that produced output but no usage + * report. Both are better served by a rough number than by the zero they would + * otherwise use — a prompt whose cost is estimated at nothing is a prompt that + * passes every gate. + */ + +import { estimateTextTokens } from '../../util/tokenEstimate.js'; + +// Re-exported so existing consumers keep one import site for all estimators. +export { estimateTextTokens } from '../../util/tokenEstimate.js'; + +/** Tokens per byte of an inline (base64) attachment; see `attachmentTokens`. */ +const ATTACHMENT_BYTES_PER_TOKEN = 150; + +/** + * What an attachment we can't measure is assumed to cost — a remote URL, or a + * `puter_path` we haven't read yet. Roughly a full-frame image at Anthropic's + * `(w*h)/750`, which is the largest an image gets before providers downscale. + */ +const UNMEASURED_ATTACHMENT_TOKENS = 1500; + +/** Base64 carries 3 bytes in every 4 characters. */ +const BASE64_BYTES_PER_CHAR = 3 / 4; + +/** Characters per token, for output we only ever saw as text. */ +const OUTPUT_CHARS_PER_TOKEN = 4; + +type ContentPart = Record; + +/** Bytes an inline data URI or bare base64 payload decodes to. */ +const inlineBytes = (value: string): number => { + const comma = value.startsWith('data:') ? value.indexOf(',') : -1; + const payload = comma === -1 ? value : value.slice(comma + 1); + return Math.floor(payload.length * BASE64_BYTES_PER_CHAR); +}; + +/** + * Tokens an attachment is worth, from whatever we can see of it. + * + * Inline payloads are measured; anything referenced by URL or FS path is + * charged the flat unmeasured estimate rather than nothing, because "we can't + * see it" and "it's free" are not the same answer. + */ +const attachmentTokens = (value: unknown): number => { + if (typeof value !== 'string' || value === '') { + return UNMEASURED_ATTACHMENT_TOKENS; + } + if (value.startsWith('data:') || !/^[a-z][a-z0-9+.-]*:/i.test(value)) { + // A data URI, or something that isn't a URI at all — a bare base64 + // payload or an FS path. Only the first is measurable; a short string + // measures short, which the floor below covers. + const bytes = inlineBytes(value); + return Math.max( + Math.ceil(bytes / ATTACHMENT_BYTES_PER_TOKEN), + value.startsWith('data:') ? 0 : UNMEASURED_ATTACHMENT_TOKENS, + ); + } + return UNMEASURED_ATTACHMENT_TOKENS; +}; + +/** + * Tokens one normalized content part is worth. + * + * Every part shape any provider accepts has to land somewhere here: the text + * ones on the text estimator, the rest on `attachmentTokens`. A part nobody + * recognises is charged the unmeasured estimate — the alternative is a content + * type nobody has taught the gate about arriving for free. + */ +const partTokens = (part: unknown): number => { + if (typeof part === 'string') return estimateTextTokens(part); + if (!part || typeof part !== 'object') return 0; + + const p = part as ContentPart; + + if (typeof p.text === 'string') return estimateTextTokens(p.text); + + // Model-output shapes replayed as history: Anthropic thinking blocks, + // our own streamed `reasoning` chunks, OpenAI refusals. All text — they + // must not fall through to the attachment default below, which would + // price a few sentences of reasoning like a full-frame image. + if (typeof p.thinking === 'string') return estimateTextTokens(p.thinking); + if (typeof p.reasoning === 'string') { + return estimateTextTokens(p.reasoning); + } + if (typeof p.refusal === 'string') return estimateTextTokens(p.refusal); + + // OpenAI-style `{ image_url: { url } }` (or the flattened string form). + if (p.image_url) { + const url = + typeof p.image_url === 'string' + ? p.image_url + : (p.image_url as ContentPart).url; + return attachmentTokens(url); + } + + // Anthropic-style `{ source: { data | url } }`. + if (p.source && typeof p.source === 'object') { + const source = p.source as ContentPart; + return attachmentTokens(source.data ?? source.url ?? source.file_id); + } + + // Puter's own reference — resolved to an upload by the provider, so its + // size isn't knowable here. + if (p.puter_path) return UNMEASURED_ATTACHMENT_TOKENS; + + if (typeof p.data === 'string' || typeof p.b64_json === 'string') { + return attachmentTokens(p.data ?? p.b64_json); + } + + // Tool traffic: arguments and results are text once serialized. + if (p.type === 'tool_use' && p.input !== undefined) { + return estimateTextTokens( + typeof p.input === 'string' ? p.input : JSON.stringify(p.input), + ); + } + if (p.type === 'tool_result') { + return estimateTextTokens( + typeof p.content === 'string' + ? p.content + : JSON.stringify(p.content ?? ''), + ); + } + if (p.type === 'compaction') { + return estimateTextTokens(String(p.encrypted_content ?? '')); + } + + return UNMEASURED_ATTACHMENT_TOKENS; +}; + +/** + * Tokens a prompt is worth, across every content part it carries. + * + * Text-only prompts land on the same v1 number the gate has always used. What + * changes for a multimodal prompt is that its attachments count for something: + * reading only `text` fields made a request carrying twenty frames of video + * look like an empty one, and an empty request is affordable to an account with + * nothing left. + */ +export const estimatePromptTokens = (messages: unknown): number => { + if (!Array.isArray(messages)) return 0; + + let tokens = 0; + for (const message of messages) { + if (typeof message === 'string') { + tokens += estimateTextTokens(message); + continue; + } + if (!message || typeof message !== 'object') continue; + + const content = (message as ContentPart).content; + if (typeof content === 'string') { + tokens += estimateTextTokens(content); + } else if (Array.isArray(content)) { + for (const part of content) tokens += partTokens(part); + } else if (content) { + tokens += partTokens(content); + } + } + return tokens; +}; + +/** + * Tokens a completion we only saw as streamed characters is worth. For a stream + * that ended without a usage report, this is what it gets billed on. + */ +export const estimateOutputTokens = (chars: number): number => { + if (!Number.isFinite(chars) || chars <= 0) return 0; + return Math.ceil(chars / OUTPUT_CHARS_PER_TOKEN); +}; diff --git a/src/backend/drivers/ai-image/providers/gemini/GeminiImageProvider.ts b/src/backend/drivers/ai-image/providers/gemini/GeminiImageProvider.ts index 69c6c70ef..03f96ea3a 100644 --- a/src/backend/drivers/ai-image/providers/gemini/GeminiImageProvider.ts +++ b/src/backend/drivers/ai-image/providers/gemini/GeminiImageProvider.ts @@ -32,6 +32,7 @@ import type { IImageProvider, } from '../../types.js'; import { isHttpUrl, toBase64DataUri } from '../../inputImage.js'; +import { estimateTextTokens } from '../../../util/tokenEstimate.js'; import { HttpError } from '@heyputer/backend/src/core/http/HttpError.js'; const MIME_SIGNATURES: Record = { @@ -441,12 +442,7 @@ export class GeminiImageProvider implements IImageProvider { if (text.length === 0) return 0; // Same approximation used by chat billing flow. - return Math.max( - 1, - Math.floor( - (text.length / 4 + text.split(/\s+/).length * (4 / 3)) / 2, - ), - ); + return Math.max(1, estimateTextTokens(text)); } #calculateTokenCostInCents( diff --git a/src/backend/drivers/ai-image/providers/openai/OpenAiImageProvider.ts b/src/backend/drivers/ai-image/providers/openai/OpenAiImageProvider.ts index 169f7c63f..2af9d27df 100644 --- a/src/backend/drivers/ai-image/providers/openai/OpenAiImageProvider.ts +++ b/src/backend/drivers/ai-image/providers/openai/OpenAiImageProvider.ts @@ -32,6 +32,7 @@ import type { } from '../../types.js'; import { OPEN_AI_IMAGE_GENERATION_MODELS } from './models.js'; import { fetchImageAsBase64, isHttpUrl } from '../../inputImage.js'; +import { estimateTextTokens } from '../../../util/tokenEstimate.js'; import { HttpError } from '@heyputer/backend/src/core/http/HttpError.js'; interface OpenAIImageUsage { @@ -45,9 +46,9 @@ interface OpenAIImageUsage { } /** - * OpenAI image generation provider for v2. - * Supports the GPT Image models (gpt-image-1, -1-mini, -1.5, -2), including - * image-to-image editing via `input_images` (the `images.edit` endpoint). + * OpenAI image generation provider for v2. Supports the GPT Image models + * (gpt-image-1, -1-mini, -1.5, -2), including image-to-image editing via + * `input_images` (the `images.edit` endpoint). */ export class OpenAiImageProvider implements IImageProvider { #meteringService: MeteringService; @@ -464,12 +465,7 @@ export class OpenAiImageProvider implements IImageProvider { if (text.length === 0) return 0; // Same approximation used by chat and Gemini image billing flows. - return Math.max( - 1, - Math.floor( - (text.length / 4 + text.split(/\s+/).length * (4 / 3)) / 2, - ), - ); + return Math.max(1, estimateTextTokens(text)); } #getCostRate(selectedModel: IImageModel, key: string): number | undefined { diff --git a/src/backend/drivers/util/tokenEstimate.ts b/src/backend/drivers/util/tokenEstimate.ts new file mode 100644 index 000000000..ef6c65810 --- /dev/null +++ b/src/backend/drivers/util/tokenEstimate.ts @@ -0,0 +1,37 @@ +/* + * Copyright (C) 2024-present Puter Technologies Inc. + * + * This file is part of Puter. + * + * Puter is free software: you can redistribute it and/or modify + * it under the terms of the GNU Affero General Public License as published + * by the Free Software Foundation, either version 3 of the License, or + * (at your option) any later version. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * GNU Affero General Public License for more details. + * + * You should have received a copy of the GNU Affero General Public License + * along with this program. If not, see . + */ + +/** + * Tokens a run of text is worth. + * + * The v1 estimator: average the two cheap approximations (characters over four, + * words times four thirds) and halve the result. It runs low on purpose — it + * gates and estimates rather than prices, and the real count arrives from the + * provider a moment later. The one estimator is shared by the chat credit gate, + * the chat unreported-stream backstop, and image-prompt pricing, so a future + * calibration moves all of them at once. + * + * @see https://help.openai.com/en/articles/4936856 + */ +export const estimateTextTokens = (text: string): number => { + if (!text) return 0; + return Math.floor( + (text.length / 4 + text.split(/\s+/).length * (4 / 3)) / 2, + ); +}; diff --git a/src/backend/server.ts b/src/backend/server.ts index c90e9e4eb..382e802a4 100644 --- a/src/backend/server.ts +++ b/src/backend/server.ts @@ -757,6 +757,11 @@ export class PuterServer { // direction: an error tagged as caused by an upstream // provider or a misbehaving client gets exposed to // the user but does not alarm at all. + // + // `noAlarm` beats both: the call site already decided + // this failure isn't worth recording (e.g. an upstream + // rate limit on a free model, where the volume tracks + // traffic and there's nothing to act on). const FORCED_ALERT_CODES = new Map([ ['upstream_rate_limited', 'info'], // Our credentials for a provider stopped working — @@ -765,6 +770,7 @@ export class PuterServer { ]); const SKIP_ALERT_PREFIXES = /^(upstream_|client_)/; const isHttp = isHttpError(err); + if (isHttp && err.noAlarm) return; const status = isHttp ? err.statusCode : 500; const legacyCode = isHttp ? (err.legacyCode ?? '') : ''; const forcedSeverity = FORCED_ALERT_CODES.get(legacyCode); diff --git a/src/backend/services/metering/MeteringService.test.ts b/src/backend/services/metering/MeteringService.test.ts index 171f2ae3e..92c5386ce 100644 --- a/src/backend/services/metering/MeteringService.test.ts +++ b/src/backend/services/metering/MeteringService.test.ts @@ -1991,6 +1991,94 @@ describe('MeteringService', () => { }); }); + // -- Credit holds -------------------------------------------------- + + describe('reserveCredits', () => { + it('takes what an in-flight operation could spend out of the spendable balance', async () => { + const before = await target.getRemainingUsage(actor); + expect(before).toBeGreaterThan(0); + + const hold = await target.reserveCredits(actor, 1000); + + expect(await target.getRemainingUsage(actor)).toBe(before - 1000); + await hold.release(); + expect(await target.getRemainingUsage(actor)).toBe(before); + }); + + it('stacks holds, so parallel operations see each other', async () => { + const before = await target.getRemainingUsage(actor); + + const first = await target.reserveCredits(actor, 400); + const second = await target.reserveCredits(actor, 600); + + expect(await target.getRemainingUsage(actor)).toBe(before - 1000); + await first.release(); + await second.release(); + }); + + it('never reports a negative balance, however much is held', async () => { + const before = await target.getRemainingUsage(actor); + const hold = await target.reserveCredits(actor, before * 10); + + expect(await target.getRemainingUsage(actor)).toBe(0); + await hold.release(); + }); + + it('leaves the reported balance alone — a hold is not usage', async () => { + const { remaining } = await target.getAllowedUsage(actor); + const hold = await target.reserveCredits(actor, 1000); + + expect((await target.getAllowedUsage(actor)).remaining).toBe( + remaining, + ); + await hold.release(); + }); + + it('releasing twice gives the budget back once', async () => { + const before = await target.getRemainingUsage(actor); + const hold = await target.reserveCredits(actor, 500); + await hold.release(); + await hold.release(); + + expect(await target.getRemainingUsage(actor)).toBe(before); + }); + + it('holds nothing for the system actor or a zero amount', async () => { + const before = await target.getRemainingUsage(actor); + await (await target.reserveCredits(actor, 0)).release(); + expect(await target.getRemainingUsage(actor)).toBe(before); + }); + + // A stream can outlive the hold's TTL; extending is what keeps its + // in-flight spend visible for the whole generation. + it('extend gives a hold another full TTL from now', async () => { + const before = await target.getRemainingUsage(actor); + const hold = await target.reserveCredits(actor, 750, { + ttlMs: 500, + }); + + // Let the original deadline lapse entirely... + await new Promise((r) => setTimeout(r, 620)); + expect(await target.getRemainingUsage(actor)).toBe(before); + + // ...extending brings the still-running operation's hold back. + await hold.extend?.(); + expect(await target.getRemainingUsage(actor)).toBe(before - 750); + + await hold.release(); + expect(await target.getRemainingUsage(actor)).toBe(before); + }); + + it('extend after release does not resurrect the hold', async () => { + const before = await target.getRemainingUsage(actor); + const hold = await target.reserveCredits(actor, 300); + await hold.release(); + await hold.extend?.(); + + expect(await target.getRemainingUsage(actor)).toBe(before); + }); + }); + // ── Resolver registration ──────────────────────────────────────── describe('resolver registration', () => { diff --git a/src/backend/services/metering/MeteringService.ts b/src/backend/services/metering/MeteringService.ts index f9bcef3cd..ebace255d 100644 --- a/src/backend/services/metering/MeteringService.ts +++ b/src/backend/services/metering/MeteringService.ts @@ -35,11 +35,13 @@ import { import { EGRESS_COSTS } from './costs'; import type { AppTotals, + CreditHold, UsageAddons, UsageByType, UsageInput, UsageRecord, } from './types'; +import { NO_CREDIT_HOLD } from './types'; import { LOCAL_UNLIMITED_USER } from '../../data/subPolicies/localUnlimitedUserPolicy.js'; import { SUB_POLICIES } from '../../data/subPolicies/index.js'; @@ -1010,9 +1012,82 @@ export class MeteringService extends PuterService { return (res as UsageByType) || ({ total: 0 } as UsageByType); } + /** + * What an actor can commit to a new operation right now. + * + * Their balance less what other operations of theirs already have in flight + * (see `reserveCredits`) — which is the number a spend decision turns on, + * and is smaller than the balance whenever the actor has several requests + * running at once. `getAllowedUsage` is the one to read for reporting a + * balance; this one is for deciding on a spend. + */ async getRemainingUsage(actor: Actor): Promise { - const { remaining } = await this.getAllowedUsage(actor); - return remaining || 0; + const [{ remaining }, held] = await Promise.all([ + this.getAllowedUsage(actor), + this.#outstandingHolds(actor), + ]); + return Math.max(0, (remaining || 0) - held); + } + + /** + * Commit part of an actor's budget to an operation that is about to run. + * + * Usage is recorded when an operation finishes, so between starting and + * finishing it is invisible to every other request that account makes — + * they all read the same balance and are each told they can spend the whole + * of it. What an account can actually overspend by is then bounded by its + * concurrency limit rather than by its budget, which for an expensive model + * is several times the allowance. + * + * A hold makes the in-flight spend visible for as long as it lasts. Take + * one for what the operation could cost at worst, before the upstream call; + * release it once the real usage has been recorded. + * + * Never throws, and a hold that couldn't be taken is a no-op handle: not + * being able to reach the cache is our problem, and turning it into failed + * requests for everyone spending money is worse than the overshoot it would + * prevent. + */ + async reserveCredits( + actor: Actor, + amount: number, + opts: { ttlMs?: number } = {}, + ): Promise { + const userId = actor?.user?.uuid; + if (!userId || isSystemActor(actor) || !(amount > 0)) { + return NO_CREDIT_HOLD; + } + + const member = await this.stores.creditHold.take( + userId, + amount, + opts.ttlMs, + ); + if (!member) return NO_CREDIT_HOLD; + + let released = false; + return { + release: async () => { + if (released) return; + released = true; + await this.stores.creditHold.release(userId, member); + }, + extend: async () => { + if (released) return; + await this.stores.creditHold.refresh( + userId, + member, + opts.ttlMs, + ); + }, + }; + } + + /** Budget this actor has committed to requests that are still running. */ + async #outstandingHolds(actor: Actor): Promise { + const userId = actor?.user?.uuid; + if (!userId || isSystemActor(actor)) return 0; + return this.stores.creditHold.outstanding(userId); } async getAllowedUsage(actor: Actor): Promise<{ diff --git a/src/backend/services/metering/types.ts b/src/backend/services/metering/types.ts index 686c58724..3d5170a72 100644 --- a/src/backend/services/metering/types.ts +++ b/src/backend/services/metering/types.ts @@ -53,3 +53,27 @@ export interface AppTotals { total: number; count: number; } + +/** + * Budget committed to an operation that hasn't finished yet. + * + * Released by the code that took it, on every path out — including failure. A + * hold nobody releases expires on its own, so a lost release costs the account + * the use of that budget for a while rather than forever. + */ +export interface CreditHold { + release(): Promise; + /** + * Push the hold's deadline out for an operation still running — a stream + * can outlive the default TTL, and a hold that expires mid-operation + * reopens the overspend window it was taken to close. Absent on the no-op + * hold; callers renew with `hold.extend?.()`. + */ + extend?(): Promise; +} + +/** + * The hold that holds nothing — for paths that take no hold but still release + * one. + */ +export const NO_CREDIT_HOLD: CreditHold = { release: async () => {} }; diff --git a/src/backend/stores/index.ts b/src/backend/stores/index.ts index ccd6c9c19..c6e06b834 100644 --- a/src/backend/stores/index.ts +++ b/src/backend/stores/index.ts @@ -21,6 +21,7 @@ import { AppFeedbackStore } from './appFeedback/AppFeedbackStore.js'; import { AppStore } from './app/AppStore.js'; import { FSEntryStore } from './fs/FSEntryStore.js'; import { GroupStore } from './group/GroupStore.js'; +import { CreditHoldStore } from './metering/CreditHoldStore.js'; import { MeteringBufferStore } from './metering/MeteringBufferStore.js'; import { NotificationStore } from './notification/NotificationStore.js'; import { OIDCStore } from './oidc/OIDCStore.js'; @@ -44,6 +45,7 @@ declare module './types.js' { interface IPuterStoreInstances { kv: SystemKVStore; meteringBuffer: MeteringBufferStore; + creditHold: CreditHoldStore; user: UserStore; app: AppStore; appFeedback: AppFeedbackStore; @@ -71,6 +73,7 @@ declare module './types.js' { export const puterStores = { kv: SystemKVStore, meteringBuffer: MeteringBufferStore, + creditHold: CreditHoldStore, user: UserStore, app: AppStore, appFeedback: AppFeedbackStore, diff --git a/src/backend/stores/metering/CreditHoldStore.test.ts b/src/backend/stores/metering/CreditHoldStore.test.ts new file mode 100644 index 000000000..1acb8ce10 --- /dev/null +++ b/src/backend/stores/metering/CreditHoldStore.test.ts @@ -0,0 +1,144 @@ +/* + * Copyright (C) 2024-present Puter Technologies Inc. + * + * This file is part of Puter. + * + * Puter is free software: you can redistribute it and/or modify + * it under the terms of the GNU Affero General Public License as published + * by the Free Software Foundation, either version 3 of the License, or + * (at your option) any later version. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * GNU Affero General Public License for more details. + * + * You should have received a copy of the GNU Affero General Public License + * along with this program. If not, see . + */ + +import { afterAll, beforeAll, describe, expect, it, vi } from 'vitest'; + +import { PuterServer } from '../../server.js'; +import { setupTestServer } from '../../testUtil.js'; +import type { CreditHoldStore } from './CreditHoldStore.js'; + +let server: PuterServer; +let store: CreditHoldStore; + +const user = () => `hold-user-${Math.random().toString(36).slice(2)}`; + +beforeAll(async () => { + server = await setupTestServer(); + store = server.stores.creditHold as CreditHoldStore; +}); + +afterAll(async () => { + await server?.shutdown(); +}); + +describe('CreditHoldStore', () => { + it('sums what an actor has taken', async () => { + const id = user(); + expect(await store.outstanding(id)).toBe(0); + + await store.take(id, 1000); + await store.take(id, 250); + + expect(await store.outstanding(id)).toBe(1250); + }); + + it('gives budget back on release', async () => { + const id = user(); + const first = await store.take(id, 1000); + await store.take(id, 500); + + await store.release(id, first); + + expect(await store.outstanding(id)).toBe(500); + }); + + it('releases at most once', async () => { + const id = user(); + const member = await store.take(id, 900); + await store.release(id, member); + await store.release(id, member); + + expect(await store.outstanding(id)).toBe(0); + }); + + it('keeps actors separate', async () => { + const [a, b] = [user(), user()]; + await store.take(a, 700); + + expect(await store.outstanding(b)).toBe(0); + }); + + // A deployment that dies mid-request never releases; without a deadline + // the account would be short that budget until the month turned over. + it('drops holds nobody released once they expire', async () => { + const id = user(); + await store.take(id, 5000, 50); + expect(await store.outstanding(id)).toBe(5000); + + await vi.waitFor( + async () => { + expect(await store.outstanding(id)).toBe(0); + }, + { timeout: 2000, interval: 25 }, + ); + }); + + // All of an actor's holds share one redis key, so a short-lived hold must + // never pull the key's expiry in under a longer-lived one — key expiry + // drops every hold in the set at once. + it('never shortens the hold set’s expiry under a longer-lived hold', async () => { + const id = user(); + await store.take(id, 1000); // default 10-minute TTL + await store.take(id, 500, 50); // much shorter TTL + + const redis = server.clients.redis as unknown as { + pttl: (key: string) => Promise; + }; + const pttl = await redis.pttl(`meter:holds:{${id}}`); + // Still on the long hold's clock, not the short one's. + expect(pttl).toBeGreaterThan(500_000); + }); + + it('refresh pushes a hold’s deadline out for a still-running operation', async () => { + const id = user(); + const member = await store.take(id, 800, 50); + + await store.refresh(id, member); // default 10-minute TTL + + // Well past the original 50ms deadline, the hold is still counted. + await new Promise((r) => setTimeout(r, 100)); + expect(await store.outstanding(id)).toBe(800); + }); + + it('ignores amounts that aren’t worth holding', async () => { + const id = user(); + expect(await store.take(id, 0)).toBeNull(); + expect(await store.take(id, -5)).toBeNull(); + expect(await store.take(id, Number.NaN)).toBeNull(); + expect(await store.outstanding(id)).toBe(0); + }); + + // Holds gate spending, so an unreachable cache has to mean "no holds", not + // "no spending" — the alternative is a cache blip presenting as every + // account being out of credit. + it('reads zero rather than failing when the cache is unreachable', async () => { + const id = user(); + await store.take(id, 4000); + + const redis = server.clients.redis as unknown as { + creditHoldSum: () => Promise; + }; + const spy = vi + .spyOn(redis, 'creditHoldSum') + .mockRejectedValue(new Error('cache down')); + + expect(await store.outstanding(id)).toBe(0); + spy.mockRestore(); + }); +}); diff --git a/src/backend/stores/metering/CreditHoldStore.ts b/src/backend/stores/metering/CreditHoldStore.ts new file mode 100644 index 000000000..012623858 --- /dev/null +++ b/src/backend/stores/metering/CreditHoldStore.ts @@ -0,0 +1,233 @@ +/* + * Copyright (C) 2024-present Puter Technologies Inc. + * + * This file is part of Puter. + * + * Puter is free software: you can redistribute it and/or modify + * it under the terms of the GNU Affero General Public License as published + * by the Free Software Foundation, either version 3 of the License, or + * (at your option) any later version. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * GNU Affero General Public License for more details. + * + * You should have received a copy of the GNU Affero General Public License + * along with this program. If not, see . + */ + +import { randomUUID } from 'node:crypto'; +import { metrics } from '@opentelemetry/api'; +import { PuterStore } from '../types'; + +// -- Metrics ---------------------------------------------------------- + +const meter = metrics.getMeter('puter-backend'); + +const expiredHoldsCounter = meter.createCounter('metering.holds.expired', { + description: + 'Credit holds that reached their deadline without being released', +}); + +// -- Constants -------------------------------------------------------- + +/** + * How long a hold survives without being released. + * + * This is the exposure both ways: too short and a long completion stops being + * accounted for while it is still running, too long and a request killed + * without warning keeps budget locked up. It bounds the second case only — + * releases are what normally ends a hold, and every path that takes one + * releases it. + */ +const DEFAULT_HOLD_TTL_MS = 10 * 60 * 1000; + +/** Room for the key to outlive its last hold, so reads still prune it. */ +const KEY_TTL_SLACK_MS = 60 * 1000; + +const holdsKey = (userId: string): string => `meter:holds:{${userId}}`; + +// -- Scripts ---------------------------------------------------------- + +/** + * KEYS: the actor's hold set. ARGV: member, deadline, key ttl. + * + * Members are scored by when they expire, which is what lets a read drop the + * ones nothing released — a deployment that dies mid-request leaves its hold + * behind, and nothing else would ever take it off. + * + * The key TTL only ever moves out: all holds share one key, so a take with a + * short TTL must not truncate the key under a longer-lived hold — expiry would + * drop every hold in the set, not just the new one. (PEXPIRE's GT flag can't do + * this: it refuses to put a TTL on a key that has none, which is exactly the + * state ZADD leaves a fresh key in.) + */ +const TAKE_SCRIPT = ` +local ttl = redis.call('PTTL', KEYS[1]) +redis.call('ZADD', KEYS[1], ARGV[2], ARGV[1]) +if tonumber(ARGV[3]) > ttl then + redis.call('PEXPIRE', KEYS[1], ARGV[3]) +end +return 1 +`; + +/** + * KEYS: the actor's hold set. ARGV: now. + * + * Prunes expired members, then sums what's left. Amounts ride on the member + * name (`:`) so one sorted set carries both the deadline and the + * number, and pruning and summing happen in the same pass. + */ +const SUM_SCRIPT = ` +local dropped = redis.call('ZREMRANGEBYSCORE', KEYS[1], '-inf', ARGV[1]) +local total = 0 +local members = redis.call('ZRANGE', KEYS[1], 0, -1) +for i = 1, #members do + local sep = string.find(members[i], ':', 1, true) + if sep then + total = total + (tonumber(string.sub(members[i], sep + 1)) or 0) + end +end +return { tostring(total), dropped } +`; + +type ScriptRunner = { + creditHoldTake(...args: string[]): Promise; + creditHoldSum(...args: string[]): Promise<[string, number]>; + zrem(key: string, member: string): Promise; +}; + +// -- CreditHoldStore -------------------------------------------------- + +/** + * Budget committed to requests that are still running. + * + * A balance says what an account has spent, and spending is only recorded once + * an operation finishes — so several operations starting at once all read the + * same balance and each one is told it can afford the whole of it. What they + * then spend is bounded by the concurrency limit rather than by the budget. + * + * A hold closes that window: taken before the upstream call for what the call + * could cost at worst, subtracted from the balance every other request reads, + * and released when the call is done and its real usage is recorded. Holds live + * in the cache rather than in the usage record because they are not usage — + * nothing has been spent yet, they expire on their own, and every deployment + * has to see them at once for them to mean anything. + */ +export class CreditHoldStore extends PuterStore { + #definedScripts = false; + + get #redis(): ScriptRunner { + this.#defineScripts(); + return this.clients.redis as unknown as ScriptRunner; + } + + #defineScripts(): void { + if (this.#definedScripts) return; + this.#definedScripts = true; + const client = this.clients.redis; + client.defineCommand('creditHoldTake', { + numberOfKeys: 1, + lua: TAKE_SCRIPT, + }); + client.defineCommand('creditHoldSum', { + numberOfKeys: 1, + lua: SUM_SCRIPT, + }); + } + + /** + * Commit `amount` of an actor's budget to something about to run. + * + * Returns the member to release it with, or null when the cache wouldn't + * take it. Null means uncommitted, not failed: a cache that can't be + * reached is our problem, and refusing the request over it would turn a + * cache blip into an outage for everyone spending money. + */ + async take( + userId: string, + amount: number, + ttlMs: number = DEFAULT_HOLD_TTL_MS, + ): Promise { + if (!userId || !Number.isFinite(amount) || amount <= 0) return null; + const member = `${randomUUID()}:${Math.ceil(amount)}`; + try { + await this.#redis.creditHoldTake( + holdsKey(userId), + member, + String(Date.now() + ttlMs), + String(ttlMs + KEY_TTL_SLACK_MS), + ); + return member; + } catch (e) { + console.warn( + `[metering] credit hold not taken for ${userId}: ${(e as Error).message}`, + ); + return null; + } + } + + /** + * Push a hold's deadline out for an operation still running. + * + * Re-takes the same member, so a hold that was already pruned comes back — + * which is what should happen: the request it stands for is still going. + * Failure is tolerated the same way as `take`'s. + */ + async refresh( + userId: string, + member: string | null, + ttlMs: number = DEFAULT_HOLD_TTL_MS, + ): Promise { + if (!userId || !member) return; + try { + await this.#redis.creditHoldTake( + holdsKey(userId), + member, + String(Date.now() + ttlMs), + String(ttlMs + KEY_TTL_SLACK_MS), + ); + } catch (e) { + console.warn( + `[metering] credit hold not refreshed for ${userId}: ${(e as Error).message}`, + ); + } + } + + /** Give a hold back. Safe to call twice; the second call is a no-op. */ + async release(userId: string, member: string | null): Promise { + if (!userId || !member) return; + try { + await this.#redis.zrem(holdsKey(userId), member); + } catch (e) { + // The hold expires on its own, so the account gets its budget + // back either way — just later than it should have. + console.warn( + `[metering] credit hold not released for ${userId}: ${(e as Error).message}`, + ); + } + } + + /** + * What this actor currently has committed to running requests. Zero when + * the cache can't answer — see `take` for why that direction. + */ + async outstanding(userId: string): Promise { + if (!userId) return 0; + try { + const [total, dropped] = await this.#redis.creditHoldSum( + holdsKey(userId), + String(Date.now()), + ); + if (dropped > 0) expiredHoldsCounter.add(dropped); + const held = Number(total); + return Number.isFinite(held) && held > 0 ? held : 0; + } catch (e) { + console.warn( + `[metering] credit holds unreadable for ${userId}: ${(e as Error).message}`, + ); + return 0; + } + } +} diff --git a/src/docs/src/rate-limits-and-quotas.md b/src/docs/src/rate-limits-and-quotas.md index 49d3a0834..be28a6410 100644 --- a/src/docs/src/rate-limits-and-quotas.md +++ b/src/docs/src/rate-limits-and-quotas.md @@ -31,6 +31,10 @@ What usage costs (the big three): - **AI** — priced per model and per token/second/character. `puter.ai.listModels()` reports models; the per-model rates are served by the API (`GET /metering/allCosts`) rather than printed here, because a single number would be wrong for every model. - **KV and storage operations** — small per-operation costs; reads served from cache are charged a fraction of an uncached read. +A streamed AI response that stops before the model reports its token counts — an upstream error part-way through the response, say — is still charged, on an estimate of what it streamed. A request that produced no output is not charged at all. + +AI requests you have in flight count against the balance while they run, at the most they could cost, and are reconciled to their real cost when they finish. Several expensive completions started at once therefore see each other's spend rather than each being told the whole balance is available — the later ones get `402 insufficient_funds` if the balance can't cover them all. + ## Rate limits Every limit is a rolling window, keyed per user. Where three numbers are shown they are **paid / free / anonymous** — "paid" is any subscription tier.