diff --git a/src/backend/drivers/ai-chat/providers/azure/AzureChatProvider.ts b/src/backend/drivers/ai-chat/providers/azure/AzureChatProvider.ts index e12d4e3fd..bec9db009 100644 --- a/src/backend/drivers/ai-chat/providers/azure/AzureChatProvider.ts +++ b/src/backend/drivers/ai-chat/providers/azure/AzureChatProvider.ts @@ -31,6 +31,7 @@ import { wantsCompaction, } from '../../utils/compaction.js'; import * as OpenAiUtil from '../../utils/OpenAIUtil.js'; +import { buildCostsOverride } from '../../utils/pricing.js'; import { processPuterPathUploads } from '../openai/fileUpload.js'; import { AZURE_MODELS } from './models.js'; @@ -256,10 +257,9 @@ export class AzureChatProvider implements IChatProvider { cached_tokens: cachedTokens, }; - const costsOverrideFromModel = Object.fromEntries( - Object.entries(trackedUsage).map(([k, v]) => { - return [k, v * modelUsed.costs[k]]; - }), + const costsOverrideFromModel = buildCostsOverride( + trackedUsage, + modelUsed, ); this.#meteringService.utilRecordUsageObject( diff --git a/src/backend/drivers/ai-chat/providers/azure/AzureResponsesProvider.ts b/src/backend/drivers/ai-chat/providers/azure/AzureResponsesProvider.ts index 130a4b40e..7a25b6db1 100644 --- a/src/backend/drivers/ai-chat/providers/azure/AzureResponsesProvider.ts +++ b/src/backend/drivers/ai-chat/providers/azure/AzureResponsesProvider.ts @@ -27,6 +27,7 @@ import type { S3ObjectStore } from '../../../../stores/fs/S3ObjectStore.js'; import type { IChatProvider, ICompleteArguments } from '../../types.js'; import { toOpenAiContextManagement } from '../../utils/compaction.js'; import * as OpenAiUtil from '../../utils/OpenAIUtil.js'; +import { buildCostsOverride } from '../../utils/pricing.js'; import { processPuterPathUploads } from '../openai/fileUpload.js'; import { AZURE_MODELS } from './models.js'; import { HttpError } from '@heyputer/backend/src/core/http/HttpError.js'; @@ -45,9 +46,7 @@ import { HttpError } from '@heyputer/backend/src/core/http/HttpError.js'; * NOT Azure's — Azure is subsidised for us. */ export class AzureResponsesProvider implements IChatProvider { - /** - * @type {import('openai').OpenAI} - */ + /** @type {import('openai').OpenAI} */ #openAi: OpenAI; #defaultModel = 'gpt-5-codex'; @@ -75,7 +74,8 @@ export class AzureResponsesProvider implements IChatProvider { /** * Returns an array of available AI models with their pricing information. - * Each model object includes an ID and cost details (currency, tokens, input/output rates). + * Each model object includes an ID and cost details (currency, tokens, + * input/output rates). */ models(extra_params?: { no_restrictions?: boolean }) { if (extra_params?.no_restrictions) { @@ -259,10 +259,9 @@ export class AzureResponsesProvider implements IChatProvider { (usage as any).input_tokens_details?.cached_tokens ?? 0, }; - const costsOverrideFromModel = Object.fromEntries( - Object.entries(trackedUsage).map(([k, v]) => { - return [k, v * modelUsed.costs[k]]; - }), + const costsOverrideFromModel = buildCostsOverride( + trackedUsage, + modelUsed, ); this.#meteringService.utilRecordUsageObject( diff --git a/src/backend/drivers/ai-chat/providers/gemini/GeminiChatProvider.test.ts b/src/backend/drivers/ai-chat/providers/gemini/GeminiChatProvider.test.ts index c43c8e123..193c410cb 100644 --- a/src/backend/drivers/ai-chat/providers/gemini/GeminiChatProvider.test.ts +++ b/src/backend/drivers/ai-chat/providers/gemini/GeminiChatProvider.test.ts @@ -399,6 +399,49 @@ describe('GeminiChatProvider.complete non-stream output', () => { }); }); + it('bills cached tokens at the input rate when the model prices no cache read', async () => { + // gemini-2.0-flash-lite's catalogue entry has no cached_tokens rate. + // Cached tokens are subtracted out of prompt_tokens, so pricing them + // at zero bills them nowhere. + const lite = GEMINI_MODELS.find( + (m) => m.id === 'gemini-2.0-flash-lite', + )!; + expect(lite.costs.cached_tokens).toBeUndefined(); + + const { provider } = makeProvider(); + createMock.mockResolvedValueOnce({ + choices: [ + { + message: { content: 'cached', role: 'assistant' }, + finish_reason: 'stop', + }, + ], + usage: { + prompt_tokens: 3000, + completion_tokens: 40, + prompt_tokens_details: { cached_tokens: 2900 }, + }, + }); + + await withTestActor(() => + provider.complete({ + model: 'gemini-2.0-flash-lite', + messages: [{ role: 'user', content: 'hi' }], + }), + ); + + const [, , , overrides] = recordSpy.mock.calls[0]!; + const inputRate = Number(lite.costs.prompt_tokens); + expect(overrides).toMatchObject({ + prompt_tokens: (3000 - 2900) * inputRate, + completion_tokens: 40 * Number(lite.costs.completion_tokens), + cached_tokens: 2900 * inputRate, + }); + expect( + (overrides as Record).cached_tokens, + ).toBeGreaterThan(0); + }); + it('zeroes cached_tokens when prompt_tokens_details is missing', async () => { const { provider } = makeProvider(); createMock.mockResolvedValueOnce({ @@ -595,6 +638,46 @@ describe('GeminiChatProvider.complete grounding request metering', () => { ); }); + it('charges every grounding-capable model the per-generation request fee', async () => { + // Flash-Lite serves grounded requests like the rest of its + // generation; without its own rate the fee fell through to the input + // token rate, which is several orders of magnitude below list. + const lite = GEMINI_MODELS.find( + (m) => m.id === 'gemini-2.0-flash-lite', + )!; + expect(lite.costs.grounding_requests).toBe(3_500_000); + + const { provider } = makeProvider(); + createMock.mockResolvedValueOnce({ + choices: [ + { + message: { + content: 'result', + role: 'assistant', + extra_content: { + grounding_metadata: { web_search_queries: ['foo'] }, + }, + }, + finish_reason: 'stop', + }, + ], + usage: { prompt_tokens: 10, completion_tokens: 5 }, + }); + + await withTestActor(() => + provider.complete({ + model: 'gemini-2.0-flash-lite', + messages: [{ role: 'user', content: 'search for foo' }], + }), + ); + + const [usage, , , overrides] = recordSpy.mock.calls[0]!; + expect(usage.grounding_requests).toBe(1); + expect(overrides!.grounding_requests).toBe( + Number(lite.costs.grounding_requests), + ); + }); + it('does not charge a grounding request when no grounding_metadata is present', async () => { const { provider } = makeProvider(); createMock.mockResolvedValueOnce({ diff --git a/src/backend/drivers/ai-chat/providers/gemini/GeminiChatProvider.ts b/src/backend/drivers/ai-chat/providers/gemini/GeminiChatProvider.ts index 623be3bf8..069c66c09 100644 --- a/src/backend/drivers/ai-chat/providers/gemini/GeminiChatProvider.ts +++ b/src/backend/drivers/ai-chat/providers/gemini/GeminiChatProvider.ts @@ -28,6 +28,7 @@ import { handle_completion_output, process_input_messages, } from '../../utils/OpenAIUtil.js'; +import { buildCostsOverride } from '../../utils/pricing.js'; import { GEMINI_MODELS } from './models.js'; export class GeminiChatProvider implements IChatProvider { @@ -142,10 +143,9 @@ export class GeminiChatProvider implements IChatProvider { : 0, }; - const costsOverrideFromModel = Object.fromEntries( - Object.entries(trackedUsage).map(([k, v]) => { - return [k, v * (modelUsed.costs[k] ?? 0)]; - }), + const costsOverrideFromModel = buildCostsOverride( + trackedUsage, + modelUsed, ); this.meteringService.utilRecordUsageObject( trackedUsage, diff --git a/src/backend/drivers/ai-chat/providers/gemini/models.ts b/src/backend/drivers/ai-chat/providers/gemini/models.ts index 73f4d127c..515921898 100644 --- a/src/backend/drivers/ai-chat/providers/gemini/models.ts +++ b/src/backend/drivers/ai-chat/providers/gemini/models.ts @@ -97,6 +97,8 @@ export const GEMINI_MODELS: IChatModel[] = [ tokens: 1_000_000, prompt_tokens: 8, completion_tokens: 30, + // Gemini 2.x grounding is $35 / 1,000 requests + grounding_requests: 3_500_000, }, max_tokens: 8192, }, diff --git a/src/backend/drivers/ai-chat/providers/openai/OpenAiChatCompletionsProvider.test.ts b/src/backend/drivers/ai-chat/providers/openai/OpenAiChatCompletionsProvider.test.ts index b37d2fe69..158becad0 100644 --- a/src/backend/drivers/ai-chat/providers/openai/OpenAiChatCompletionsProvider.test.ts +++ b/src/backend/drivers/ai-chat/providers/openai/OpenAiChatCompletionsProvider.test.ts @@ -407,6 +407,47 @@ describe('OpenAiChatProvider.complete non-stream output', () => { }); }); + it('bills cached tokens at the input rate when the model prices no cache read', async () => { + // o4-mini's catalogue entry has no cached_tokens rate. Cached tokens + // are subtracted out of prompt_tokens, so pricing them at zero bills + // them nowhere — the whole cached portion of the request goes free. + const o4Mini = OPEN_AI_MODELS.find((m) => m.id === 'o4-mini')!; + expect(o4Mini.costs.cached_tokens).toBeUndefined(); + + const { provider } = makeProvider(); + createMock.mockResolvedValueOnce({ + choices: [ + { + message: { content: 'cached', role: 'assistant' }, + finish_reason: 'stop', + }, + ], + usage: { + prompt_tokens: 2989, + completion_tokens: 12, + prompt_tokens_details: { cached_tokens: 2816 }, + }, + }); + + await withTestActor(() => + provider.complete({ + model: 'o4-mini', + messages: [{ role: 'user', content: 'hi' }], + }), + ); + + const [, , , overrides] = recordSpy.mock.calls[0]!; + const inputRate = Number(o4Mini.costs.prompt_tokens); + expect(overrides).toEqual({ + prompt_tokens: (2989 - 2816) * inputRate, + completion_tokens: 12 * Number(o4Mini.costs.completion_tokens), + cached_tokens: 2816 * inputRate, + }); + expect( + (overrides as Record).cached_tokens, + ).toBeGreaterThan(0); + }); + it('zeroes cached_tokens when prompt_tokens_details is missing', async () => { const { provider } = makeProvider(); createMock.mockResolvedValueOnce({ diff --git a/src/backend/drivers/ai-chat/providers/openai/OpenAiChatCompletionsProvider.ts b/src/backend/drivers/ai-chat/providers/openai/OpenAiChatCompletionsProvider.ts index 1ccc2692e..8f97012f4 100644 --- a/src/backend/drivers/ai-chat/providers/openai/OpenAiChatCompletionsProvider.ts +++ b/src/backend/drivers/ai-chat/providers/openai/OpenAiChatCompletionsProvider.ts @@ -31,6 +31,7 @@ import { wantsCompaction, } from '../../utils/compaction.js'; import * as OpenAiUtil from '../../utils/OpenAIUtil.js'; +import { buildCostsOverride } from '../../utils/pricing.js'; import { processPuterPathUploads } from './fileUpload.js'; import { OPEN_AI_MODELS } from './models.js'; import type { OpenAiResponsesChatProvider } from './OpenAiChatResponsesProvider.js'; @@ -228,10 +229,9 @@ export class OpenAiChatProvider implements IChatProvider { usage.prompt_tokens_details?.cached_tokens ?? 0, }; - const costsOverrideFromModel = Object.fromEntries( - Object.entries(trackedUsage).map(([k, v]) => { - return [k, v * modelUsed.costs[k]]; - }), + const costsOverrideFromModel = buildCostsOverride( + trackedUsage, + modelUsed, ); this.#meteringService.utilRecordUsageObject( diff --git a/src/backend/drivers/ai-chat/providers/openai/OpenAiChatResponsesProvider.test.ts b/src/backend/drivers/ai-chat/providers/openai/OpenAiChatResponsesProvider.test.ts index 4f0904add..7105833ee 100644 --- a/src/backend/drivers/ai-chat/providers/openai/OpenAiChatResponsesProvider.test.ts +++ b/src/backend/drivers/ai-chat/providers/openai/OpenAiChatResponsesProvider.test.ts @@ -445,6 +445,43 @@ describe('OpenAiResponsesChatProvider.complete non-stream output', () => { }); }); + it('bills cached tokens at the input rate when the model prices no cache read', async () => { + // gpt-5.4-pro is responses-API-only and its catalogue entry has no + // cached_tokens rate. Cached tokens are subtracted out of the input + // count, so pricing them at zero bills them nowhere. + const pro = OPEN_AI_MODELS.find((m) => m.id === 'gpt-5.4-pro')!; + expect(pro.costs.cached_tokens).toBeUndefined(); + + const { provider } = makeProvider(); + responsesCreateMock.mockResolvedValueOnce({ + output: [{ role: 'assistant' }], + output_text: 'cached', + usage: { + input_tokens: 7761, + output_tokens: 20, + input_tokens_details: { cached_tokens: 7680 }, + }, + }); + + await withTestActor(() => + provider.complete({ + model: 'gpt-5.4-pro', + messages: [{ role: 'user', content: 'hi' }], + }), + ); + + const [, , , overrides] = recordSpy.mock.calls[0]!; + const inputRate = Number(pro.costs.prompt_tokens); + expect(overrides).toEqual({ + prompt_tokens: (7761 - 7680) * inputRate, + completion_tokens: 20 * Number(pro.costs.completion_tokens), + cached_tokens: 7680 * inputRate, + }); + expect( + (overrides as Record).cached_tokens, + ).toBeGreaterThan(0); + }); + it('shapes function_call output items into OpenAI tool_calls on the response', async () => { const { provider } = makeProvider(); responsesCreateMock.mockResolvedValueOnce({ diff --git a/src/backend/drivers/ai-chat/providers/openai/OpenAiChatResponsesProvider.ts b/src/backend/drivers/ai-chat/providers/openai/OpenAiChatResponsesProvider.ts index 94605b385..f31973e26 100644 --- a/src/backend/drivers/ai-chat/providers/openai/OpenAiChatResponsesProvider.ts +++ b/src/backend/drivers/ai-chat/providers/openai/OpenAiChatResponsesProvider.ts @@ -27,21 +27,21 @@ import type { S3ObjectStore } from '../../../../stores/fs/S3ObjectStore.js'; import type { IChatProvider, ICompleteArguments } from '../../types.js'; import { toOpenAiContextManagement } from '../../utils/compaction.js'; import * as OpenAiUtil from '../../utils/OpenAIUtil.js'; +import { buildCostsOverride } from '../../utils/pricing.js'; import { processPuterPathUploads } from './fileUpload.js'; import { OPEN_AI_MODELS } from './models.js'; import { HttpError } from '@heyputer/backend/src/core/http/HttpError.js'; /** - * OpenAICompletionService class provides an interface to OpenAI's chat completion API. - * Extends BaseService to handle chat completions, message moderation, token counting, - * and streaming responses. Implements the puter-chat-completion interface and manages - * OpenAI API interactions with support for multiple models including GPT-4 variants. - * Handles usage tracking, spending records, and content moderation. + * OpenAICompletionService class provides an interface to OpenAI's chat + * completion API. Extends BaseService to handle chat completions, message + * moderation, token counting, and streaming responses. Implements the + * puter-chat-completion interface and manages OpenAI API interactions with + * support for multiple models including GPT-4 variants. Handles usage tracking, + * spending records, and content moderation. */ export class OpenAiResponsesChatProvider implements IChatProvider { - /** - * @type {import('openai').OpenAI} - */ + /** @type {import('openai').OpenAI} */ #openAi: OpenAI; #defaultModel = 'gpt-5-nano'; @@ -66,7 +66,8 @@ export class OpenAiResponsesChatProvider implements IChatProvider { /** * Returns an array of available AI models with their pricing information. - * Each model object includes an ID and cost details (currency, tokens, input/output rates). + * Each model object includes an ID and cost details (currency, tokens, + * input/output rates). */ models(extra_params) { if (extra_params?.no_restrictions) { @@ -252,10 +253,9 @@ export class OpenAiResponsesChatProvider implements IChatProvider { (usage as any).input_tokens_details?.cached_tokens ?? 0, }; - const costsOverrideFromModel = Object.fromEntries( - Object.entries(trackedUsage).map(([k, v]) => { - return [k, v * modelUsed.costs[k]]; - }), + const costsOverrideFromModel = buildCostsOverride( + trackedUsage, + modelUsed, ); this.#meteringService.utilRecordUsageObject( diff --git a/src/backend/drivers/ai-chat/utils/pricing.test.ts b/src/backend/drivers/ai-chat/utils/pricing.test.ts new file mode 100644 index 000000000..4f1f01807 --- /dev/null +++ b/src/backend/drivers/ai-chat/utils/pricing.test.ts @@ -0,0 +1,142 @@ +/* + * 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 type { IChatModel } from '../types.js'; +import { buildCostsOverride, usdPerMToken } from './pricing.js'; + +const model = (costs: Record): IChatModel => + ({ + id: 'test-model', + costs_currency: 'usd-cents', + input_cost_key: 'prompt_tokens', + output_cost_key: 'completion_tokens', + costs, + max_tokens: 1024, + }) as IChatModel; + +describe('usdPerMToken', () => { + it('always emits a cached_tokens row, defaulting to zero', () => { + expect(usdPerMToken(1, 2)).toEqual({ + tokens: 1_000_000, + prompt_tokens: 100, + completion_tokens: 200, + cached_tokens: 0, + }); + expect(usdPerMToken(1, 2, 0.5).cached_tokens).toBe(50); + }); +}); + +describe('buildCostsOverride', () => { + it('multiplies each usage key by its own declared rate', () => { + const overrides = buildCostsOverride( + { prompt_tokens: 90, completion_tokens: 50, cached_tokens: 10 }, + model({ prompt_tokens: 110, completion_tokens: 440, cached_tokens: 55 }), + ); + + expect(overrides).toEqual({ + prompt_tokens: 90 * 110, + completion_tokens: 50 * 440, + cached_tokens: 10 * 55, + }); + }); + + it('prices an undeclared key at the input rate rather than giving it away', () => { + // A model whose catalogue entry omits cached_tokens: the cached count + // has already been subtracted out of prompt_tokens, so pricing it at + // zero bills it nowhere at all. + const overrides = buildCostsOverride( + { prompt_tokens: 173, completion_tokens: 12, cached_tokens: 2816 }, + model({ prompt_tokens: 110, completion_tokens: 440 }), + ); + + expect(overrides.cached_tokens).toBe(2816 * 110); + expect(overrides.cached_tokens).toBeGreaterThan(0); + }); + + it('prices an undeclared output-denominated key at the output rate', () => { + const overrides = buildCostsOverride( + { prompt_tokens: 10, completion_tokens: 20, thinking_tokens: 30 }, + model({ prompt_tokens: 8, completion_tokens: 30 }), + ); + + expect(overrides.thinking_tokens).toBe(30 * 30); + }); + + it('honours an explicitly declared zero rate', () => { + // An explicit zero is a pricing decision — usually "already billed + // inside another row" — and must not be overridden by the fallback. + const overrides = buildCostsOverride( + { prompt_tokens: 10, cached_tokens: 99 }, + model({ prompt_tokens: 8, completion_tokens: 30, cached_tokens: 0 }), + ); + + expect(overrides.cached_tokens).toBe(0); + }); + + it('falls back to zero only when the model prices nothing at all', () => { + const overrides = buildCostsOverride( + { prompt_tokens: 10, cached_tokens: 5 }, + model({}), + ); + + expect(overrides).toEqual({ prompt_tokens: 0, cached_tokens: 0 }); + }); + + it('skips the tokens scale descriptor', () => { + const overrides = buildCostsOverride( + { prompt_tokens: 10, tokens: 1_000_000 }, + model({ prompt_tokens: 8, completion_tokens: 30 }), + ); + + expect(overrides).toEqual({ prompt_tokens: 80 }); + }); + + it('never emits a non-finite value for a model with a broken cost table', () => { + const overrides = buildCostsOverride( + { prompt_tokens: 10, completion_tokens: 20, cached_tokens: 30 }, + model({ + prompt_tokens: Number.NaN, + completion_tokens: Number.POSITIVE_INFINITY, + }), + ); + + for (const value of Object.values(overrides)) { + expect(Number.isFinite(value)).toBe(true); + } + }); + + it('resolves rates through the model default keys when none are declared', () => { + const overrides = buildCostsOverride( + { input_tokens: 10, output_tokens: 20, cached_tokens: 5 }, + { + id: 'defaults', + costs_currency: 'usd-cents', + costs: { input_tokens: 3, output_tokens: 9 }, + max_tokens: 1024, + } as IChatModel, + ); + + expect(overrides).toEqual({ + input_tokens: 30, + output_tokens: 180, + cached_tokens: 15, + }); + }); +}); diff --git a/src/backend/drivers/ai-chat/utils/pricing.ts b/src/backend/drivers/ai-chat/utils/pricing.ts index 4f9e2e8dc..1d9f3fe3a 100644 --- a/src/backend/drivers/ai-chat/utils/pricing.ts +++ b/src/backend/drivers/ai-chat/utils/pricing.ts @@ -17,7 +17,7 @@ * along with this program. If not, see . */ -import type { ModelCost } from '../types.js'; +import type { IChatModel, ModelCost } from '../types.js'; const CENTS_PER_USD = 100; const MTOK = 1_000_000; @@ -38,3 +38,52 @@ export const usdPerMToken = ( completion_tokens: outputUsd * CENTS_PER_USD, cached_tokens: cachedReadUsd * CENTS_PER_USD, }); + +const isRate = (value: unknown): value is number => + typeof value === 'number' && Number.isFinite(value); + +/** + * Prices a tracked-usage object against a model's cost table. + * + * A usage key the model doesn't price falls back to the model's output rate + * when it is output-denominated and its input rate otherwise — never to zero. + * Pricing an unpriced key at zero gives the unit away, and a provider that + * subtracts cached tokens out of the prompt count has already removed them from + * the key that would otherwise have caught them. The fallback mirrors the rate + * resolution behind the reported `usd_cents`, so the ledger and the figure + * quoted to the caller agree. + */ +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 costs = model.costs ?? {}; + const inputRate = isRate(costs[inputKey]) ? costs[inputKey] : undefined; + const outputRate = isRate(costs[outputKey]) ? costs[outputKey] : undefined; + + const isOutputKey = (key: string) => + key === outputKey || + key === 'output_tokens' || + key === 'completion_tokens' || + key === 'thinking_tokens'; + + const overrides: Record = {}; + for (const [key, amount] of Object.entries(trackedUsage)) { + // `tokens` is a scale descriptor ("costs expressed per N tokens"), + // not a per-unit rate. + if (key === 'tokens') continue; + + const rate = isRate(costs[key]) + ? costs[key] + : ((isOutputKey(key) ? outputRate : inputRate) ?? 0); + + overrides[key] = amount * rate; + } + + return overrides; +}; diff --git a/src/backend/drivers/ai-image/ImageGenerationDriver.test.ts b/src/backend/drivers/ai-image/ImageGenerationDriver.test.ts index 580313ab0..bb03971e6 100644 --- a/src/backend/drivers/ai-image/ImageGenerationDriver.test.ts +++ b/src/backend/drivers/ai-image/ImageGenerationDriver.test.ts @@ -128,6 +128,13 @@ vi.mock('replicate', () => { return { default: Replicate }; }); +const { secureFetchMock } = vi.hoisted(() => ({ secureFetchMock: vi.fn() })); + +vi.mock('../../util/secureHttp.js', async (importOriginal) => ({ + ...(await importOriginal()), + secureFetch: secureFetchMock, +})); + // ── Test harness ──────────────────────────────────────────────────── let server: PuterServer; @@ -163,6 +170,7 @@ beforeEach(() => { googleAIGenerateImagesMock.mockReset(); togetherImagesGenerateMock.mockReset(); replicateRunMock.mockReset(); + secureFetchMock.mockReset(); fetchSpy = vi.spyOn(globalThis, 'fetch') as MockInstance; eventEmitSpy = vi.spyOn(server.clients.event, 'emit') as MockInstance< (...args: unknown[]) => unknown @@ -514,7 +522,7 @@ describe('ImageGenerationDriver.generate puter_output_path', () => { openaiImagesGenerateMock.mockResolvedValueOnce({ data: [{ url: 'https://oai/img.png' }], }); - fetchSpy.mockResolvedValueOnce( + secureFetchMock.mockResolvedValueOnce( new Response(Buffer.from('fake-png'), { status: 200, headers: { 'content-type': 'image/png' }, @@ -546,7 +554,7 @@ describe('ImageGenerationDriver.generate puter_output_path', () => { openaiImagesGenerateMock.mockResolvedValueOnce({ data: [{ url: 'https://oai/img.png' }], }); - fetchSpy.mockResolvedValueOnce( + secureFetchMock.mockResolvedValueOnce( new Response(Buffer.from('fake-png'), { status: 200, headers: { 'content-type': 'image/png' }, @@ -577,6 +585,13 @@ describe('ImageGenerationDriver.generate puter_output_path', () => { expect(meta.path).toBe('/testuser/photos/out.png'); expect(meta.contentType).toBe('image/png'); expect(meta.overwrite).toBe(true); + + // The result URL is downloaded through the SSRF-guarded fetch, not + // the unguarded global one — its body lands in the user's FS. + expect(secureFetchMock).toHaveBeenCalledWith('https://oai/img.png', { + skipProxy: true, + }); + expect(fetchSpy).not.toHaveBeenCalled(); }); it('does not forward puter_output_path to the upstream provider call', async () => { @@ -589,7 +604,7 @@ describe('ImageGenerationDriver.generate puter_output_path', () => { openaiImagesGenerateMock.mockResolvedValueOnce({ data: [{ url: 'https://oai/img.png' }], }); - fetchSpy.mockResolvedValueOnce( + secureFetchMock.mockResolvedValueOnce( new Response(Buffer.from('fake-png'), { status: 200, headers: { 'content-type': 'image/png' }, diff --git a/src/backend/drivers/ai-image/ImageGenerationDriver.ts b/src/backend/drivers/ai-image/ImageGenerationDriver.ts index 3bac30081..b8ae10d20 100644 --- a/src/backend/drivers/ai-image/ImageGenerationDriver.ts +++ b/src/backend/drivers/ai-image/ImageGenerationDriver.ts @@ -25,6 +25,7 @@ import { Context } from '../../core/context.js'; import { HttpError } from '../../core/http/HttpError.js'; import type { Actor } from '../../core/actor.js'; import { PuterDriver } from '../types.js'; +import { secureFetch } from '../../util/secureHttp.js'; import { AI_CONCURRENT, AI_RATE_LIMIT } from '../util/aiLimits.js'; import { CloudflareImageProvider } from './providers/cloudflare/CloudflareImageProvider.js'; import { GeminiImageProvider } from './providers/gemini/GeminiImageProvider.js'; @@ -392,7 +393,11 @@ export class ImageGenerationDriver extends PuterDriver { header.match(/data:(.*?);/)?.[1] ?? 'application/octet-stream'; buffer = Buffer.from(result.substring(commaIdx + 1), 'base64'); } else { - const response = await fetch(result); + // Provider-minted URL, but fetched with the same SSRF guards as + // the input paths: it reaches an unauthenticated GET whose body + // lands in the user's filesystem. skipProxy because generated + // media is ours to download directly, not user input to screen. + const response = await secureFetch(result, { skipProxy: true }); if (!response.ok) { throw new HttpError( 502, diff --git a/src/backend/drivers/ai-tts/providers/elevenlabs/ElevenLabsTTSProvider.test.ts b/src/backend/drivers/ai-tts/providers/elevenlabs/ElevenLabsTTSProvider.test.ts index bc321963f..bfe553bc7 100644 --- a/src/backend/drivers/ai-tts/providers/elevenlabs/ElevenLabsTTSProvider.test.ts +++ b/src/backend/drivers/ai-tts/providers/elevenlabs/ElevenLabsTTSProvider.test.ts @@ -217,6 +217,40 @@ describe('ElevenLabsTTSProvider.synthesize argument validation', () => { ).rejects.toMatchObject({ statusCode: 400 }); expect(fetchSpy).not.toHaveBeenCalled(); }); + + it('rejects a model the cost table cannot price, before paying the vendor', async () => { + // An unpriced id resolved to a zero rate, which made the credit gate + // pass for anyone and recorded the synthesis as free — while the id + // was forwarded to the vendor and billed to us. + const provider = makeProvider(); + + await expect( + withTestActor(() => + provider.synthesize({ text: 'hello', model: 'eleven_flash_v2' }), + ), + ).rejects.toMatchObject({ + statusCode: 400, + fields: { key: 'model', got: 'eleven_flash_v2' }, + }); + + expect(fetchSpy).not.toHaveBeenCalled(); + expect(hasCreditsSpy).not.toHaveBeenCalled(); + expect(incrementUsageSpy).not.toHaveBeenCalled(); + }); + + it('accepts a priced model that the engine listing does not advertise', async () => { + const provider = makeProvider(); + fetchSpy.mockResolvedValueOnce(audioResponse()); + + await withTestActor(() => + provider.synthesize({ text: 'hi', model: 'eleven_turbo_v2' }), + ); + + expect(fetchSpy).toHaveBeenCalledTimes(1); + const [, usageType, , cost] = incrementUsageSpy.mock.calls[0]!; + expect(usageType).toBe('elevenlabs:eleven_turbo_v2:character'); + expect(cost).toBe(ELEVENLABS_TTS_COSTS['eleven_turbo_v2'] * 2); + }); }); // ── Credit gate ───────────────────────────────────────────────────── diff --git a/src/backend/drivers/ai-tts/providers/elevenlabs/ElevenLabsTTSProvider.ts b/src/backend/drivers/ai-tts/providers/elevenlabs/ElevenLabsTTSProvider.ts index 196ade439..92bc59539 100644 --- a/src/backend/drivers/ai-tts/providers/elevenlabs/ElevenLabsTTSProvider.ts +++ b/src/backend/drivers/ai-tts/providers/elevenlabs/ElevenLabsTTSProvider.ts @@ -216,12 +216,28 @@ export class ElevenLabsTTSProvider extends TTSProvider { const voiceId = voiceArg || this.defaultVoiceId; const modelId = modelArg || DEFAULT_MODEL; + + // Gate on the cost table rather than the advertised model list: an id + // we can't price is an id we can't bill for, and the vendor bills us + // for it either way. + if (!Object.hasOwn(ELEVENLABS_TTS_COSTS, modelId)) { + const expected = Object.keys(ELEVENLABS_TTS_COSTS); + throw new HttpError( + 400, + `Invalid model: ${modelId}. Expected: ${expected.join(', ')}`, + { + legacyCode: 'field_invalid', + fields: { key: 'model', expected, got: modelId }, + }, + ); + } + const desiredFormat = output_format || response_format || DEFAULT_OUTPUT_FORMAT; const actor = Context.get('actor')!; const usageKey = `elevenlabs:${modelId}:character`; - const ucentsPerChar = ELEVENLABS_TTS_COSTS[modelId] ?? 0; + const ucentsPerChar = ELEVENLABS_TTS_COSTS[modelId]; const totalCost = ucentsPerChar * text.length; const usageAllowed = await this.meteringService.hasEnoughCredits( diff --git a/src/backend/drivers/ai-video/VideoGenerationDriver.test.ts b/src/backend/drivers/ai-video/VideoGenerationDriver.test.ts index 867d45bff..3c723a08e 100644 --- a/src/backend/drivers/ai-video/VideoGenerationDriver.test.ts +++ b/src/backend/drivers/ai-video/VideoGenerationDriver.test.ts @@ -120,6 +120,13 @@ vi.mock('together-ai', () => { return { Together, default: Together }; }); +const { secureFetchMock } = vi.hoisted(() => ({ secureFetchMock: vi.fn() })); + +vi.mock('../../util/secureHttp.js', async (importOriginal) => ({ + ...(await importOriginal()), + secureFetch: secureFetchMock, +})); + // ── Test harness ──────────────────────────────────────────────────── let server: PuterServer; @@ -148,6 +155,7 @@ beforeEach(() => { geminiGenerateVideosMock.mockReset(); togetherVideosCreateMock.mockReset(); togetherVideosRetrieveMock.mockReset(); + secureFetchMock.mockReset(); hasCreditsSpy = vi.spyOn(server.services.metering, 'hasEnoughCredits'); hasCreditsSpy.mockResolvedValue(true); vi.spyOn(server.services.metering, 'getRemainingUsage').mockResolvedValue( @@ -557,6 +565,47 @@ describe('VideoGenerationDriver.generate puter_output_path', () => { ).toBe('/testuser/videos/clip.mp4'); }); + it('downloads a URL result through the SSRF-guarded fetch before writing it to FS', async () => { + const aclCheckSpy = vi.spyOn(server.services.acl, 'check'); + aclCheckSpy.mockResolvedValueOnce(true); + + const fsWriteSpy = vi.spyOn(server.services.fs, 'write'); + fsWriteSpy.mockResolvedValueOnce(undefined as never); + + togetherVideosCreateMock.mockResolvedValueOnce({ id: 'tg-job' }); + togetherVideosRetrieveMock.mockResolvedValueOnce({ + id: 'tg-job', + status: 'completed', + outputs: { video_url: 'https://together/out.mp4' }, + }); + secureFetchMock.mockResolvedValueOnce( + new Response(Buffer.from('fake-mp4'), { + status: 200, + headers: { 'content-type': 'video/mp4' }, + }), + ); + + await withTestUser(() => + driver.generate({ + prompt: 'hi', + model: 'togetherai:minimax/video-01-director', + puter_output_path: '/testuser/videos/clip.mp4', + } as never), + ); + + expect(secureFetchMock).toHaveBeenCalledWith( + 'https://together/out.mp4', + { skipProxy: true }, + ); + expect(fsWriteSpy).toHaveBeenCalledTimes(1); + const [, writeArg] = fsWriteSpy.mock.calls[0]!; + const meta = ( + writeArg as { fileMetadata: { path: string; contentType: string } } + ).fileMetadata; + expect(meta.path).toBe('/testuser/videos/clip.mp4'); + expect(meta.contentType).toBe('video/mp4'); + }); + it('writes stream result to FS and returns a new stream to caller', async () => { const aclCheckSpy = vi.spyOn(server.services.acl, 'check'); aclCheckSpy.mockResolvedValueOnce(true); diff --git a/src/backend/drivers/ai-video/VideoGenerationDriver.ts b/src/backend/drivers/ai-video/VideoGenerationDriver.ts index 967917eb1..e3d73eb0c 100644 --- a/src/backend/drivers/ai-video/VideoGenerationDriver.ts +++ b/src/backend/drivers/ai-video/VideoGenerationDriver.ts @@ -24,6 +24,7 @@ import { Context } from '../../core/context.js'; import { HttpError } from '../../core/http/HttpError.js'; import type { Actor } from '../../core/actor.js'; import { PuterDriver } from '../types.js'; +import { secureFetch } from '../../util/secureHttp.js'; import { AI_CONCURRENT, AI_RATE_LIMIT } from '../util/aiLimits.js'; import { GeminiVideoProvider } from './providers/gemini/GeminiVideoProvider.js'; import { OpenAIVideoProvider } from './providers/openai/OpenAIVideoProvider.js'; @@ -399,7 +400,14 @@ export class VideoGenerationDriver extends PuterDriver { contentType = header.match(/data:(.*?);/)?.[1] ?? 'video/mp4'; buffer = Buffer.from(result.substring(commaIdx + 1), 'base64'); } else { - const response = await fetch(result); + // Provider-minted URL, but fetched with the same SSRF guards + // as the input paths: it reaches an unauthenticated GET whose + // body lands in the user's filesystem. skipProxy because + // generated media is ours to download directly, not user + // input to screen. + const response = await secureFetch(result, { + skipProxy: true, + }); if (!response.ok) { throw new HttpError( 502,