fix PUT-1444 (#3515)

src/backend/drivers/ai-chat/ChatCompletionDriver.routing.test.ts

@@ -0,0 +1,193 @@
+/**
+ * Copyright (C) 2024-present Puter Technologies Inc.
+ *
+ * This file is part of Puter.
+ *
+ * Puter is free software: you can redistribute it and/or modify it under the
+ * terms of the GNU Affero General Public License as published by the Free
+ * Software Foundation, either version 3 of the License, or (at your option) any
+ * later version.
+ *
+ * This program is distributed in the hope that it will be useful, but WITHOUT
+ * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
+ * FOR A PARTICULAR PURPOSE. See the GNU Affero General Public License for more
+ * details.
+ *
+ * You should have received a copy of the GNU Affero General Public License
+ * along with this program. If not, see
+ * [https://www.gnu.org/licenses/](https://www.gnu.org/licenses/).
+ */
+
+/**
+ * Which provider actually serves a model that several providers advertise.
+ *
+ * Lives apart from ChatCompletionDriver.test.ts because `vi.mock` is
+ * file-scoped and hoisted — mocking the OpenAI SDK and axios here would
+ * otherwise leak into every test in that file. Both mocks sit at the real
+ * network egress points, so the driver's registration, model-map build and
+ * resolution all run for real.
+ */
+
+import {
+    afterAll,
+    beforeAll,
+    describe,
+    expect,
+    it,
+    vi,
+    type MockInstance,
+} from 'vitest';
+
+import { HttpError } from '../../core/http/HttpError.js';
+import { PuterServer } from '../../server.js';
+import { setupTestServer } from '../../testUtil.js';
+import { kv } from '../../util/kvSingleton.js';
+import { withTestActor } from '../integrationTestUtil.js';
+import { ChatCompletionDriver } from './ChatCompletionDriver.js';
+
+// -- OpenAI SDK mock ------------------------------------------------
+// Gemini reaches Google through `new openai.OpenAI()` (default export) and
+// the gateway through the named one, so both must resolve to the same ctor.
+
+const { createMock } = vi.hoisted(() => ({ createMock: vi.fn() }));
+
+vi.mock('openai', () => {
+    const OpenAICtor = vi.fn().mockImplementation(function (
+        this: Record<string, unknown>,
+    ) {
+        this.chat = { completions: { create: createMock } };
+    });
+    return { OpenAI: OpenAICtor, default: { OpenAI: OpenAICtor } };
+});
+
+// -- axios mock (gateway model catalog) -----------------------------
+
+const { axiosRequestMock } = vi.hoisted(() => ({ axiosRequestMock: vi.fn() }));
+
+vi.mock('axios', () => ({
+    default: { request: axiosRequestMock },
+    request: axiosRequestMock,
+}));
+
+// -- Harness --------------------------------------------------------
+
+let server: PuterServer;
+let driver: ChatCompletionDriver;
+
+const INFRON_KV_KEY = 'infronChat:models';
+
+// Google lists gemini-2.5-flash input at $0.30/MTok. The gateway quotes a
+// floor price across its upstream routes, so it undercuts — which is exactly
+// the condition that used to hand it the traffic.
+const GATEWAY_CATALOG = [
+    {
+        id: 'google/gemini-2.5-flash',
+        display_name: 'Google: Gemini 2.5 Flash',
+        category_type: 'LLM',
+        supported_endpoint_types: ['openai'],
+        context_length: 1_048_576,
+        max_output_tokens: 65_536,
+        min_prompt_price: 0.15,
+        min_completion_price: 1.0,
+    },
+    {
+        // Only the gateway carries this one — no first-party counterpart.
+        id: 'google/gemini-2.5-flash-image-preview',
+        display_name: 'Google: Gemini 2.5 Flash Image Preview',
+        category_type: 'LLM',
+        supported_endpoint_types: ['openai'],
+        context_length: 32_768,
+        max_output_tokens: 8_192,
+        min_prompt_price: 0.3,
+        min_completion_price: 2.5,
+    },
+];
+
+beforeAll(async () => {
+    server = await setupTestServer();
+    kv.del?.(INFRON_KV_KEY);
+    axiosRequestMock.mockResolvedValue({ data: { data: GATEWAY_CATALOG } });
+
+    // Built once, not per-test: `#buildModelMap` mutates the catalogs
+    // providers hand back (lowercasing ids, pushing `puterId` onto the
+    // shared `aliases` array), and GeminiChatProvider returns its
+    // module-level GEMINI_MODELS by reference.
+    driver = new ChatCompletionDriver(
+        {
+            providers: {
+                gemini: { apiKey: 'test-key' },
+                infron: { apiKey: 'test-key' },
+                ollama: { enabled: false },
+            },
+        } as never,
+        server.clients,
+        server.stores,
+        server.services,
+    );
+    driver.onServerStart();
+    // `onServerStart` doesn't await `#buildModelMap`, and the gateway
+    // catalog resolves on a microtask — poll until both providers land.
+    for (let i = 0; i < 200; i++) {
+        const ids = await driver.list();
+        if (ids.some((id) => id.startsWith('infron:'))) break;
+        await new Promise((r) => setTimeout(r, 5));
+    }
+});
+
+afterAll(async () => {
+    await server?.shutdown();
+});
+
+/**
+ * Route a request and report who was tried, in order. Forcing the upstream to
+ * reject is what makes the whole chain observable: the driver records every
+ * attempt on the thrown error, and `attempts[0]` is who it chose first.
+ */
+const attemptsFor = async (model: string) => {
+    createMock.mockRejectedValue(new Error('upstream down'));
+    let caught: HttpError | undefined;
+    try {
+        await withTestActor(() =>
+            driver.complete({
+                model,
+                messages: [{ role: 'user', content: 'hi' }],
+            }),
+        );
+    } catch (e) {
+        caught = e as HttpError;
+    }
+    expect(caught).toBeInstanceOf(HttpError);
+    return (caught as unknown as { fields: { attempts: { model: string; provider: string }[] } })
+        .fields.attempts;
+};
+
+describe('ChatCompletionDriver gemini routing', () => {
+    it('serves gemini models from Google, with the gateway only as fallback', async () => {
+        const attempts = await attemptsFor('gemini-2.5-flash');
+
+        expect(attempts[0]).toMatchObject({
+            provider: 'gemini',
+            model: 'gemini-2.5-flash',
+        });
+        expect(attempts[1]).toMatchObject({
+            provider: 'infron',
+            model: 'infron:google/gemini-2.5-flash',
+        });
+    });
+
+    it('routes the prefixed and puterId forms to Google too', async () => {
+        for (const alias of [
+            'google/gemini-2.5-flash',
+            'google:google/gemini-2.5-flash',
+        ]) {
+            const attempts = await attemptsFor(alias);
+            expect(attempts[0].provider).toBe('gemini');
+        }
+    });
+
+    it('still routes models only the gateway carries to the gateway', async () => {
+        const attempts = await attemptsFor('google/gemini-2.5-flash-image-preview');
+
+        expect(attempts[0]).toMatchObject({ provider: 'infron' });
+    });
+});

src/backend/drivers/ai-chat/ChatCompletionDriver.ts

@@ -61,6 +61,10 @@ import {
     normalize_messages,
     normalize_single_message,
 } from './utils/Messages.js';
+import {
+    AGGREGATOR_PROVIDERS,
+    compareModelPreference,
+} from './utils/modelRouting.js';
 import { AIChatStream } from './utils/Streaming.js';
 
 const MAX_FALLBACKS = 4; // includes first attempt
@@ -1019,16 +1023,9 @@ export class ChatCompletionDriver extends PuterDriver {
     // -- Model map ---------------------------------------------------
 
     async #buildModelMap() {
-        const AGGREGATORS = new Set([
-            'together-ai',
-            'openrouter',
-            'infron',
-            'neuralwatt',
-        ]);
-
         for (const providerName in this.#providers) {
             const provider = this.#providers[providerName];
-            const isAggregator = AGGREGATORS.has(providerName);
+            const isAggregator = AGGREGATOR_PROVIDERS.has(providerName);
 
             for (const model of await provider.models()) {
                 model.id = model.id.trim().toLowerCase();
@@ -1058,8 +1055,11 @@ export class ChatCompletionDriver extends PuterDriver {
                             existing !== this.#modelIdMap[model.id]
                         ) {
                             if (existing.some((m) => m.provider === 'gemini')) {
-                                // Gemini exception — let the aggregator
-                                // entry through.
+                                // Gemini is the one vendor whose resold
+                                // duplicates we keep, so a Google outage has
+                                // somewhere to fall back to. Ranking (see
+                                // `compareModelPreference`) keeps the direct
+                                // provider ahead of them.
                                 continue;
                             }
                             skip = true;
@@ -1098,22 +1098,7 @@ export class ChatCompletionDriver extends PuterDriver {
                     }
                 }
 
-                // Sort: together-ai always last; then cheapest input-cost
-                // first; ties break by shorter id (usually the official
-                // name over a long aggregator-qualified one).
-                this.#modelIdMap[model.id].sort((a, b) => {
-                    const aAgg = a.provider === 'together-ai';
-                    const bAgg = b.provider === 'together-ai';
-                    if (aAgg !== bAgg) return aAgg ? 1 : -1;
-                    const aCost = a.costs[
-                        (a.input_cost_key as string) || 'input_tokens'
-                    ] as number;
-                    const bCost = b.costs[
-                        (b.input_cost_key as string) || 'input_tokens'
-                    ] as number;
-                    if (aCost === bCost) return a.id.length - b.id.length;
-                    return aCost - bCost;
-                });
+                this.#modelIdMap[model.id].sort(compareModelPreference);
             }
         }
     }

src/backend/drivers/ai-chat/utils/modelRouting.test.ts

@@ -0,0 +1,135 @@
+/**
+ * Copyright (C) 2024-present Puter Technologies Inc.
+ *
+ * This file is part of Puter.
+ *
+ * Puter is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU Affero General Public License as published
+ * by the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU Affero General Public License for more details.
+ *
+ * You should have received a copy of the GNU Affero General Public License
+ * along with this program.  If not, see <https://www.gnu.org/licenses/>.
+ */
+
+import { describe, expect, it } from 'vitest';
+import { GEMINI_MODELS } from '../providers/gemini/models.js';
+import type { IChatModel } from '../types.js';
+import { compareModelPreference } from './modelRouting.js';
+
+// `#buildModelMap` mutates the catalogs providers hand back, and
+// `GeminiChatProvider.models()` returns the module-level `GEMINI_MODELS` by
+// reference — clone so these fixtures can't be perturbed by another suite.
+const geminiModel = (id: string, provider = 'gemini'): IChatModel => {
+    const found = GEMINI_MODELS.find((m) => m.id === id);
+    if (!found) throw new Error(`no such gemini model: ${id}`);
+    return { ...structuredClone(found), provider };
+};
+
+// Mirrors how the reseller providers coerce a catalog entry: `<gateway>:` on
+// the id, `input_cost_key: 'prompt'`, and prices as microcents per token.
+const resoldModel = (
+    provider: string,
+    catalogId: string,
+    promptCost: number,
+): IChatModel =>
+    ({
+        id: `${provider}:${catalogId}`,
+        name: `${catalogId} (${provider})`,
+        aliases: [catalogId, catalogId.split('/').slice(1).join('/')],
+        costs_currency: 'usd-cents',
+        input_cost_key: 'prompt',
+        output_cost_key: 'completion',
+        costs: { tokens: 1_000_000, prompt: promptCost, completion: 100 },
+        provider,
+    }) as IChatModel;
+
+const winner = (...models: IChatModel[]) =>
+    [...models].sort(compareModelPreference)[0];
+
+describe('compareModelPreference', () => {
+    it('serves the vendor directly even when a reseller quotes a lower price', () => {
+        // Google lists gemini-2.5-flash input at 30 microcents/token; the
+        // gateway advertises a floor price across its upstream routes and
+        // undercuts it. Price must not decide who serves the request.
+        const direct = geminiModel('gemini-2.5-flash');
+        const resold = resoldModel('infron', 'google/gemini-2.5-flash', 15);
+
+        expect(direct.costs.prompt_tokens).toBeGreaterThan(
+            resold.costs.prompt as number,
+        );
+        expect(winner(resold, direct).provider).toBe('gemini');
+        expect(winner(direct, resold).provider).toBe('gemini');
+    });
+
+    it('ranks every reseller behind the vendor, not just one of them', () => {
+        const direct = geminiModel('gemini-3.1-pro-preview');
+        const bucket = [
+            resoldModel('openrouter', 'google/gemini-3.1-pro-preview', 90),
+            resoldModel('infron', 'google/gemini-3.1-pro-preview', 80),
+            resoldModel('together-ai', 'google/gemini-3.1-pro-preview', 70),
+            resoldModel('neuralwatt', 'google/gemini-3.1-pro-preview', 60),
+            direct,
+        ];
+
+        expect(
+            bucket.sort(compareModelPreference).map((m) => m.provider),
+        ).toEqual([
+            'gemini',
+            'neuralwatt',
+            'infron',
+            'openrouter',
+            'together-ai',
+        ]);
+    });
+
+    it('keeps together-ai strictly behind the other resellers', () => {
+        // A flat direct/reseller split would let these tie and re-order by
+        // price; together-ai stays last regardless of how cheap it quotes.
+        const together = resoldModel('together-ai', 'meta/llama-4', 1);
+        const openrouter = resoldModel('openrouter', 'meta/llama-4', 500);
+
+        expect(winner(together, openrouter).provider).toBe('openrouter');
+    });
+
+    it('still orders two direct providers by cheapest input cost', () => {
+        const cheap = geminiModel('gemini-2.5-flash-lite');
+        const pricey = geminiModel('gemini-2.5-pro');
+
+        expect(cheap.costs.prompt_tokens).toBeLessThan(
+            pricey.costs.prompt_tokens as number,
+        );
+        expect(winner(pricey, cheap).id).toBe('gemini-2.5-flash-lite');
+    });
+
+    it('breaks price ties on the shorter id', () => {
+        const short = { ...geminiModel('gemini-2.5-flash'), id: 'gemini-x' };
+        const long = {
+            ...geminiModel('gemini-2.5-flash'),
+            id: 'some-vendor/gemini-x-2025-preview',
+            provider: 'azure-openai',
+        };
+
+        expect(winner(long, short).id).toBe('gemini-x');
+    });
+
+    it('leaves a reseller serving models no vendor provider carries', () => {
+        // The image-preview models are absent from GEMINI_MODELS, so the
+        // gateway is the only route and must stay the winner.
+        expect(
+            GEMINI_MODELS.some((m) => m.id === 'gemini-2.5-flash-image-preview'),
+        ).toBe(false);
+
+        const onlyRoute = resoldModel(
+            'infron',
+            'google/gemini-2.5-flash-image-preview',
+            20,
+        );
+        expect(winner(onlyRoute).provider).toBe('infron');
+    });
+});

src/backend/drivers/ai-chat/utils/modelRouting.ts

@@ -0,0 +1,67 @@
+/**
+ * Copyright (C) 2024-present Puter Technologies Inc.
+ *
+ * This file is part of Puter.
+ *
+ * Puter is free software: you can redistribute it and/or modify it under the
+ * terms of the GNU Affero General Public License as published by the Free
+ * Software Foundation, either version 3 of the License, or (at your option) any
+ * later version.
+ *
+ * This program is distributed in the hope that it will be useful, but WITHOUT
+ * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
+ * FOR A PARTICULAR PURPOSE. See the GNU Affero General Public License for more
+ * details.
+ *
+ * You should have received a copy of the GNU Affero General Public License
+ * along with this program. If not, see
+ * [https://www.gnu.org/licenses/](https://www.gnu.org/licenses/).
+ */
+
+import type { IChatModel } from '../types.js';
+
+/**
+ * Providers that resell other vendors' models rather than serving their own.
+ * Their catalogs duplicate models we already reach directly, so their entries
+ * are kept only as fallback routes.
+ */
+export const AGGREGATOR_PROVIDERS = new Set([
+    'together-ai',
+    'openrouter',
+    'infron',
+    'neuralwatt',
+]);
+
+// Lower rank is served first. `together-ai` sits behind the other resellers —
+// a pre-existing guarantee this ranking extends rather than replaces.
+const providerRank = (provider?: string): number => {
+    if (provider === 'together-ai') return 2;
+    if (provider && AGGREGATOR_PROVIDERS.has(provider)) return 1;
+    return 0;
+};
+
+/**
+ * Orders the candidates that share a model bucket; the first one gets served.
+ *
+ * Direct vendors outrank resellers regardless of quoted price. Resellers
+ * advertise a floor price across their upstream routes, so on price alone they
+ * undercut the vendor's list price and capture traffic for models we hold a
+ * direct integration for. Within a rank, cheapest input cost wins and ties
+ * break by shorter id — usually the official name over a qualified one.
+ */
+export const compareModelPreference = (
+    a: IChatModel,
+    b: IChatModel,
+): number => {
+    const rankDiff = providerRank(a.provider) - providerRank(b.provider);
+    if (rankDiff !== 0) return rankDiff;
+
+    const aCost = a.costs[
+        (a.input_cost_key as string) || 'input_tokens'
+    ] as number;
+    const bCost = b.costs[
+        (b.input_cost_key as string) || 'input_tokens'
+    ] as number;
+    if (aCost === bCost) return a.id.length - b.id.length;
+    return aCost - bCost;
+};