lee/puter
1
0 Fork 0
mirror of https://github.com/HeyPuter/puter.git synced 2026-10-21 02:52:06 +00:00
Code Issues Actions 16 Packages Projects Releases Wiki Activity

Bound unary Argon requests and reject unsupported content

Browse Source
This commit is contained in:
404oops authored and GitHub committed 2026-10-16 21:22:01 -07:00
1 parent 2bc4106e13
commit 97487e7ce1
8 files changed
+620 -181
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

src/backend/drivers/ai-chat/ChatCompletionDriver.test.ts
+281 -107
View File
@@ -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<IChatProvider['requestPricing']>;
}).requestPricing = () => ({
(
FakeChatProvider.prototype as unknown as {
requestPricing: NonNullable<IChatProvider['requestPricing']>;
}
).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<string, unknown>)
.requestPricing;
delete (
FakeChatProvider.prototype as unknown as Record<string, unknown>
).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<IChatProvider['requestPricing']>;
}).requestPricing = () => ({ extraCost: 1_000_000 });
(
FakeChatProvider.prototype as unknown as {
requestPricing: NonNullable<IChatProvider['requestPricing']>;
}
).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<string, unknown>
).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<IChatProvider['requestPricing']>;
}
).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<string, unknown>)
.requestPricing;
delete (
FakeChatProvider.prototype as unknown as Record<string, unknown>
).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<string, unknown>
| undefined;
Record<string, unknown> | 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;
}
});
src/backend/drivers/ai-chat/ChatCompletionDriver.ts
+6 -4
View File
@@ -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)),
src/backend/drivers/ai-chat/providers/gemini/GeminiChatProvider.native.test.ts
+268 -59
View File
@@ -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<ICompleteArguments>) => {
});
await (
result as {
init_chat_stream: (p: { chatStream: AIChatStream }) => Promise<void>;
init_chat_stream: (p: {
chatStream: AIChatStream;
}) => Promise<void>;
}
).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<MeteringService['reserveCredits']>;
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<ICompleteArguments>) =>
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(
src/backend/drivers/ai-chat/providers/gemini/GeminiChatProvider.ts
+34 -3
View File
@@ -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
src/backend/drivers/ai-chat/providers/gemini/geminiNative.ts
+15 -3
View File
@@ -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,
src/backend/drivers/ai-chat/providers/gemini/models.ts
+4 -2
View File
@@ -24,12 +24,14 @@ export const GEMINI_NATIVE_MODEL_IDS: ReadonlySet<string> = 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'],
src/backend/drivers/ai-chat/types.ts
+10 -3
View File
@@ -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`).
src/docs/src/rate-limits-and-quotas.md
+2
View File
@@ -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.