Merge pull request #3987 from HeyPuter/fix/put-1874-chat-test-mode

Fix chat test mode routing to fake provider

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));

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.

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({

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;

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(