diff --git a/src/backend/drivers/ai-chat/ChatCompletionDriver.ts b/src/backend/drivers/ai-chat/ChatCompletionDriver.ts index c7015567a..45e9a070d 100644 --- a/src/backend/drivers/ai-chat/ChatCompletionDriver.ts +++ b/src/backend/drivers/ai-chat/ChatCompletionDriver.ts @@ -36,6 +36,7 @@ import { requestAiCosts, withAiCostFactor, } from '../util/aiCostFactor.js'; +import { CompletionAccountingError } from './utils/meterChatUsage.js'; import { AI_CONCURRENT, AI_RATE_LIMIT } from '../util/aiLimits.js'; import { isCreditExhaustion as isUpstreamCreditExhaustion, @@ -909,9 +910,13 @@ export class ChatCompletionDriver extends PuterDriver { // mark the route unhealthy and walk the fallback chain. abortSignal?.throwIfAborted(); - // A withheld completion was still a completion — charged, final, - // not a route failure worth another (billed) attempt elsewhere. - if (isModerationRefusal(e)) throw e; + // Completed inference is terminal even when moderation withholds it + // or accounting fails. Retrying would buy another completion. + if ( + e instanceof CompletionAccountingError || + isModerationRefusal(e) + ) + throw e; // Our own request validation (policy rejection, bad tool shape) // is never a route problem — no other provider would accept it // either, so there is nothing a fallback would fix. An empty @@ -978,7 +983,11 @@ export class ChatCompletionDriver extends PuterDriver { await hold.release(); hold = NO_CREDIT_HOLD; abortSignal?.throwIfAborted(); - if (isModerationRefusal(fbErr)) throw fbErr; + if ( + fbErr instanceof CompletionAccountingError || + isModerationRefusal(fbErr) + ) + throw fbErr; const attempt = recordFailure(fallback, fbErr); // A fallback's own 4xx is a translation gap between // vendors, not the caller's fault — stop climbing the @@ -1153,12 +1162,14 @@ export class ChatCompletionDriver extends PuterDriver { // Fire-and-forget — the stream writes happen async while the // response is being piped to the client. (async () => { + let accountingFailed = false; try { await init({ chatStream }); if (!passthrough.writableEnded) { throw new Error('Stream ended before completion'); } } catch (e) { + accountingFailed = e instanceof CompletionAccountingError; if (!chatStream.detached) { const attempt = toAttempt(model.id, model.provider!, e); chatStream.write( @@ -1167,10 +1178,14 @@ export class ChatCompletionDriver extends PuterDriver { message: sanitizeUpstreamMessage( e instanceof Error ? e.message : String(e), ), - ...(attempt.upstreamType - ? { code: attempt.upstreamType } - : {}), - status: streamErrorStatus(attempt), + ...(accountingFailed + ? { code: 'completion_accounting_failed' } + : attempt.upstreamType + ? { code: attempt.upstreamType } + : {}), + status: accountingFailed + ? 502 + : streamErrorStatus(attempt), })}\n`, ); } @@ -1181,9 +1196,14 @@ export class ChatCompletionDriver extends PuterDriver { abortSignal?.removeEventListener('abort', detachStream); // Providers report usage the moment they meter it (see // `AIChatStream.reportUsage`); a stream that never got - // there was never charged for. + // there was never charged for. Explicit accounting failures must + // stay uncharged; only interrupted work uses the estimate. try { - if (!blocked && !chatStream.reportedUsage) { + if ( + !blocked && + !accountingFailed && + !chatStream.reportedUsage + ) { this.#meterUnreportedStream({ actor, chatStream, diff --git a/src/backend/drivers/ai-chat/providers/moonshot/MoonshotProvider.test.ts b/src/backend/drivers/ai-chat/providers/moonshot/MoonshotProvider.test.ts index 725ecd778..161d87196 100644 --- a/src/backend/drivers/ai-chat/providers/moonshot/MoonshotProvider.test.ts +++ b/src/backend/drivers/ai-chat/providers/moonshot/MoonshotProvider.test.ts @@ -33,7 +33,7 @@ * endpoint. */ -import { Writable } from 'node:stream'; +import { Readable, Writable } from 'node:stream'; import { afterAll, afterEach, @@ -47,13 +47,20 @@ import { } from 'vitest'; import { ChatCompletionDriver } from '../../ChatCompletionDriver.js'; -import { clearUnhealthyRoutes } from '../../utils/providerHealth.js'; +import { + clearUnhealthyRoutes, + isRouteUnhealthy, +} from '../../utils/providerHealth.js'; import { SYSTEM_ACTOR } from '../../../../core/actor.js'; import type { MeteringService } from '../../../../services/metering/MeteringService.js'; import { PuterServer } from '../../../../server.js'; import { setupTestServer } from '../../../../testUtil.js'; -import { withTestActor } from '../../../integrationTestUtil.js'; +import { + readStreamResult, + withTestActor, +} from '../../../integrationTestUtil.js'; import { AIChatStream } from '../../utils/Streaming.js'; +import { FakeChatProvider } from '../FakeChatProvider.js'; import { MOONSHOT_MODELS } from './models.js'; import { MoonshotProvider } from './MoonshotProvider.js'; import { modelLookupNames } from '../../utils/modelRouting.js'; @@ -1318,6 +1325,107 @@ describe('Kimi K3.1 Pro', () => { }, ); + it.each([ + [false, null, null, false], + [false, '10', '41', false], + [true, null, null, false], + [true, '10', '41', false], + [false, null, null, true], + [false, '10', '41', true], + ] as const)( + 'terminates accounting failures (stream=%s, TTL=%s/%s, fallback=%s) without retry or estimated charge', + async (stream, oneHour, fiveMinute, startWithFake) => { + const model = MOONSHOT_MODELS.find((m) => m.id === models[0])!; + // A real alternative route ensures a retry would be observable. + vi.spyOn(FakeChatProvider.prototype, 'models').mockResolvedValue([ + model, + ]); + const fallback = vi.spyOn(FakeChatProvider.prototype, 'complete'); + if (startWithFake) + fallback.mockRejectedValueOnce(new Error('Connection failed')); + const d = new ChatCompletionDriver( + { + providers: { + moonshot: { apiKey: 'test-key' }, + ollama: { enabled: false }, + }, + } as never, + server.clients, + server.stores, + server.services, + ); + await d.onServerStart(); + const actor = { + user: { + uuid: `kimi-accounting-${Math.random()}`, + username: 'kimi-accounting', + email: 'kimi@test.com', + }, + } as never; + const take = vi.spyOn(server.stores.creditHold, 'take'); + const release = vi.spyOn(server.stores.creditHold, 'release'); + mockResponse( + stream + ? asAsyncIterable([ + { choices: [{ delta: { content: 'answer' } }] }, + { choices: [], usage: structuredClone(usage) }, + ]) + : response(), + oneHour, + fiveMinute, + ); + const completion = withTestActor( + () => + d.complete({ + model: models[0]!, + provider: startWithFake ? 'fake-chat' : 'moonshot', + messages: [{ role: 'user', content: 'hi' }], + max_tokens: 10, + stream, + }), + actor, + ); + if (stream) { + const result = (await completion) as unknown as { + stream: Readable; + }; + const events = (await readStreamResult(result)) + .toString() + .trim() + .split('\n') + .map((line) => JSON.parse(line)); + expect(events).toContainEqual( + expect.objectContaining({ + type: 'error', + code: 'completion_accounting_failed', + status: 502, + message: expect.stringContaining( + 'No charge was recorded', + ), + }), + ); + } else { + await expect(completion).rejects.toMatchObject({ + statusCode: 502, + legacyCode: 'completion_accounting_failed', + message: expect.stringContaining('No charge was recorded'), + }); + } + await vi.waitFor(() => + expect(release).toHaveBeenCalledTimes(startWithFake ? 2 : 1), + ); + expect(take).toHaveBeenCalledTimes(startWithFake ? 2 : 1); + expect(take.mock.calls[0]![1]).toBeGreaterThan(0); + expect( + await server.services.metering.getOutstandingHolds(actor), + ).toBe(0); + expect(createMock).toHaveBeenCalledTimes(1); + expect(fallback).toHaveBeenCalledTimes(startWithFake ? 1 : 0); + expect(recordSpy).not.toHaveBeenCalled(); + expect(isRouteUnhealthy('moonshot', models[0]!)).toBe(false); + }, + ); + it('derives one-hour writes from the five-minute header and total writes', async () => { const { provider } = makeProvider(); mockResponse(response(), null, '40'); diff --git a/src/backend/drivers/ai-chat/providers/moonshot/MoonshotProvider.ts b/src/backend/drivers/ai-chat/providers/moonshot/MoonshotProvider.ts index e313c3ba9..ee9957420 100644 --- a/src/backend/drivers/ai-chat/providers/moonshot/MoonshotProvider.ts +++ b/src/backend/drivers/ai-chat/providers/moonshot/MoonshotProvider.ts @@ -29,6 +29,7 @@ import type { ICompleteArguments, PuterMessage, } from '../../types.js'; +import { CompletionAccountingError } from '../../utils/meterChatUsage.js'; import { openAICompatParams } from '../../utils/openaiParams.js'; import type { MeteringService } from '../../../../services/metering/MeteringService.js'; import { sdkClientOptions } from '../../utils/sdkClient.js'; @@ -186,13 +187,7 @@ export class MoonshotProvider extends OpenAICompatProvider { fiveMinuteHeader === '' || oneHourHeader === '' ) { - throw new HttpError( - 502, - 'Moonshot cache-write TTL usage is unavailable or inconsistent', - { - legacyCode: 'bad_response', - }, - ); + throw new CompletionAccountingError(); } return { usage: { diff --git a/src/backend/drivers/ai-chat/utils/meterChatUsage.ts b/src/backend/drivers/ai-chat/utils/meterChatUsage.ts index 38d5e253d..06962bffe 100644 --- a/src/backend/drivers/ai-chat/utils/meterChatUsage.ts +++ b/src/backend/drivers/ai-chat/utils/meterChatUsage.ts @@ -17,11 +17,29 @@ * along with this program. If not, see . */ +import { HttpError } from '../../../core/http/HttpError.js'; import type { Actor } from '../../../core/actor.js'; import type { MeteringService } from '../../../services/metering/MeteringService.js'; import type { IChatModel } from '../types.js'; import { buildCostsOverride } from './pricing.js'; +/** + * Upstream inference finished, but no reliable charge can be recorded. + * Terminal: retrying would buy another completion, and estimating would hide + * the accounting failure. Only raise this before recording any charge. + */ +export class CompletionAccountingError extends HttpError { + constructor() { + super( + 502, + 'The completion finished, but its usage could not be verified. No charge was recorded.', + { + legacyCode: 'completion_accounting_failed', + }, + ); + } +} + /** Microcents per US dollar. */ const MICROCENTS_PER_USD = 100_000_000; diff --git a/src/docs/src/AI/chat.md b/src/docs/src/AI/chat.md index adc86f7e4..fed36e2be 100755 --- a/src/docs/src/AI/chat.md +++ b/src/docs/src/AI/chat.md @@ -51,7 +51,7 @@ Select `kimi-k3.1-pro` or `kimi-k3.1-pro-highspeed` with `provider: 'moonshotai' The context window is 1,048,576 tokens. Use `max_tokens` to set an output cap up to 1,048,576 tokens, subject to the remaining context window and available credits. Preserve the complete returned assistant message, including its `reasoning` and tool calls, when continuing a conversation. -Usage separates uncached input, cached reads, cache writes by TTL, and output. Cache writes are billed once when a prefix is stored. The upstream cache uses a five-minute TTL by default. Requests reserve credits at the one-hour cache-write rate because an existing prefix can retain that TTL. Actual charges use the returned usage. If writes are reported without a usable TTL breakdown, the request fails rather than estimating their charge. For current rates, see the [AI models list](https://developer.puter.com/ai/models/). +Usage separates uncached input, cached reads, cache writes by TTL, and output. Cache writes are billed once when a prefix is stored. The upstream cache uses a five-minute TTL by default. Requests reserve credits at the one-hour cache-write rate because an existing prefix can retain that TTL. Actual charges use the returned usage. If writes are reported without a usable TTL breakdown, the request ends with `completion_accounting_failed`: no charge is recorded, the credit hold is released, and Puter does not retry another provider. A stream may already have delivered content before this terminal error. For current rates, see the [AI models list](https://developer.puter.com/ai/models/). #### `testMode` (Boolean) (Optional)