From 0b1a4283f542e91940edbc3df0aa94c66f5e2783 Mon Sep 17 00:00:00 2001 From: 404oops Date: Tue, 29 Sep 2026 12:17:33 +0200 Subject: [PATCH] fix(ai): honor chat test mode --- .../ChatCompletionDriver.edges.test.ts | 19 +++++++++ .../ai-chat/ChatCompletionDriver.test.ts | 40 +++++++++++++++++++ .../drivers/ai-chat/ChatCompletionDriver.ts | 26 +++++++++--- src/backend/drivers/ai-chat/types.ts | 1 + src/puter-js/tests/api/suites/ai.suite.ts | 10 +++++ 5 files changed, 90 insertions(+), 6 deletions(-) diff --git a/src/backend/drivers/ai-chat/ChatCompletionDriver.edges.test.ts b/src/backend/drivers/ai-chat/ChatCompletionDriver.edges.test.ts index faa0cc269..2894a632f 100644 --- a/src/backend/drivers/ai-chat/ChatCompletionDriver.edges.test.ts +++ b/src/backend/drivers/ai-chat/ChatCompletionDriver.edges.test.ts @@ -167,6 +167,25 @@ const errorFor = async ( // -- Provider registration ------------------------------------------- describe('ChatCompletionDriver provider registration', () => { + it('routes test mode to fake-chat without calling the selected upstream provider', async () => { + const upstream = vi.spyOn(AzureChatProvider.prototype, 'complete'); + const fake = vi.spyOn(FakeChatProvider.prototype, 'complete'); + + const result = await withTestActor(() => + fullDriver.complete({ + model: 'gpt-4o', + provider: 'azure-openai', + messages: [{ role: 'user', content: 'hi' }], + test_mode: true, + }), + ); + + expect(upstream).not.toHaveBeenCalled(); + expect(fake).toHaveBeenCalledOnce(); + expect(fake.mock.calls[0]![0].model).toBe('fake'); + expect('message' in result && result.message.model).toBe('fake'); + }); + it('registers a model surface spanning every credentialed provider', async () => { const models = await fullDriver.models(); const providers = new Set(models.map((m) => m.provider)); diff --git a/src/backend/drivers/ai-chat/ChatCompletionDriver.test.ts b/src/backend/drivers/ai-chat/ChatCompletionDriver.test.ts index 4cd4fbeef..94765e3f8 100644 --- a/src/backend/drivers/ai-chat/ChatCompletionDriver.test.ts +++ b/src/backend/drivers/ai-chat/ChatCompletionDriver.test.ts @@ -180,6 +180,46 @@ describe('ChatCompletionDriver model catalog', () => { // ── Auth + model resolution ───────────────────────────────────────── 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, 'withAiCostFactor'); + const complete = vi.spyOn(FakeChatProvider.prototype, 'complete'); + + const result = await withTestActor(() => + driver.complete({ + model: 'costly', + messages: [{ role: 'user', content: 'hi' }], + test_mode: true, + }), + ); + + expect(creditGate).not.toHaveBeenCalled(); + expect(complete).toHaveBeenCalledOnce(); + expect(complete.mock.calls[0]![0].model).toBe('fake'); + expect('message' in result && result.message.model).toBe('fake'); + expect('usage' in result && result.usage).toMatchObject({ + input_tokens: 0, + output_tokens: 1, + }); + }); + + it('streams a fake completion in test mode', async () => { + const complete = vi.spyOn(FakeChatProvider.prototype, 'complete'); + const result = await withTestActor(() => + driver.complete({ + model: 'costly', + messages: [{ role: 'user', content: 'hi' }], + stream: true, + test_mode: true, + }), + ); + + expect(complete.mock.calls[0]![0].model).toBe('fake'); + 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' })); + }); + it('throws 401 when no actor is in context', async () => { // Note: not wrapped in `withTestActor` — `Context.get('actor')` // returns undefined. diff --git a/src/backend/drivers/ai-chat/ChatCompletionDriver.ts b/src/backend/drivers/ai-chat/ChatCompletionDriver.ts index eaeccd8cf..168b2fdcd 100644 --- a/src/backend/drivers/ai-chat/ChatCompletionDriver.ts +++ b/src/backend/drivers/ai-chat/ChatCompletionDriver.ts @@ -493,8 +493,20 @@ export class ChatCompletionDriver extends PuterDriver { } } - // Skipped for blocked requests since fake-chat is free and the user - // shouldn't see a billing error in place of the abuse page. + if (args.test_mode === true && !blocked) { + const fakeModel = this.#resolveModel('fake', 'fake-chat'); + if (!fakeModel) { + throw new HttpError(500, 'Test chat provider unavailable', { + legacyCode: 'internal_error', + }); + } + model = fakeModel; + } + + const useFakeProvider = blocked || args.test_mode === true; + + // Fake responses need no credit hold, including blocked prompts and + // test requests. // // The gate hands back a hold on what this attempt could cost, which // stands in for its usage until the real numbers land. It is released @@ -502,7 +514,7 @@ export class ChatCompletionDriver extends PuterDriver { // where "done" is the stream draining rather than this method // returning. let hold: CreditHold = NO_CREDIT_HOLD; - if (!blocked) { + if (!useFakeProvider) { hold = await this.#applyCreditGate(actor, model, args, { promptTokenEstimate, requestedMaxTokens, @@ -537,7 +549,9 @@ export class ChatCompletionDriver extends PuterDriver { }; try { - if (!blocked) await this.#resolvePuterPaths(provider, args, actor); + if (!useFakeProvider) { + await this.#resolvePuterPaths(provider, args, actor); + } res = await provider.complete({ ...args, model: model.id, @@ -576,7 +590,7 @@ export class ChatCompletionDriver extends PuterDriver { // be capped against what is actually left. // The previous attempt released its hold when it failed, so // this one starts from nothing held. - if (!blocked) { + if (!useFakeProvider) { hold = await this.#applyCreditGate(actor, fallback, args, { promptTokenEstimate, requestedMaxTokens, @@ -586,7 +600,7 @@ export class ChatCompletionDriver extends PuterDriver { tried.add(routeId(fallback.provider!, fallback.id)); try { - if (!blocked) { + if (!useFakeProvider) { await this.#resolvePuterPaths(fbProvider, args, actor); } res = await fbProvider.complete({ diff --git a/src/backend/drivers/ai-chat/types.ts b/src/backend/drivers/ai-chat/types.ts index 5a50b31b2..0d8bf5c3f 100644 --- a/src/backend/drivers/ai-chat/types.ts +++ b/src/backend/drivers/ai-chat/types.ts @@ -75,6 +75,7 @@ export interface ICompleteArguments { provider?: string; stream?: boolean; model: string; + test_mode?: boolean; tools?: unknown[]; tool_choice?: unknown; parallel_tool_calls?: boolean; diff --git a/src/puter-js/tests/api/suites/ai.suite.ts b/src/puter-js/tests/api/suites/ai.suite.ts index ae4024061..a6ef3a2a0 100644 --- a/src/puter-js/tests/api/suites/ai.suite.ts +++ b/src/puter-js/tests/api/suites/ai.suite.ts @@ -105,6 +105,16 @@ export default suite('ai', { t.assert.ok(textOf(result).length > 0, 'message should contain text'); }, + 'chat testMode returns a sample without using the selected model': async (t) => { + useApiToken(t); + const result = await t.puter.ai.chat('Hello there', { + model: 'costly', + }, true); + t.assert.equal((result.message as { model?: string }).model, 'fake'); + t.assert.equal((result.usage as Usage).input_tokens, 0); + t.assert.ok(textOf(result).length > 0); + }, + 'chat accepts a messages array': async (t) => { useApiToken(t); const result = await t.puter.ai.chat(