From 97487e7ce1758b2d0b246d9fee60445fa57038e1 Mon Sep 17 00:00:00 2001 From: 404oops Date: Fri, 16 Oct 2026 03:19:18 +0200 Subject: [PATCH] Bound unary Argon requests and reject unsupported content --- .../ai-chat/ChatCompletionDriver.test.ts | 388 +++++++++++++----- .../drivers/ai-chat/ChatCompletionDriver.ts | 10 +- .../gemini/GeminiChatProvider.native.test.ts | 327 ++++++++++++--- .../providers/gemini/GeminiChatProvider.ts | 37 +- .../ai-chat/providers/gemini/geminiNative.ts | 18 +- .../ai-chat/providers/gemini/models.ts | 6 +- src/backend/drivers/ai-chat/types.ts | 13 +- src/docs/src/rate-limits-and-quotas.md | 2 + 8 files changed, 620 insertions(+), 181 deletions(-) diff --git a/src/backend/drivers/ai-chat/ChatCompletionDriver.test.ts b/src/backend/drivers/ai-chat/ChatCompletionDriver.test.ts index 09e3b53f1..c2c691a04 100644 --- a/src/backend/drivers/ai-chat/ChatCompletionDriver.test.ts +++ b/src/backend/drivers/ai-chat/ChatCompletionDriver.test.ts @@ -20,19 +20,17 @@ /** * Offline unit tests for ChatCompletionDriver. * - * Boots a real PuterServer (in-memory sqlite + dynamo + s3 + mock - * redis) so the driver runs against the live `EventClient`, - * `MeteringService`, stores, and `FSService`. The driver is - * exercised through its public `complete()` interface against the - * always-available `FakeChatProvider` — no real upstream API keys - * are needed because no `config.providers.*` entries are set, so + * Boots a real PuterServer (in-memory sqlite + dynamo + s3 + mock redis) so the + * driver runs against the live `EventClient`, `MeteringService`, stores, and + * `FSService`. The driver is exercised through its public `complete()` + * interface against the always-available `FakeChatProvider` — no real upstream + * API keys are needed because no `config.providers.*` entries are set, so * `fake-chat` is the only provider registered. * - * Provider behaviour (request shape, streaming dialect) is covered - * separately by each provider's test file. This file pins down - * driver-level behaviour: auth gate, model resolution, validation - * event routing, credit/quota gates, max_tokens cap, fallback, - * event emission and cost calculation. + * Provider behaviour (request shape, streaming dialect) is covered separately + * by each provider's test file. This file pins down driver-level behaviour: + * auth gate, model resolution, validation event routing, credit/quota gates, + * max_tokens cap, fallback, event emission and cost calculation. */ import type { Readable } from 'node:stream'; import { @@ -256,7 +254,10 @@ describe('ChatCompletionDriver model map', () => { describe('ChatCompletionDriver.complete auth and model resolution', () => { it('uses a free fake completion in test mode without taking a credit hold', async () => { - const creditGate = vi.spyOn(server.services.metering, 'getUsageHeadroom'); + const creditGate = vi.spyOn( + server.services.metering, + 'getUsageHeadroom', + ); const complete = vi.spyOn(FakeChatProvider.prototype, 'complete'); const result = await withTestActor(() => @@ -290,10 +291,15 @@ describe('ChatCompletionDriver.complete auth and model resolution', () => { ); expect(complete.mock.calls[0]![0].model).toBe('fake'); - const streamResult = result as unknown as { dataType: string; stream: Readable }; + const streamResult = result as unknown as { + dataType: string; + stream: Readable; + }; expect(streamResult.dataType).toBe('stream'); const chunks = await collectStream(streamResult.stream); - expect(chunks).toContainEqual(expect.objectContaining({ type: 'text' })); + expect(chunks).toContainEqual( + expect.objectContaining({ type: 'text' }), + ); }); it('throws 401 when no actor is in context', async () => { @@ -368,7 +374,10 @@ describe('ChatCompletionDriver.complete auth and model resolution', () => { const completeSpy = vi.spyOn(FakeChatProvider.prototype, 'complete'); completeSpy.mockResolvedValueOnce({ - message: { role: 'assistant', content: [{ type: 'text', text: 'ok' }] }, + message: { + role: 'assistant', + content: [{ type: 'text', text: 'ok' }], + }, usage: {}, finish_reason: 'stop', } as never); @@ -617,7 +626,11 @@ describe('ChatCompletionDriver.complete events and cost emission', () => { id: 'priced', aliases: [], costs_currency: 'usd-cents', - costs: { input_tokens: 1000, output_tokens: 2000, web_search_requests: 1_000_000 }, + costs: { + input_tokens: 1000, + output_tokens: 2000, + web_search_requests: 1_000_000, + }, long_context_pricing: { threshold: 5, input_multiplier: 2, @@ -633,7 +646,11 @@ describe('ChatCompletionDriver.complete events and cost emission', () => { role: 'assistant', content: [{ type: 'text', text: 'ok' }], }, - usage: { input_tokens: 10, output_tokens: 7, web_search_requests: 1 }, + usage: { + input_tokens: 10, + output_tokens: 7, + web_search_requests: 1, + }, finish_reason: 'stop', } as never); @@ -751,7 +768,10 @@ 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, 'getUsageHeadroom').mockResolvedValue({ balance: 0, held: 0 }); + vi.spyOn( + server.services.metering, + 'getUsageHeadroom', + ).mockResolvedValue({ balance: 0, held: 0 }); await expect( withTestActor(() => @@ -786,7 +806,10 @@ describe('ChatCompletionDriver.complete credit gate and max_tokens cap', () => { // Approx input cost is tiny (very short prompt), so allowed // output ≈ 100_000 / 2000 = 50 tokens, comfortably below the // caller's 10_000 ceiling. - vi.spyOn(server.services.metering, 'getUsageHeadroom').mockResolvedValue({ balance: 100_000, held: 0 }); + vi.spyOn( + server.services.metering, + 'getUsageHeadroom', + ).mockResolvedValue({ balance: 100_000, held: 0 }); const completeSpy = vi .spyOn(FakeChatProvider.prototype, 'complete') @@ -813,62 +836,102 @@ describe('ChatCompletionDriver.complete credit gate and max_tokens cap', () => { expect(passed.max_tokens!).toBeGreaterThan(0); }); - it.each(['multipliers', 'catalog overrides'])('caps `max_tokens` at the conditional output rate using %s', async pricingKind => { + it.each(['multipliers', 'catalog overrides'])( + 'caps `max_tokens` at the conditional output rate using %s', + async (pricingKind) => { + vi.spyOn( + FakeChatProvider.prototype, + 'models', + ).mockResolvedValueOnce([ + { + id: 'capme', + aliases: [], + costs_currency: 'usd-cents', + costs: { input_tokens: 1000, output_tokens: 2000 }, + ...(pricingKind === 'multipliers' + ? { + long_context_pricing: { + threshold: 10, + input_multiplier: 2, + output_multiplier: 1.5, + }, + } + : { + pricingOverrides: [ + { + minPromptTokens: 10, + costs: { + input_tokens: 2000, + output_tokens: 3000, + }, + }, + ], + }), + max_tokens: 8192, + }, + ]); + const d = await makeDriver(); + + // A ~100-token prompt is past the threshold. 1_000_000 microcents at + // 2000 * 1.5 per output token leaves at most 333 output tokens before + // the prompt is paid for; the standard rate would allow ~450 after it. + vi.spyOn( + server.services.metering, + 'getUsageHeadroom', + ).mockResolvedValue({ balance: 1_000_000, held: 0 }); + + 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(() => + d.complete({ + model: 'capme', + messages: [{ role: 'user', content: 'x'.repeat(400) }], + max_tokens: 10_000, + }), + ); + + const passed = completeSpy.mock.calls[0]![0] as ICompleteArguments; + expect(passed.max_tokens!).toBeLessThanOrEqual(333); + expect(passed.max_tokens!).toBeGreaterThan(0); + }, + ); + + it('reserves a catalog fixed request fee before calling the provider', async () => { vi.spyOn(FakeChatProvider.prototype, 'models').mockResolvedValueOnce([ { - id: 'capme', - aliases: [], + id: 'request-fee', costs_currency: 'usd-cents', - costs: { input_tokens: 1000, output_tokens: 2000 }, - ...(pricingKind === 'multipliers' ? { long_context_pricing: { - threshold: 10, - input_multiplier: 2, - output_multiplier: 1.5, - } } : { pricingOverrides: [{ minPromptTokens: 10, costs: { input_tokens: 2000, output_tokens: 3000 } }] }), - max_tokens: 8192, + max_tokens: 1000, + costs: { + input_tokens: 1, + output_tokens: 1, + request: 1_000_001, + }, }, ]); const d = await makeDriver(); - - // A ~100-token prompt is past the threshold. 1_000_000 microcents at - // 2000 * 1.5 per output token leaves at most 333 output tokens before - // the prompt is paid for; the standard rate would allow ~450 after it. - vi.spyOn(server.services.metering, 'getUsageHeadroom').mockResolvedValue({ balance: 1_000_000, held: 0 }); - - 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(() => - d.complete({ - model: 'capme', - messages: [{ role: 'user', content: 'x'.repeat(400) }], - max_tokens: 10_000, - }), - ); - - const passed = completeSpy.mock.calls[0]![0] as ICompleteArguments; - expect(passed.max_tokens!).toBeLessThanOrEqual(333); - expect(passed.max_tokens!).toBeGreaterThan(0); - }); - - it('reserves a catalog fixed request fee before calling the provider', async () => { - vi.spyOn(FakeChatProvider.prototype, 'models').mockResolvedValueOnce([{ - id: 'request-fee', costs_currency: 'usd-cents', max_tokens: 1000, - costs: { input_tokens: 1, output_tokens: 1, request: 1_000_001 }, - }]); - const d = await makeDriver(); - vi.spyOn(server.services.metering, 'getUsageHeadroom').mockResolvedValue({ balance: 1_000_000, held: 0 }); + vi.spyOn( + server.services.metering, + 'getUsageHeadroom', + ).mockResolvedValue({ balance: 1_000_000, held: 0 }); const completeSpy = vi.spyOn(FakeChatProvider.prototype, 'complete'); - await expect(withTestActor(() => d.complete({ model: 'request-fee', messages: [{ role: 'user', content: 'hi' }] }))) - .rejects.toMatchObject({ statusCode: 402 }); + await expect( + withTestActor(() => + d.complete({ + model: 'request-fee', + messages: [{ role: 'user', content: 'hi' }], + }), + ), + ).rejects.toMatchObject({ statusCode: 402 }); expect(completeSpy).not.toHaveBeenCalled(); }); @@ -889,7 +952,10 @@ describe('ChatCompletionDriver.complete credit gate and max_tokens cap', () => { // 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, 'getUsageHeadroom').mockResolvedValue({ balance: 100, held: 0 }); + vi.spyOn( + server.services.metering, + 'getUsageHeadroom', + ).mockResolvedValue({ balance: 100, held: 0 }); const completeSpy = vi.spyOn(FakeChatProvider.prototype, 'complete'); @@ -927,23 +993,31 @@ describe('ChatCompletionDriver.complete credit gate and max_tokens cap', () => { // extraCost alone (50_000) eats most of a 100_000-microcent balance; // the fast cost keys (half the standard rate) are what is left to // afford output with. - (FakeChatProvider.prototype as unknown as { - requestPricing: NonNullable; - }).requestPricing = () => ({ + ( + FakeChatProvider.prototype as unknown as { + requestPricing: NonNullable; + } + ).requestPricing = () => ({ inputKey: 'fast_input_tokens', outputKey: 'fast_output_tokens', extraCost: 50_000, }); try { - vi.spyOn(server.services.metering, 'getUsageHeadroom').mockResolvedValue({ + vi.spyOn( + server.services.metering, + 'getUsageHeadroom', + ).mockResolvedValue({ balance: 100_000, held: 0, }); const completeSpy = vi .spyOn(FakeChatProvider.prototype, 'complete') .mockResolvedValueOnce({ - message: { role: 'assistant', content: [{ type: 'text', text: 'ok' }] }, + message: { + role: 'assistant', + content: [{ type: 'text', text: 'ok' }], + }, usage: { input_tokens: 1, output_tokens: 1 }, finish_reason: 'stop', } as never); @@ -963,12 +1037,13 @@ describe('ChatCompletionDriver.complete credit gate and max_tokens cap', () => { expect(passed.max_tokens!).toBeLessThanOrEqual(50); expect(passed.max_tokens!).toBeGreaterThan(30); } finally { - delete (FakeChatProvider.prototype as unknown as Record) - .requestPricing; + delete ( + FakeChatProvider.prototype as unknown as Record + ).requestPricing; } }); - it('rejects with 402 when a low balance cannot cover a tool\'s extraCost reservation', async () => { + it("caps `max_tokens` and the hold at the provider's requestPricing maxTokens", async () => { vi.spyOn(FakeChatProvider.prototype, 'models').mockResolvedValueOnce([ { id: 'capme', @@ -980,16 +1055,87 @@ describe('ChatCompletionDriver.complete credit gate and max_tokens cap', () => { ]); const d = await makeDriver(); - (FakeChatProvider.prototype as unknown as { - requestPricing: NonNullable; - }).requestPricing = () => ({ extraCost: 1_000_000 }); + ( + FakeChatProvider.prototype as unknown as { + requestPricing: NonNullable; + } + ).requestPricing = () => ({ maxTokens: 100 }); try { - vi.spyOn(server.services.metering, 'getUsageHeadroom').mockResolvedValue({ + vi.spyOn( + server.services.metering, + 'getUsageHeadroom', + ).mockResolvedValue({ + balance: 100_000_000, + held: 0, + }); + const reserve = vi.spyOn( + server.services.metering, + 'reserveCredits', + ); + 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(() => + d.complete({ + model: 'capme', + messages: [{ role: 'user', content: 'hi' }], + max_tokens: 10_000, + }), + ); + + const passed = completeSpy.mock.calls[0]![0] as ICompleteArguments; + expect(passed.max_tokens).toBe(100); + // The balance and the request both allow more; the hold is the + // tiny prompt plus 100 output tokens at 2000. + const [, held] = reserve.mock.calls[0]!; + expect(held).toBeGreaterThanOrEqual(100 * 2000); + expect(held).toBeLessThan(101 * 2000); + } finally { + delete ( + FakeChatProvider.prototype as unknown as Record + ).requestPricing; + } + }); + + it("rejects with 402 when a low balance cannot cover a tool's extraCost reservation", async () => { + vi.spyOn(FakeChatProvider.prototype, 'models').mockResolvedValueOnce([ + { + id: 'capme', + aliases: [], + costs_currency: 'usd-cents', + costs: { input_tokens: 1000, output_tokens: 2000 }, + max_tokens: 8192, + }, + ]); + const d = await makeDriver(); + + ( + FakeChatProvider.prototype as unknown as { + requestPricing: NonNullable; + } + ).requestPricing = () => ({ extraCost: 1_000_000 }); + + try { + vi.spyOn( + server.services.metering, + 'getUsageHeadroom', + ).mockResolvedValue({ balance: 100, held: 0, }); - const completeSpy = vi.spyOn(FakeChatProvider.prototype, 'complete'); + const completeSpy = vi.spyOn( + FakeChatProvider.prototype, + 'complete', + ); await expect( withTestActor(() => @@ -999,11 +1145,15 @@ describe('ChatCompletionDriver.complete credit gate and max_tokens cap', () => { tools: [{ type: 'web_search' }], }), ), - ).rejects.toMatchObject({ statusCode: 402, legacyCode: 'insufficient_funds' }); + ).rejects.toMatchObject({ + statusCode: 402, + legacyCode: 'insufficient_funds', + }); expect(completeSpy).not.toHaveBeenCalled(); } finally { - delete (FakeChatProvider.prototype as unknown as Record) - .requestPricing; + delete ( + FakeChatProvider.prototype as unknown as Record + ).requestPricing; } }); @@ -1011,14 +1161,20 @@ describe('ChatCompletionDriver.complete credit gate and max_tokens cap', () => { // unexempted `web_search_preview` became an unnamed function tool, which // OpenAI rejects and the provider's search pricing no longer recognizes. it('passes web_search_preview through to the provider unchanged', async () => { - vi.spyOn(server.services.metering, 'getUsageHeadroom').mockResolvedValue({ + vi.spyOn( + server.services.metering, + 'getUsageHeadroom', + ).mockResolvedValue({ balance: 100_000_000, held: 0, }); const completeSpy = vi .spyOn(FakeChatProvider.prototype, 'complete') .mockResolvedValueOnce({ - message: { role: 'assistant', content: [{ type: 'text', text: 'ok' }] }, + message: { + role: 'assistant', + content: [{ type: 'text', text: 'ok' }], + }, usage: { input_tokens: 1, output_tokens: 1 }, finish_reason: 'stop', } as never); @@ -1103,7 +1259,10 @@ describe('ChatCompletionDriver.complete credit gate and max_tokens cap', () => { }, ]); const d = await makeDriver(); - vi.spyOn(server.services.metering, 'getUsageHeadroom').mockResolvedValue({ balance: 1_000_000, held: 0 }); + vi.spyOn( + server.services.metering, + 'getUsageHeadroom', + ).mockResolvedValue({ balance: 1_000_000, held: 0 }); vi.spyOn( server.services.metering, 'getActorSubscription', @@ -1134,7 +1293,10 @@ describe('ChatCompletionDriver.complete credit gate and max_tokens cap', () => { }, ]); const d = await makeDriver(); - vi.spyOn(server.services.metering, 'getUsageHeadroom').mockResolvedValue({ balance: 1_000_000, held: 0 }); + vi.spyOn( + server.services.metering, + 'getUsageHeadroom', + ).mockResolvedValue({ balance: 1_000_000, held: 0 }); vi.spyOn( server.services.metering, 'getActorSubscription', @@ -1173,7 +1335,10 @@ 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, 'getUsageHeadroom').mockResolvedValue({ balance: 1_000_000, held: 0 }); + vi.spyOn( + server.services.metering, + 'getUsageHeadroom', + ).mockResolvedValue({ balance: 1_000_000, held: 0 }); await expect( withTestActor(() => @@ -1303,9 +1468,10 @@ describe('ChatCompletionDriver Claude usage.thinking_tokens', () => { const d = await makeClaudeDriver(); // Affords exactly 300 output tokens at 1000 µ¢ each: above the real // 233 output tokens, below 233 + 110 thinking_tokens (343). - vi.spyOn(server.services.metering, 'getUsageHeadroom').mockResolvedValue( - { balance: 300_000, held: 0 }, - ); + vi.spyOn( + server.services.metering, + 'getUsageHeadroom', + ).mockResolvedValue({ balance: 300_000, held: 0 }); vi.spyOn(ClaudeProvider.prototype, 'complete').mockResolvedValueOnce({ message: { role: 'assistant', content: 'hi' }, finish_reason: 'stop', @@ -1323,8 +1489,7 @@ describe('ChatCompletionDriver Claude usage.thinking_tokens', () => { messages: [{ role: 'user', content: 'hi' }], }); return Context.get('driverMetadata') as - | Record - | undefined; + Record | undefined; }); expect(metadata?.usage_limited).toBeUndefined(); @@ -1370,7 +1535,10 @@ describe('ChatCompletionDriver.complete normalization', () => { ])( 'normalizes a raw wire-form tool_choice (%j) before it reaches the provider', async (wireToolChoice, normalized) => { - const completeSpy = vi.spyOn(FakeChatProvider.prototype, 'complete'); + const completeSpy = vi.spyOn( + FakeChatProvider.prototype, + 'complete', + ); completeSpy.mockResolvedValueOnce({ message: { role: 'assistant', content: 'ok' }, usage: {}, @@ -1672,9 +1840,7 @@ describe('ChatCompletionDriver.complete OpenAI-shape normalization', () => { content_type: 'application/x-ndjson', }); // Drain so the fake provider's populator finishes cleanly. - await collectStream( - (res as unknown as { stream: Readable }).stream, - ); + await collectStream((res as unknown as { stream: Readable }).stream); }); it('a blocked prompt rerouted to fake-chat keeps its historical native shape', async () => { @@ -1733,8 +1899,11 @@ describe('ChatCompletionDriver.complete fallback and error envelope', () => { legacyCode: 'upstream_timeout', message: 'AI provider timed out', }); - const attempts = (caught as unknown as { fields: { attempts: { timedOut?: boolean }[] } }) - .fields.attempts; + const attempts = ( + caught as unknown as { + fields: { attempts: { timedOut?: boolean }[] }; + } + ).fields.attempts; expect(attempts).toHaveLength(1); expect(attempts[0].timedOut).toBe(true); }); @@ -1829,8 +1998,9 @@ describe('ChatCompletionDriver.complete fallback and error envelope', () => { expect(caught).toBeInstanceOf(HttpError); expect(caught!.statusCode).toBe(500); expect(caught!.message).toBe('All providers failed'); - const attempts = (caught as unknown as { fields: { attempts: unknown[] } }) - .fields.attempts; + const attempts = ( + caught as unknown as { fields: { attempts: unknown[] } } + ).fields.attempts; expect(Array.isArray(attempts)).toBe(true); // No second provider serves `fake`, so we record exactly one attempt // before falling out of the loop. @@ -2020,8 +2190,9 @@ describe('ChatCompletionDriver.countTokens', () => { it('prefers the provider’s own countTokens hook when it exists', async () => { const hook = vi.fn().mockResolvedValue(42); - (FakeChatProvider.prototype as unknown as { countTokens?: unknown }).countTokens = - hook; + ( + FakeChatProvider.prototype as unknown as { countTokens?: unknown } + ).countTokens = hook; try { const result = await driver[COUNT_TOKENS]({ model: 'fake', @@ -2034,8 +2205,11 @@ describe('ChatCompletionDriver.countTokens', () => { provider: 'fake-chat', }); } finally { - delete (FakeChatProvider.prototype as unknown as { countTokens?: unknown }) - .countTokens; + delete ( + FakeChatProvider.prototype as unknown as { + countTokens?: unknown; + } + ).countTokens; } }); diff --git a/src/backend/drivers/ai-chat/ChatCompletionDriver.ts b/src/backend/drivers/ai-chat/ChatCompletionDriver.ts index 7c90ff726..00073b9b6 100644 --- a/src/backend/drivers/ai-chat/ChatCompletionDriver.ts +++ b/src/backend/drivers/ai-chat/ChatCompletionDriver.ts @@ -1404,10 +1404,11 @@ export class ChatCompletionDriver extends PuterDriver { const { promptTokenEstimate, requestedMaxTokens, pricingTime } = estimates; // A provider can name its own cost keys for this request (Claude's - // fast-mode rate card) and size an extra reservation the per-token - // math can't see (a web-search/advisor tool that could run several - // times before the model answers) — the same clamps it applies when - // building the request, so the hold and the request agree. + // fast-mode rate card), cap its output, and size an extra reservation + // the per-token math can't see (a web-search/advisor tool that could + // run several times before the model answers) — the same clamps it + // applies when building the request, so the hold and the request + // agree. const pricing = this.#providers[model.provider!]?.requestPricing?.(args, model, { promptTokenEstimate, @@ -1475,6 +1476,7 @@ export class ChatCompletionDriver extends PuterDriver { const limit = Math.min( requestedMaxTokens ?? Number.POSITIVE_INFINITY, outputCeiling(model, promptTokenEstimate), + positiveOrInfinity(pricing.maxTokens), ); const cap = Math.floor( Math.min(limit, affordableOutputTokens(remainingCredits)), diff --git a/src/backend/drivers/ai-chat/providers/gemini/GeminiChatProvider.native.test.ts b/src/backend/drivers/ai-chat/providers/gemini/GeminiChatProvider.native.test.ts index b0806c90c..815bad499 100644 --- a/src/backend/drivers/ai-chat/providers/gemini/GeminiChatProvider.native.test.ts +++ b/src/backend/drivers/ai-chat/providers/gemini/GeminiChatProvider.native.test.ts @@ -18,9 +18,9 @@ */ /** - * GeminiChatProvider models served over the native `generateContent` API. - * Boots a real PuterServer so metering is recorded end to end; the Google and - * OpenAI SDKs are mocked at the module boundary, which is the network egress. + * GeminiChatProvider models served over the native `generateContent` API. Boots + * a real PuterServer so metering is recorded end to end; the Google and OpenAI + * SDKs are mocked at the module boundary, which is the network egress. */ import { Writable } from 'node:stream'; @@ -40,7 +40,11 @@ import { Context } from '../../../../core/context.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 { ChatCompletionDriver } from '../../ChatCompletionDriver.js'; import type { ICompleteArguments, UsageDetails } from '../../types.js'; import { normalize_messages } from '../../utils/Messages.js'; import { trackedOutputTokens } from '../../utils/pricing.js'; @@ -181,7 +185,9 @@ const runStream = async (args: Partial) => { }); await ( result as { - init_chat_stream: (p: { chatStream: AIChatStream }) => Promise; + init_chat_stream: (p: { + chatStream: AIChatStream; + }) => Promise; } ).init_chat_stream({ chatStream }); const events = chunks @@ -258,7 +264,10 @@ describe('gemini-4-argon catalog and routing', () => { thinking_tokens: 3, grounding_requests: 0, }); - expect(result.usageDetails).toEqual({ inputTokens: 10, outputTokens: 1 }); + expect(result.usageDetails).toEqual({ + inputTokens: 10, + outputTokens: 1, + }); }); it.each([ @@ -358,7 +367,8 @@ describe('gemini-4-argon thinking levels', () => { async (effort) => { await expect( complete({ - reasoning_effort: effort as ICompleteArguments['reasoning_effort'], + reasoning_effort: + effort as ICompleteArguments['reasoning_effort'], }), ).rejects.toMatchObject({ statusCode: 400, @@ -487,7 +497,9 @@ describe('gemini-4-argon text completion', () => { { type: 'text', text: 'What is this?' }, { type: 'image_url', - image_url: { url: 'data:image/png;base64,iVBORw0K' }, + image_url: { + url: 'data:image/png;base64,iVBORw0K', + }, }, ], }, @@ -505,24 +517,70 @@ describe('gemini-4-argon text completion', () => { }); }); +describe('gemini-4-argon unsupported content', () => { + const AUDIO = { + type: 'input_audio', + input_audio: { data: 'UklGRg==', format: 'wav' }, + }; + const FILE = { + type: 'file', + file: { + filename: 'notes.pdf', + file_data: 'data:application/pdf;base64,JVBERi0=', + }, + }; + + it.each([ + [ + 'text and audio', + 'input_audio', + [{ type: 'text', text: 'Transcribe this' }, AUDIO], + ], + ['audio only', 'input_audio', [AUDIO]], + [ + 'text and a file', + 'file', + [{ type: 'text', text: 'Summarize this' }, FILE], + ], + ])( + 'rejects %s before calling Google', + async (_label, partType, content) => { + await expect( + complete({ messages: [{ role: 'user', content }] }), + ).rejects.toMatchObject({ + statusCode: 400, + legacyCode: 'bad_request', + message: `${partType} content is not supported by this model`, + }); + expect(generateContent).not.toHaveBeenCalled(); + }, + ); +}); + describe('gemini-4-argon pricing by date', () => { it.each([ ['2026-11-30T23:59:59.999Z', 200, 1000, 10], ['2026-12-01T00:00:00.000Z', 400, 2000, 20], ['2027-01-15T08:00:00.000Z', 400, 2000, 20], - ])('prices a request at %s at the matching rates', async (iso, input, output, cached) => { - generateContent.mockResolvedValueOnce( - textResponse('OK', { ...LIVE_USAGE, cachedContentTokenCount: 6 }), - ); - await complete({}, new Date(iso)); - const [, , , costs] = recordSpy.mock.calls[0]!; - expect(costs).toEqual({ - prompt_tokens: 30 * input, - completion_tokens: 5 * output, - thinking_tokens: 69 * output, - cached_tokens: 6 * cached, - }); - }); + ])( + 'prices a request at %s at the matching rates', + async (iso, input, output, cached) => { + generateContent.mockResolvedValueOnce( + textResponse('OK', { + ...LIVE_USAGE, + cachedContentTokenCount: 6, + }), + ); + await complete({}, new Date(iso)); + const [, , , costs] = recordSpy.mock.calls[0]!; + expect(costs).toEqual({ + prompt_tokens: 30 * input, + completion_tokens: 5 * output, + thinking_tokens: 69 * output, + cached_tokens: 6 * cached, + }); + }, + ); it('settles at the instant the driver priced the attempt at, not the clock', async () => { vi.useFakeTimers({ @@ -550,8 +608,13 @@ describe('gemini-4-argon streaming', () => { generateContentStream.mockResolvedValueOnce( asAsyncIterable([ { - candidates: [{ content: { role: 'model', parts: [{ text: 'O' }] } }], - usageMetadata: { promptTokenCount: 36, totalTokenCount: 36 }, + candidates: [ + { content: { role: 'model', parts: [{ text: 'O' }] } }, + ], + usageMetadata: { + promptTokenCount: 36, + totalTokenCount: 36, + }, }, { candidates: [ @@ -597,7 +660,9 @@ describe('gemini-4-argon streaming', () => { generateContentStream.mockResolvedValueOnce( asAsyncIterable([ { - candidates: [{ content: { role: 'model', parts: [{ text: 'O' }] } }], + candidates: [ + { content: { role: 'model', parts: [{ text: 'O' }] } }, + ], usageMetadata: { promptTokenCount: 5, candidatesTokenCount: 1, @@ -634,30 +699,36 @@ describe('gemini-4-argon streaming', () => { it.each([ ['no usage metadata', undefined], ['empty usage metadata', {}], - ])('leaves a stream with %s unreported so the driver estimates it', async (_label, usageMetadata) => { - generateContentStream.mockResolvedValueOnce( - asAsyncIterable([ - { - candidates: [ - { - content: { role: 'model', parts: [{ text: 'OK' }] }, - finishReason: 'STOP', - }, - ], - usageMetadata, - }, - ]), - ); + ])( + 'leaves a stream with %s unreported so the driver estimates it', + async (_label, usageMetadata) => { + generateContentStream.mockResolvedValueOnce( + asAsyncIterable([ + { + candidates: [ + { + content: { + role: 'model', + parts: [{ text: 'OK' }], + }, + finishReason: 'STOP', + }, + ], + usageMetadata, + }, + ]), + ); - const { events, chatStream } = await runStream({}); + const { events, chatStream } = await runStream({}); - expect(events.filter((e) => e.type === 'text')).toHaveLength(1); - // Null is what makes the driver charge its output estimate. - expect(chatStream.reportedUsage).toBeNull(); - expect( - recordSpy.mock.calls.flatMap(([usage]) => Object.values(usage)), - ).toEqual([]); - }); + expect(events.filter((e) => e.type === 'text')).toHaveLength(1); + // Null is what makes the driver charge its output estimate. + expect(chatStream.reportedUsage).toBeNull(); + expect( + recordSpy.mock.calls.flatMap(([usage]) => Object.values(usage)), + ).toEqual([]); + }, + ); it('streams a tool call with its thought signature', async () => { generateContentStream.mockResolvedValueOnce( @@ -700,6 +771,103 @@ describe('gemini-4-argon streaming', () => { }); }); +describe('gemini-4-argon output cap', () => { + it.each([ + ['no max_tokens', undefined], + ['a larger max_tokens', 65_536], + ])( + 'caps a non-streamed request with %s at 32,768 output tokens', + async (_label, maxTokens) => { + generateContent.mockResolvedValueOnce(textResponse('OK')); + await complete({ max_tokens: maxTokens }); + expect(requestOf().config.maxOutputTokens).toBe(32_768); + }, + ); + + it('sends a stream the max_tokens it asked for', async () => { + generateContentStream.mockResolvedValueOnce( + asAsyncIterable([textResponse('OK')]), + ); + await runStream({ max_tokens: 65_536 }); + expect(requestOf(generateContentStream).config.maxOutputTokens).toBe( + 65_536, + ); + }); + + describe('through the driver', () => { + let driver: ChatCompletionDriver; + let reserve: MockInstance; + + beforeEach(async () => { + // Launch rates, so a hold converts to output tokens exactly. + vi.useFakeTimers({ toFake: ['Date'], now: LAUNCH }); + driver = new ChatCompletionDriver( + { + providers: { + gemini: { apiKey: 'test-key' }, + ollama: { enabled: false }, + }, + } as never, + server.clients, + server.stores, + server.services, + ); + await driver.onServerStart(); + // Enough balance that only the request and the model bound output. + vi.spyOn( + server.services.metering, + 'getUsageHeadroom', + ).mockResolvedValue({ balance: Number.MAX_SAFE_INTEGER, held: 0 }); + reserve = vi.spyOn(server.services.metering, 'reserveCredits'); + }); + + afterEach(() => { + driver.onServerShutdown(); + vi.useRealTimers(); + }); + + const viaDriver = (args: Partial) => + withTestActor(() => + driver.complete({ + model: 'gemini-4-argon', + messages: [{ role: 'user', content: 'Say OK' }], + ...args, + } as ICompleteArguments), + ); + + // The two-token prompt costs less than one output token. + const heldOutputTokens = () => + Math.floor( + reserve.mock.calls[0]![1] / ARGON.costs.completion_tokens, + ); + + it.each([ + ['no max_tokens', 32_768, undefined], + ['a larger max_tokens', 32_768, 65_536], + ['a smaller max_tokens', 1_000, 1_000], + ])( + 'sends and holds a non-streamed request with %s at %i output tokens', + async (_label, expected, maxTokens) => { + generateContent.mockResolvedValueOnce(textResponse('OK')); + await viaDriver({ max_tokens: maxTokens }); + expect(requestOf().config.maxOutputTokens).toBe(expected); + expect(heldOutputTokens()).toBe(expected); + }, + ); + + it('sends and holds a stream at the full 65,536-token limit', async () => { + generateContentStream.mockResolvedValueOnce( + asAsyncIterable([textResponse('OK')]), + ); + await readStreamResult(await viaDriver({ stream: true })); + expect( + requestOf(generateContentStream).config.maxOutputTokens, + ).toBe(65_536); + expect(heldOutputTokens()).toBe(65_536); + }); + }); +}); + describe('gemini-4-argon finish reasons', () => { const LIVE_METERED = { prompt_tokens: 36, @@ -732,7 +900,7 @@ describe('gemini-4-argon finish reasons', () => { ['UNEXPECTED_TOOL_CALL', 'UNEXPECTED_TOOL_CALL'], ['LANGUAGE', 'LANGUAGE'], ['OTHER', 'OTHER'], - ['CONTINUATION', 'CONTINUATION'], + ['CONTINUATION', 'length'], ['SOME_FUTURE_REASON', 'SOME_FUTURE_REASON'], ] as const; @@ -765,6 +933,38 @@ describe('gemini-4-argon finish reasons', () => { }); }, ); + + // The raw API puts a resume token beside CONTINUATION; nothing here can + // use it, so it must not reach the caller. + const continued = { + candidates: [ + { + content: { role: 'model', parts: [{ text: 'Partial' }] }, + finishReason: 'CONTINUATION', + continuationToken: 'resume-token', + }, + ], + usageMetadata: LIVE_USAGE, + }; + + it('ends a CONTINUATION cut as length without its resume token', async () => { + generateContent.mockResolvedValueOnce(continued); + const result = await complete({}); + expect(result).toMatchObject({ finish_reason: 'length' }); + expect(JSON.stringify(result)).not.toContain('resume-token'); + }); + + it('streams a CONTINUATION cut as a max_tokens stop without its resume token', async () => { + generateContentStream.mockResolvedValueOnce( + asAsyncIterable([continued]), + ); + const { events } = await runStream({}); + expect(events.find((e) => e.type === 'usage')).toMatchObject({ + finish_reason: 'length', + stopReason: 'max_tokens', + }); + expect(JSON.stringify(events)).not.toContain('resume-token'); + }); }); describe('gemini-4-argon blocked prompts', () => { @@ -807,19 +1007,28 @@ describe('gemini-4-argon blocked prompts', () => { }); it.each([ - ['an unspecified block reason', { blockReason: 'BLOCKED_REASON_UNSPECIFIED' }], - ['the API docs spelling of it', { blockReason: 'BLOCK_REASON_UNSPECIFIED' }], + [ + 'an unspecified block reason', + { blockReason: 'BLOCKED_REASON_UNSPECIFIED' }, + ], + [ + 'the API docs spelling of it', + { blockReason: 'BLOCK_REASON_UNSPECIFIED' }, + ], ['no prompt feedback', undefined], - ])('still rejects a response with no candidates and %s', async (_label, promptFeedback) => { - generateContent.mockResolvedValueOnce({ - ...(promptFeedback ? { promptFeedback } : {}), - usageMetadata: PROMPT_USAGE, - }); - await expect(complete({})).rejects.toMatchObject({ - statusCode: 400, - legacyCode: 'bad_response', - }); - }); + ])( + 'still rejects a response with no candidates and %s', + async (_label, promptFeedback) => { + generateContent.mockResolvedValueOnce({ + ...(promptFeedback ? { promptFeedback } : {}), + usageMetadata: PROMPT_USAGE, + }); + await expect(complete({})).rejects.toMatchObject({ + statusCode: 400, + legacyCode: 'bad_response', + }); + }, + ); it('streams a blocked prompt as a content_filter end with its usage', async () => { generateContentStream.mockResolvedValueOnce( diff --git a/src/backend/drivers/ai-chat/providers/gemini/GeminiChatProvider.ts b/src/backend/drivers/ai-chat/providers/gemini/GeminiChatProvider.ts index 7b8c5d8de..04b6f406e 100644 --- a/src/backend/drivers/ai-chat/providers/gemini/GeminiChatProvider.ts +++ b/src/backend/drivers/ai-chat/providers/gemini/GeminiChatProvider.ts @@ -41,13 +41,26 @@ import { type UsageSource, } from '../OpenAICompatProvider.js'; import { geminiNativeClient } from './geminiNative.js'; -import { GEMINI_MODELS, GEMINI_NATIVE_MODEL_IDS } from './models.js'; +import { + GEMINI_MODELS, + GEMINI_NATIVE_MODEL_IDS, + GEMINI_ARGON_NON_STREAM_MAX_TOKENS, +} from './models.js'; type GroundingContent = { grounding_metadata?: unknown }; /** The thinking levels the native models accept; medium when unset. */ const NATIVE_REASONING_EFFORTS = new Set(['low', 'medium', 'high']); +/** Google's cap on non-streamed native output; streams keep the catalog's. */ +const nativeOutputCap = ( + args: ICompleteArguments, + model: IChatModel, +): number | undefined => + model.id === 'gemini-4-argon' && !args.stream + ? GEMINI_ARGON_NON_STREAM_MAX_TOKENS + : undefined; + const COMPAT_BASE_URL = 'https://generativelanguage.googleapis.com/v1beta/openai/'; @@ -114,6 +127,15 @@ export class GeminiChatProvider extends OpenAICompatProvider { }); } + /** Sizes the credit hold to the same output cap `vendorParams` sends. */ + requestPricing( + args: ICompleteArguments, + model: IChatModel, + ): { maxTokens?: number } { + const maxTokens = nativeOutputCap(args, model); + return maxTokens === undefined ? {} : { maxTokens }; + } + protected override prepareMessages(messages: PuterMessage[]) { for (const message of messages) delete message.cache_control; return messages; @@ -134,8 +156,18 @@ export class GeminiChatProvider extends OpenAICompatProvider { ); } const hasTools = Array.isArray(params.tools) && params.tools.length > 0; + const outputCap = nativeOutputCap(args, model); + const requested = params.max_completion_tokens as number | undefined; return { ...params, + ...(outputCap !== undefined + ? { + max_completion_tokens: Math.min( + requested ?? outputCap, + outputCap, + ), + } + : {}), ...(effort ? { reasoning_effort: effort } : {}), ...(hasTools && args.tool_choice ? { tool_choice: toolChoiceToWire(args.tool_choice, 'chat') } @@ -152,8 +184,7 @@ export class GeminiChatProvider extends OpenAICompatProvider { if (GEMINI_NATIVE_MODEL_IDS.has(model.id)) { return nativeMeteredUsage( source.usage as unknown as - | GenerateContentResponseUsageMetadata - | undefined, + GenerateContentResponseUsageMetadata | undefined, ); } // Non-stream grounding metadata lives in choices[0].message; a stream diff --git a/src/backend/drivers/ai-chat/providers/gemini/geminiNative.ts b/src/backend/drivers/ai-chat/providers/gemini/geminiNative.ts index 81150d954..4149c5fa6 100644 --- a/src/backend/drivers/ai-chat/providers/gemini/geminiNative.ts +++ b/src/backend/drivers/ai-chat/providers/gemini/geminiNative.ts @@ -122,8 +122,10 @@ const contentParts = (content: ChatContent): Part[] => { for (const part of content) { if (typeof part === 'string') { if (part) parts.push({ text: part }); - } else if (part.type === 'text' && typeof part.text === 'string') { - if (part.text) parts.push({ text: part.text }); + } else if (part.type === 'text') { + if (typeof part.text === 'string' && part.text) { + parts.push({ text: part.text }); + } } else if (part.type === 'image_url') { // http(s) images were already inlined as data URLs. const uri = parseDataUri(part.image_url?.url ?? ''); @@ -132,6 +134,14 @@ const contentParts = (content: ChatContent): Part[] => { ? { inlineData: { mimeType: uri.mimeType, data: uri.data } } : unsupportedMediaTextPart('image could not be loaded'), ); + } else { + // Audio and file parts have no mapping here; dropping them would + // send Google a different prompt than the caller wrote. + throw new HttpError( + 400, + `${part.type ?? 'untyped'} content is not supported by this model`, + { legacyCode: 'bad_request' }, + ); } } return parts; @@ -280,7 +290,9 @@ const finishReasonOf = ( sawToolCall: boolean, ): string | undefined => { if (!reason) return undefined; - if (reason === 'MAX_TOKENS') return 'length'; + // CONTINUATION is a cut at Google's output limit, and the SDK drops the + // token that would resume it, so it reads as any other truncation. + if (reason === 'MAX_TOKENS' || reason === 'CONTINUATION') return 'length'; if (FILTERED_FINISH_REASONS.has(reason)) return 'content_filter'; return promoteStopForToolCalls( reason === 'STOP' ? 'stop' : reason, diff --git a/src/backend/drivers/ai-chat/providers/gemini/models.ts b/src/backend/drivers/ai-chat/providers/gemini/models.ts index e86477599..9bfb28162 100644 --- a/src/backend/drivers/ai-chat/providers/gemini/models.ts +++ b/src/backend/drivers/ai-chat/providers/gemini/models.ts @@ -24,12 +24,14 @@ export const GEMINI_NATIVE_MODEL_IDS: ReadonlySet = new Set([ 'gemini-4-argon', ]); -// Hardcoded from https://models.dev/api.json +// Google's unary model metadata reports this cap; streaming keeps the catalog cap. +export const GEMINI_ARGON_NON_STREAM_MAX_TOKENS = 32_768; + export const GEMINI_MODELS: IChatModel[] = [ { puterId: 'google:google/gemini-4-argon', id: 'gemini-4-argon', - // The driver has no audio or PDF part to forward. + // The native adapter forwards only text and image parts. modalities: { input: ['text', 'image'], output: ['text'], diff --git a/src/backend/drivers/ai-chat/types.ts b/src/backend/drivers/ai-chat/types.ts index be03981e4..09c18b403 100644 --- a/src/backend/drivers/ai-chat/types.ts +++ b/src/backend/drivers/ai-chat/types.ts @@ -308,8 +308,8 @@ export interface IChatProvider { /** * `model` is this provider's own catalog entry for `arg.model`, as the * driver resolved it; a provider resolves `arg.model` itself without it. - * `pricingTime` is the instant the driver's credit gate priced this - * attempt at; usage metered for the attempt is priced at the same instant. + * `pricingTime` is the instant the driver's credit gate priced this attempt + * at; usage metered for the attempt is priced at the same instant. */ complete( arg: ICompleteArguments, @@ -334,12 +334,19 @@ export interface IChatProvider { /** * Request-specific pricing the credit gate must reserve for (web search, * fast mode, the advisor tool), in µ¢, before the AI cost factor. + * `maxTokens` is an output cap the provider holds this request to; the gate + * caps `max_tokens` and sizes the hold to it as well. */ requestPricing?( args: ICompleteArguments, model: IChatModel, est: { promptTokenEstimate: number }, - ): { inputKey?: string; outputKey?: string; extraCost?: number }; + ): { + inputKey?: string; + outputKey?: string; + extraCost?: number; + maxTokens?: number; + }; /** * Exact prompt token count, when the vendor offers one (e.g. Claude's * `count_tokens`). diff --git a/src/docs/src/rate-limits-and-quotas.md b/src/docs/src/rate-limits-and-quotas.md index 1dfadd1e4..623f6ce91 100644 --- a/src/docs/src/rate-limits-and-quotas.md +++ b/src/docs/src/rate-limits-and-quotas.md @@ -74,6 +74,8 @@ An input over its limit is rejected before it reaches the model, as is an image Chat requests time out after 10 minutes, including streamed responses. See the [AI model directory](https://developer.puter.com/ai/models/) for model specifications. +A `gemini-4-argon` response that isn't streamed is capped at 32,768 output tokens, and a larger `max_tokens` is lowered to that cap. Streaming uses the model-directory limit. + Text-to-speech calls to ElevenLabs, xAI and Speechify time out after 2 minutes; xAI speech-to-text and voice-changer calls after 10 minutes, upload included. A call that times out before any audio or transcript comes back fails with `504 upstream_timeout`. The OpenAI- and Anthropic-compatible endpoints (`/puterai/openai/v1/*`, `/puterai/anthropic/v1/messages`) require a paid plan; a free account gets `402 subscription_required`. The same models are available to every account through `puter.ai.*` and `/drivers/call`, and the model catalogue endpoints are open to everyone.