|
| 1 | +import assert from 'node:assert/strict'; |
| 2 | +import { createRequire } from 'node:module'; |
| 3 | +import test from 'node:test'; |
| 4 | + |
| 5 | +import { getAiGatewayRewrites } from '../../apps/web/ai-gateway-rewrites.mjs'; |
| 6 | + |
| 7 | +const require = createRequire(new URL('../../apps/web/next.config.mjs', import.meta.url)); |
| 8 | +const { getPathMatch } = require('next/dist/shared/lib/router/utils/path-match'); |
| 9 | +const { prepareDestination } = require('next/dist/shared/lib/router/utils/prepare-destination'); |
| 10 | + |
| 11 | +function rewrite(url: string, env: NodeJS.ProcessEnv) { |
| 12 | + const request = new URL(url, 'http://localhost'); |
| 13 | + for (const rule of getAiGatewayRewrites(env)) { |
| 14 | + const params = getPathMatch(rule.source)(request.pathname); |
| 15 | + if (params) { |
| 16 | + return prepareDestination({ |
| 17 | + destination: rule.destination, |
| 18 | + params, |
| 19 | + query: Object.fromEntries(request.searchParams), |
| 20 | + appendParamsToQuery: true, |
| 21 | + }).parsedDestination; |
| 22 | + } |
| 23 | + } |
| 24 | + return undefined; |
| 25 | +} |
| 26 | + |
| 27 | +const routes = [ |
| 28 | + ...['gateway', 'openrouter'].flatMap(prefix => [ |
| 29 | + ...['chat/completions', 'responses', 'messages', 'audio/transcriptions', 'transcription-models'] |
| 30 | + .flatMap(endpoint => [endpoint, `v1/${endpoint}`]) |
| 31 | + .map(endpoint => [`/api/${prefix}/${endpoint}`, `/api/v1/${endpoint.replace(/^v1\//, '')}`]), |
| 32 | + ...['models-by-provider', 'providers', 'embeddings'].map(endpoint => [ |
| 33 | + `/api/${prefix}/${endpoint}`, |
| 34 | + `/api/v1/${endpoint}`, |
| 35 | + ]), |
| 36 | + ...['models', 'v1/models'].map(endpoint => [ |
| 37 | + `/api/${prefix}/${endpoint}/anthropic/claude-sonnet-4/endpoints`, |
| 38 | + '/api/v1/models/anthropic/claude-sonnet-4/endpoints', |
| 39 | + ]), |
| 40 | + [`/api/${prefix}/models`, '/api/v1/models'], |
| 41 | + ]), |
| 42 | + ['/api/gateway/v1/models', '/api/v1/models'], |
| 43 | + ['/api/gateway/embedding-models', '/api/v1/embedding-models'], |
| 44 | + ['/api/openrouter/models/validate', '/api/v1/models/validate'], |
| 45 | + ['/api/gateway/typesafe/v1/systemone', '/api/v1/systemone'], |
| 46 | + ['/api/fim/completions', '/api/v1/fim/completions'], |
| 47 | + ['/api/edit/completions', '/api/v1/edit/completions'], |
| 48 | + ['/api/organizations/org-123/models', '/api/v1/organizations/org-123/models'], |
| 49 | + ['/api/organizations/org-123/models/validate', '/api/v1/organizations/org-123/models/validate'], |
| 50 | +]; |
| 51 | + |
| 52 | +for (const globalBackend of ['true', 'false']) { |
| 53 | + void test(`rewrites every legacy gateway route in production with GLOBAL_KILO_BACKEND=${globalBackend}`, () => { |
| 54 | + const env = { |
| 55 | + NODE_ENV: 'production', |
| 56 | + VERCEL_ENV: 'production', |
| 57 | + GLOBAL_KILO_BACKEND: globalBackend, |
| 58 | + AI_GATEWAY_PORT: '9999', |
| 59 | + KILO_PORT_OFFSET: '2500', |
| 60 | + }; |
| 61 | + for (const [source, destination] of routes) { |
| 62 | + const result = rewrite(`${source}?source=extension&stream=true`, env); |
| 63 | + assert.equal(result?.protocol, 'https:', source); |
| 64 | + assert.equal(result?.hostname, 'ai-gateway.kilo.ai', source); |
| 65 | + assert.equal(result?.pathname, destination, source); |
| 66 | + assert.deepEqual(result?.query, { source: 'extension', stream: 'true' }, source); |
| 67 | + } |
| 68 | + }); |
| 69 | +} |
| 70 | + |
| 71 | +void test('rewrites to the local gateway with default, offset, and explicit ports', () => { |
| 72 | + for (const [env, port] of [ |
| 73 | + [{}, '3010'], |
| 74 | + [{ KILO_PORT_OFFSET: '2500' }, '5510'], |
| 75 | + [{ KILO_PORT_OFFSET: '2500', AI_GATEWAY_PORT: '6000' }, '6000'], |
| 76 | + ] as const) { |
| 77 | + for (const [source, destination] of routes) { |
| 78 | + const result = rewrite(source, { |
| 79 | + ...env, |
| 80 | + NODE_ENV: 'development', |
| 81 | + VERCEL_ENV: 'production', |
| 82 | + }); |
| 83 | + assert.equal(result?.protocol, 'http:', source); |
| 84 | + assert.equal(result?.hostname, 'localhost', source); |
| 85 | + assert.equal(result?.port, port, source); |
| 86 | + assert.equal(result?.pathname, destination, source); |
| 87 | + } |
| 88 | + } |
| 89 | +}); |
| 90 | + |
| 91 | +void test('leaves preview, test, and non-Vercel production handlers on the web app', () => { |
| 92 | + for (const env of [ |
| 93 | + { NODE_ENV: 'production', VERCEL_ENV: 'preview' }, |
| 94 | + { NODE_ENV: 'production', VERCEL_ENV: 'development' }, |
| 95 | + { NODE_ENV: 'production' }, |
| 96 | + { NODE_ENV: 'test' }, |
| 97 | + ]) { |
| 98 | + assert.deepEqual(getAiGatewayRewrites(env), []); |
| 99 | + } |
| 100 | +}); |
| 101 | + |
| 102 | +void test('does not move unrelated organization, Exa, marketplace, or cron routes to the gateway', () => { |
| 103 | + for (const source of [ |
| 104 | + '/api/organizations', |
| 105 | + '/api/organizations/org-123', |
| 106 | + '/api/organizations/org-123/defaults', |
| 107 | + '/api/organizations/org-123/modes', |
| 108 | + '/api/organizations/org-123/tokens', |
| 109 | + '/api/exa/search', |
| 110 | + '/api/marketplace/mcps', |
| 111 | + '/api/cron/sync-model-stats', |
| 112 | + ]) { |
| 113 | + assert.equal(rewrite(source, { NODE_ENV: 'production', VERCEL_ENV: 'production' }), undefined); |
| 114 | + } |
| 115 | +}); |
| 116 | + |
| 117 | +void test('preserves encoded model paths and trailing slashes', () => { |
| 118 | + const result = rewrite('/api/gateway/v1/models/provider/model%2Fvariant/endpoints/', { |
| 119 | + NODE_ENV: 'production', |
| 120 | + VERCEL_ENV: 'production', |
| 121 | + }); |
| 122 | + assert.equal(result?.pathname, '/api/v1/models/provider/model%2Fvariant/endpoints'); |
| 123 | +}); |
0 commit comments