From ce6358099a6b657db7e997d8c2ddc6f429abfc39 Mon Sep 17 00:00:00 2001 From: Matt Rubens <2600+mrubens@users.noreply.github.com> Date: Thu, 16 Jul 2026 16:39:35 -0400 Subject: [PATCH 1/7] Add inference gateway proxying model provider APIs for task sandboxes Adds /api/inference/:provider/* to the API server. Task sandboxes authenticate with their run-scoped token; the gateway resolves the deployment's provider API key server-side (runtime env first, then persisted deployment env vars), injects it upstream, and streams the response back. Upstream paths are allow-listed per provider so run tokens can only reach inference endpoints. Covers OpenRouter, Anthropic, OpenAI, and Google Gemini. Providers that need request signing (Vertex) or client-managed OAuth refresh (ChatGPT subscriptions) are not proxied yet. --- apps/api/src/handlers/index.ts | 3 + .../__tests__/inference-gateway.test.ts | 290 ++++++++++++++++++ apps/api/src/handlers/inference/index.ts | 252 +++++++++++++++ apps/api/src/handlers/inference/registry.ts | 100 ++++++ apps/api/src/handlers/mcp/proxy-utils.ts | 2 +- apps/api/src/route-policies.ts | 9 + apps/api/src/server.ts | 2 + packages/db/src/lib/model-runtime-config.ts | 26 ++ 8 files changed, 683 insertions(+), 1 deletion(-) create mode 100644 apps/api/src/handlers/inference/__tests__/inference-gateway.test.ts create mode 100644 apps/api/src/handlers/inference/index.ts create mode 100644 apps/api/src/handlers/inference/registry.ts diff --git a/apps/api/src/handlers/index.ts b/apps/api/src/handlers/index.ts index 3c72a3a56..33ac2be76 100644 --- a/apps/api/src/handlers/index.ts +++ b/apps/api/src/handlers/index.ts @@ -22,6 +22,9 @@ export { trpc } from './trpc'; export { mcp } from './mcp'; export { mcpRouting } from './mcp/routing'; +// inference gateway +export { inference } from './inference'; + // task runs export { taskRunsRouter } from './task-runs'; diff --git a/apps/api/src/handlers/inference/__tests__/inference-gateway.test.ts b/apps/api/src/handlers/inference/__tests__/inference-gateway.test.ts new file mode 100644 index 000000000..b9b383e76 --- /dev/null +++ b/apps/api/src/handlers/inference/__tests__/inference-gateway.test.ts @@ -0,0 +1,290 @@ +import { Hono } from 'hono'; +import type { AuthTokenContext, RunTokenContext } from '@roomote/types'; + +import type { Variables } from '../../../types'; + +const { mockFindTaskRun, mockResolveModelProviderEnvValue } = vi.hoisted( + () => ({ + mockFindTaskRun: vi.fn(), + mockResolveModelProviderEnvValue: vi.fn(), + }), +); + +vi.mock('@roomote/db/server', () => ({ + db: { + query: { + taskRuns: { findFirst: mockFindTaskRun }, + }, + }, + taskRuns: { id: 'id' }, + eq: vi.fn((column: unknown, value: unknown) => ({ column, value })), + resolveModelProviderEnvValue: mockResolveModelProviderEnvValue, +})); + +import { inference } from '../index'; + +function createApp(authContext: Variables['authContext']) { + const app = new Hono<{ Variables: Variables }>(); + + app.use('*', async (c, next) => { + c.set('authContext', authContext); + await next(); + }); + + app.route('/api/inference', inference); + return app; +} + +function createRunToken(overrides?: Partial): RunTokenContext { + return { + runId: 42, + userId: null, + principal: 'deployment', + tokenType: 'run', + version: 1, + ...(overrides ?? {}), + }; +} + +function createUserToken(): AuthTokenContext { + return { + userId: 'user-1', + tokenType: 'auth', + } as AuthTokenContext; +} + +function stubUpstreamFetch(response?: Response) { + const fetchMock = vi.fn().mockResolvedValue( + response ?? + new Response(JSON.stringify({ id: 'msg_1' }), { + status: 200, + headers: { 'content-type': 'application/json' }, + }), + ); + vi.stubGlobal('fetch', fetchMock); + return fetchMock; +} + +async function postMessages( + app: Hono<{ Variables: Variables }>, + path = '/api/inference/anthropic/v1/messages', + headers: Record = {}, +) { + return app.request(path, { + method: 'POST', + headers: { + 'content-type': 'application/json', + authorization: 'Bearer run-token-value', + ...headers, + }, + body: JSON.stringify({ model: 'claude-sonnet-5', max_tokens: 16 }), + }); +} + +describe('inference gateway', () => { + beforeEach(() => { + vi.clearAllMocks(); + vi.unstubAllGlobals(); + mockFindTaskRun.mockResolvedValue({ id: 42 }); + mockResolveModelProviderEnvValue.mockResolvedValue('provider-secret-key'); + }); + + it('rejects user auth tokens', async () => { + const fetchMock = stubUpstreamFetch(); + const response = await postMessages(createApp(createUserToken())); + + expect(response.status).toBe(403); + expect(fetchMock).not.toHaveBeenCalled(); + }); + + it('rejects requests with no auth context', async () => { + const fetchMock = stubUpstreamFetch(); + const response = await postMessages(createApp(undefined)); + + expect(response.status).toBe(403); + expect(fetchMock).not.toHaveBeenCalled(); + }); + + it('rejects run tokens whose task run no longer exists', async () => { + const fetchMock = stubUpstreamFetch(); + mockFindTaskRun.mockResolvedValue(undefined); + + const response = await postMessages(createApp(createRunToken())); + + expect(response.status).toBe(404); + expect(fetchMock).not.toHaveBeenCalled(); + }); + + it('rejects unknown providers', async () => { + const fetchMock = stubUpstreamFetch(); + const response = await postMessages( + createApp(createRunToken()), + '/api/inference/not-a-provider/v1/messages', + ); + + expect(response.status).toBe(404); + expect(fetchMock).not.toHaveBeenCalled(); + }); + + it('rejects upstream paths outside the provider allowlist', async () => { + const fetchMock = stubUpstreamFetch(); + const response = await postMessages( + createApp(createRunToken()), + '/api/inference/anthropic/v1/organizations/members', + ); + + expect(response.status).toBe(403); + expect(fetchMock).not.toHaveBeenCalled(); + }); + + it('rejects sibling paths that share an allowed prefix without a separator', async () => { + const fetchMock = stubUpstreamFetch(); + const response = await postMessages( + createApp(createRunToken()), + '/api/inference/anthropic/v1/messages-admin', + ); + + expect(response.status).toBe(403); + expect(fetchMock).not.toHaveBeenCalled(); + }); + + it('returns 404 when the provider key is not configured', async () => { + const fetchMock = stubUpstreamFetch(); + mockResolveModelProviderEnvValue.mockResolvedValue(undefined); + + const response = await postMessages(createApp(createRunToken())); + + expect(response.status).toBe(404); + expect(fetchMock).not.toHaveBeenCalled(); + }); + + it('forwards to the upstream with the provider key and strips the run token', async () => { + const fetchMock = stubUpstreamFetch(); + const response = await postMessages( + createApp(createRunToken()), + undefined, + { + 'anthropic-version': '2023-06-01', + }, + ); + + expect(response.status).toBe(200); + expect(mockResolveModelProviderEnvValue).toHaveBeenCalledWith( + 'ANTHROPIC_API_KEY', + ); + expect(fetchMock).toHaveBeenCalledTimes(1); + + const [url, init] = fetchMock.mock.calls[0] as [string, RequestInit]; + expect(url).toBe('https://api.anthropic.com/v1/messages'); + + const headers = new Headers(init.headers); + expect(headers.get('x-api-key')).toBe('provider-secret-key'); + expect(headers.get('authorization')).toBeNull(); + expect(headers.get('anthropic-version')).toBe('2023-06-01'); + + const forwardedBody = await new Response(init.body).text(); + expect(JSON.parse(forwardedBody)).toEqual({ + model: 'claude-sonnet-5', + max_tokens: 16, + }); + }); + + it('sends Bearer-prefixed keys for bearer-auth providers', async () => { + const fetchMock = stubUpstreamFetch(); + const response = await postMessages( + createApp(createRunToken()), + '/api/inference/openrouter/v1/chat/completions', + ); + + expect(response.status).toBe(200); + expect(mockResolveModelProviderEnvValue).toHaveBeenCalledWith( + 'OPENROUTER_API_KEY', + ); + + const [url, init] = fetchMock.mock.calls[0] as [string, RequestInit]; + expect(url).toBe('https://openrouter.ai/api/v1/chat/completions'); + + const headers = new Headers(init.headers); + expect(headers.get('authorization')).toBe('Bearer provider-secret-key'); + }); + + it('preserves the query string on upstream requests', async () => { + const fetchMock = stubUpstreamFetch(); + const app = createApp(createRunToken()); + + const response = await app.request( + '/api/inference/google/v1beta/models/gemini-2.5-pro:streamGenerateContent?alt=sse', + { + method: 'POST', + headers: { + 'content-type': 'application/json', + authorization: 'Bearer run-token-value', + }, + body: JSON.stringify({}), + }, + ); + + expect(response.status).toBe(200); + + const [url, init] = fetchMock.mock.calls[0] as [string, RequestInit]; + expect(url).toBe( + 'https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-pro:streamGenerateContent?alt=sse', + ); + + const headers = new Headers(init.headers); + expect(headers.get('x-goog-api-key')).toBe('provider-secret-key'); + }); + + it('streams the upstream response body and status through', async () => { + const upstreamBody = 'event: message_start\ndata: {}\n\n'; + stubUpstreamFetch( + new Response(upstreamBody, { + status: 200, + headers: { 'content-type': 'text/event-stream' }, + }), + ); + + const response = await postMessages(createApp(createRunToken())); + + expect(response.status).toBe(200); + expect(response.headers.get('content-type')).toBe('text/event-stream'); + expect(await response.text()).toBe(upstreamBody); + }); + + it('passes upstream error statuses through', async () => { + stubUpstreamFetch( + new Response(JSON.stringify({ error: { type: 'overloaded_error' } }), { + status: 529, + headers: { 'content-type': 'application/json' }, + }), + ); + + const response = await postMessages(createApp(createRunToken())); + + expect(response.status).toBe(529); + }); + + it('returns 502 when the upstream fetch fails', async () => { + const fetchMock = vi + .fn() + .mockRejectedValue(new Error('connect ECONNREFUSED')); + vi.stubGlobal('fetch', fetchMock); + + const response = await postMessages(createApp(createRunToken())); + + expect(response.status).toBe(502); + }); + + it('rejects unsupported methods', async () => { + const fetchMock = stubUpstreamFetch(); + const app = createApp(createRunToken()); + + const response = await app.request('/api/inference/anthropic/v1/messages', { + method: 'DELETE', + headers: { authorization: 'Bearer run-token-value' }, + }); + + expect(response.status).toBe(405); + expect(fetchMock).not.toHaveBeenCalled(); + }); +}); diff --git a/apps/api/src/handlers/inference/index.ts b/apps/api/src/handlers/inference/index.ts new file mode 100644 index 000000000..9e20b4e3f --- /dev/null +++ b/apps/api/src/handlers/inference/index.ts @@ -0,0 +1,252 @@ +import { Hono } from 'hono'; + +import { formatSingleLineLog } from '@roomote/types'; +import { + db, + eq, + resolveModelProviderEnvValue, + taskRuns, +} from '@roomote/db/server'; + +import type { Variables } from '../../types'; +import { fetchWithLongLivedStreamDispatcher } from '../long-lived-fetch'; +import { createLoggedProxyResponseBody } from '../proxy-response-stream'; +import { + buildProxyResponseHeaders, + isRunTokenContext, +} from '../mcp/proxy-utils'; +import { + formatProviderAuthHeaderValue, + getInferenceProvider, + isInferencePathAllowed, +} from './registry'; + +/** + * Client headers never forwarded upstream. The client's own credential + * headers (its bearer token is the Roomote run token) are replaced with the + * provider key; hop-by-hop and proxy-added routing headers are dropped so the + * upstream sees a clean direct request. + */ +const REQUEST_HEADER_DENYLIST = new Set([ + 'authorization', + 'x-api-key', + 'x-goog-api-key', + 'cookie', + 'host', + 'connection', + 'content-length', + 'keep-alive', + 'proxy-authenticate', + 'proxy-authorization', + 'te', + 'trailer', + 'transfer-encoding', + 'upgrade', + 'expect', + 'forwarded', + 'x-forwarded-for', + 'x-forwarded-host', + 'x-forwarded-proto', + 'x-real-ip', +]); + +function buildUpstreamRequestHeaders( + requestHeaders: Headers, + authHeaderName: string, + authHeaderValue: string, +): Headers { + const headers = new Headers(); + + for (const [key, value] of requestHeaders.entries()) { + if (!REQUEST_HEADER_DENYLIST.has(key.toLowerCase())) { + headers.set(key, value); + } + } + + headers.set(authHeaderName, authHeaderValue); + + return headers; +} + +/** + * The upstream path is everything after the provider segment. The router + * only matches `/:provider/*`, so the marker is always present. + */ +function extractUpstreamPath(pathname: string, providerId: string): string { + const marker = `/${providerId}/`; + const markerIndex = pathname.indexOf(marker); + + return markerIndex === -1 + ? '' + : pathname.slice(markerIndex + marker.length - 1); +} + +/** + * Inference gateway: forwards LLM API traffic from task sandboxes to model + * providers, injecting the deployment's provider key server-side. Sandboxes + * authenticate with their run-scoped token and never hold the key itself. + * + * Requests stream through in both directions; inference responses are + * long-lived SSE streams, so upstream fetches use the no-body-timeout + * dispatcher and rely on caller disconnects for cleanup. + */ +export const inference = new Hono<{ Variables: Variables }>(); + +inference.on(['POST', 'GET'], '/:provider/*', async (c) => { + const startedAt = Date.now(); + const requestId = crypto.randomUUID(); + const method = c.req.method; + const providerId = c.req.param('provider'); + const pathname = new URL(c.req.url).pathname; + const logPrefix = `[Inference Gateway:${providerId}]`; + + const auth = c.get('authContext'); + + if (!auth || !isRunTokenContext(auth)) { + return c.json( + { error: 'The inference gateway requires a task run token' }, + 403, + ); + } + + const taskRun = await db.query.taskRuns.findFirst({ + columns: { id: true }, + where: eq(taskRuns.id, auth.runId), + }); + + if (!taskRun) { + return c.json({ error: 'Task run not found for this token' }, 404); + } + + const provider = getInferenceProvider(providerId); + + if (!provider) { + return c.json({ error: `Unknown inference provider "${providerId}"` }, 404); + } + + const upstreamPath = extractUpstreamPath(pathname, providerId); + + if (!isInferencePathAllowed(provider, upstreamPath)) { + console.warn( + formatSingleLineLog(`${logPrefix} Rejected disallowed upstream path`, { + requestId, + method, + runId: auth.runId, + upstreamPath, + }), + ); + + return c.json( + { + error: `Path is not allowed through the ${provider.name} inference gateway`, + }, + 403, + ); + } + + let apiKey: string | undefined; + + try { + apiKey = await resolveModelProviderEnvValue(provider.envVarName); + } catch (error) { + console.error( + formatSingleLineLog(`${logPrefix} Failed to resolve provider key`, { + requestId, + method, + runId: auth.runId, + error: error instanceof Error ? error.message : String(error), + }), + ); + + return c.json( + { error: `Failed to resolve the ${provider.name} API key` }, + 500, + ); + } + + if (!apiKey) { + return c.json( + { + error: `No ${provider.name} API key is configured for this deployment`, + }, + 404, + ); + } + + const search = new URL(c.req.url).search; + const upstreamUrl = `${provider.upstreamBaseUrl}${upstreamPath}${search}`; + + try { + const upstreamResponse = await fetchWithLongLivedStreamDispatcher( + upstreamUrl, + { + method, + headers: buildUpstreamRequestHeaders( + c.req.raw.headers, + provider.authHeader.name, + formatProviderAuthHeaderValue(provider, apiKey), + ), + body: c.req.raw.body, + signal: c.req.raw.signal, + // Required by undici when streaming a request body. + ...(c.req.raw.body ? { duplex: 'half' as const } : {}), + }, + ); + + if (!upstreamResponse.ok) { + console.warn( + formatSingleLineLog(`${logPrefix} Upstream returned non-OK status`, { + requestId, + method, + runId: auth.runId, + upstreamPath, + status: upstreamResponse.status, + statusText: upstreamResponse.statusText, + elapsedMs: Date.now() - startedAt, + }), + ); + } + + return new Response( + createLoggedProxyResponseBody({ + body: upstreamResponse.body, + logPrefix: `${logPrefix} Upstream response stream failed`, + getLogFields: () => ({ + requestId, + method, + runId: auth.runId, + upstreamPath, + status: upstreamResponse.status, + elapsedMs: Date.now() - startedAt, + }), + trackingContext: { + route: `inference:${providerId}`, + method, + path: pathname, + requestId, + }, + }), + { + status: upstreamResponse.status, + headers: buildProxyResponseHeaders(upstreamResponse.headers), + }, + ); + } catch (error) { + console.error( + formatSingleLineLog(`${logPrefix} Upstream fetch failed`, { + requestId, + method, + runId: auth.runId, + upstreamPath, + elapsedMs: Date.now() - startedAt, + error: error instanceof Error ? error.message : String(error), + }), + ); + + return c.json({ error: `Failed to reach the ${provider.name} API` }, 502); + } +}); + +inference.all('/:provider/*', (c) => + c.json({ error: 'Method not allowed' }, 405), +); diff --git a/apps/api/src/handlers/inference/registry.ts b/apps/api/src/handlers/inference/registry.ts new file mode 100644 index 000000000..61e130a9f --- /dev/null +++ b/apps/api/src/handlers/inference/registry.ts @@ -0,0 +1,100 @@ +import type { SetupModelProviderId } from '@roomote/types'; + +/** + * How an upstream provider expects its API key: the header to set and whether + * the value is Bearer-prefixed or sent raw. + */ +interface InferenceProviderAuthHeader { + name: string; + scheme?: 'bearer'; +} + +interface InferenceProviderDefinition { + id: SetupModelProviderId; + name: string; + upstreamBaseUrl: string; + /** Deployment env var holding the provider API key. */ + envVarName: string; + authHeader: InferenceProviderAuthHeader; + /** + * Upstream path prefixes the gateway forwards. Everything else is rejected + * so a run token can only reach inference endpoints, never the provider's + * account, billing, or admin surface. + */ + allowedPathPrefixes: readonly string[]; +} + +/** + * Providers reachable through the inference gateway. Key-authenticated + * HTTP-proxyable providers only: providers that need request signing + * (Vertex service accounts) or client-managed OAuth refresh (ChatGPT + * subscriptions) are handled separately. + */ +const INFERENCE_PROVIDERS: readonly InferenceProviderDefinition[] = [ + { + id: 'openrouter', + name: 'OpenRouter', + upstreamBaseUrl: 'https://openrouter.ai/api', + envVarName: 'OPENROUTER_API_KEY', + authHeader: { name: 'authorization', scheme: 'bearer' }, + allowedPathPrefixes: ['/v1'], + }, + { + id: 'anthropic', + name: 'Anthropic', + upstreamBaseUrl: 'https://api.anthropic.com', + envVarName: 'ANTHROPIC_API_KEY', + authHeader: { name: 'x-api-key' }, + allowedPathPrefixes: ['/v1/messages', '/v1/models'], + }, + { + id: 'openai', + name: 'OpenAI', + upstreamBaseUrl: 'https://api.openai.com', + envVarName: 'OPENAI_API_KEY', + authHeader: { name: 'authorization', scheme: 'bearer' }, + allowedPathPrefixes: [ + '/v1/chat/completions', + '/v1/responses', + '/v1/embeddings', + '/v1/models', + ], + }, + { + id: 'google', + name: 'Google Gemini', + upstreamBaseUrl: 'https://generativelanguage.googleapis.com', + envVarName: 'GEMINI_API_KEY', + authHeader: { name: 'x-goog-api-key' }, + allowedPathPrefixes: ['/v1beta/models', '/v1/models'], + }, +]; + +export function getInferenceProvider( + providerId: string, +): InferenceProviderDefinition | undefined { + return INFERENCE_PROVIDERS.find((provider) => provider.id === providerId); +} + +/** + * A path is allowed when it equals an allowed prefix or nests under it with a + * `/` separator, so `/v1/messages-admin` cannot ride on a `/v1/messages` + * allowance while `/v1beta/models/gemini-2.5-pro:streamGenerateContent` + * passes under `/v1beta/models`. + */ +export function isInferencePathAllowed( + provider: InferenceProviderDefinition, + upstreamPath: string, +): boolean { + return provider.allowedPathPrefixes.some( + (prefix) => + upstreamPath === prefix || upstreamPath.startsWith(`${prefix}/`), + ); +} + +export function formatProviderAuthHeaderValue( + provider: InferenceProviderDefinition, + apiKey: string, +): string { + return provider.authHeader.scheme === 'bearer' ? `Bearer ${apiKey}` : apiKey; +} diff --git a/apps/api/src/handlers/mcp/proxy-utils.ts b/apps/api/src/handlers/mcp/proxy-utils.ts index 5b728e392..96c4f790b 100644 --- a/apps/api/src/handlers/mcp/proxy-utils.ts +++ b/apps/api/src/handlers/mcp/proxy-utils.ts @@ -243,7 +243,7 @@ function buildProxyRequestHeaders( return headers; } -function buildProxyResponseHeaders(upstreamHeaders: Headers): Headers { +export function buildProxyResponseHeaders(upstreamHeaders: Headers): Headers { const headers = new Headers(); const excludedHeaders = new Set([ diff --git a/apps/api/src/route-policies.ts b/apps/api/src/route-policies.ts index 616b3db31..2f05e80aa 100644 --- a/apps/api/src/route-policies.ts +++ b/apps/api/src/route-policies.ts @@ -214,6 +214,15 @@ export const ROUTE_POLICY_RULES: readonly RoutePolicyRule[] = [ rateLimits: WEBHOOK_RATE_LIMITS, }, + // Inference gateway: task sandboxes call model providers through this + // proxy with their run-scoped token; the provider key is injected + // server-side and never enters the sandbox. + { + name: 'inference', + match: { type: 'prefix', path: '/api/inference' }, + policy: 'task-token', + }, + // Router-facing MCP endpoints: accept user auth tokens (LLM router // gathering context before a run exists) and task run tokens. { diff --git a/apps/api/src/server.ts b/apps/api/src/server.ts index b1fb553cf..0a51a1980 100644 --- a/apps/api/src/server.ts +++ b/apps/api/src/server.ts @@ -45,6 +45,7 @@ import { teams, telegram, discord, + inference, mcp, mcpRouting, taskRunsRouter, @@ -192,6 +193,7 @@ export function createApiApp(): ApiApp { app.route('/api/webhooks/teams', teams); app.route('/api/webhooks/telegram', telegram); app.route('/api/internal/discord', discord); + app.route('/api/inference', inference); app.route('/api/mcp', mcp); app.route('/api/mcp-routing', mcpRouting); app.route('/api/task-runs', taskRunsRouter); diff --git a/packages/db/src/lib/model-runtime-config.ts b/packages/db/src/lib/model-runtime-config.ts index 56d0984bb..27f1930bf 100644 --- a/packages/db/src/lib/model-runtime-config.ts +++ b/packages/db/src/lib/model-runtime-config.ts @@ -123,6 +123,32 @@ function resolveProviderKeyNames({ ]; } +/** + * Resolve a single model-provider env value with the same precedence the task + * runtime uses: the runtime process env first, then the persisted (encrypted) + * deployment environment variables. + */ +export async function resolveModelProviderEnvValue( + envVarName: string, + options: { + runtimeEnv?: Partial>; + executor?: DatabaseOrTransaction; + } = {}, +): Promise { + const runtimeEnv = options.runtimeEnv ?? process.env; + const runtimeValue = normalizeConfiguredValue(runtimeEnv[envVarName]); + + if (runtimeValue) { + return runtimeValue; + } + + const persistedEnvVars = await loadPersistedDeploymentEnvVars( + options.executor ?? db, + ); + + return normalizeConfiguredValue(persistedEnvVars[envVarName]); +} + export async function resolveEffectiveModelRuntimeEnv( options: { runtimeEnv?: Partial>; From 593d95d122939c23af05dc8123f5bae71fada1c0 Mon Sep 17 00:00:00 2001 From: Matt Rubens <2600+mrubens@users.noreply.github.com> Date: Thu, 16 Jul 2026 20:58:11 +0000 Subject: [PATCH 2/7] fix: harden inference gateway provider access --- .../__tests__/inference-gateway.test.ts | 26 ++++++++-- apps/api/src/handlers/inference/index.ts | 2 +- apps/api/src/handlers/inference/registry.ts | 45 +++++++++-------- .../__tests__/model-runtime-config.test.ts | 48 +++++++++++++++++++ packages/db/src/lib/environment-variables.ts | 3 +- packages/db/src/lib/model-runtime-config.ts | 35 ++++++++++---- 6 files changed, 125 insertions(+), 34 deletions(-) create mode 100644 packages/db/src/lib/__tests__/model-runtime-config.test.ts diff --git a/apps/api/src/handlers/inference/__tests__/inference-gateway.test.ts b/apps/api/src/handlers/inference/__tests__/inference-gateway.test.ts index b9b383e76..3c4ea6bba 100644 --- a/apps/api/src/handlers/inference/__tests__/inference-gateway.test.ts +++ b/apps/api/src/handlers/inference/__tests__/inference-gateway.test.ts @@ -148,6 +148,20 @@ describe('inference gateway', () => { expect(fetchMock).not.toHaveBeenCalled(); }); + it('rejects provider account paths nested below former inference prefixes', async () => { + const fetchMock = stubUpstreamFetch(); + const app = createApp(createRunToken()); + + const [anthropicResponse, openRouterResponse] = await Promise.all([ + postMessages(app, '/api/inference/anthropic/v1/messages/batches'), + postMessages(app, '/api/inference/openrouter/v1/key'), + ]); + + expect(anthropicResponse.status).toBe(403); + expect(openRouterResponse.status).toBe(403); + expect(fetchMock).not.toHaveBeenCalled(); + }); + it('returns 404 when the provider key is not configured', async () => { const fetchMock = stubUpstreamFetch(); mockResolveModelProviderEnvValue.mockResolvedValue(undefined); @@ -169,9 +183,9 @@ describe('inference gateway', () => { ); expect(response.status).toBe(200); - expect(mockResolveModelProviderEnvValue).toHaveBeenCalledWith( + expect(mockResolveModelProviderEnvValue).toHaveBeenCalledWith([ 'ANTHROPIC_API_KEY', - ); + ]); expect(fetchMock).toHaveBeenCalledTimes(1); const [url, init] = fetchMock.mock.calls[0] as [string, RequestInit]; @@ -197,9 +211,9 @@ describe('inference gateway', () => { ); expect(response.status).toBe(200); - expect(mockResolveModelProviderEnvValue).toHaveBeenCalledWith( + expect(mockResolveModelProviderEnvValue).toHaveBeenCalledWith([ 'OPENROUTER_API_KEY', - ); + ]); const [url, init] = fetchMock.mock.calls[0] as [string, RequestInit]; expect(url).toBe('https://openrouter.ai/api/v1/chat/completions'); @@ -233,6 +247,10 @@ describe('inference gateway', () => { const headers = new Headers(init.headers); expect(headers.get('x-goog-api-key')).toBe('provider-secret-key'); + expect(mockResolveModelProviderEnvValue).toHaveBeenCalledWith([ + 'GEMINI_API_KEY', + 'GOOGLE_GENERATIVE_AI_API_KEY', + ]); }); it('streams the upstream response body and status through', async () => { diff --git a/apps/api/src/handlers/inference/index.ts b/apps/api/src/handlers/inference/index.ts index 9e20b4e3f..e31f32380 100644 --- a/apps/api/src/handlers/inference/index.ts +++ b/apps/api/src/handlers/inference/index.ts @@ -147,7 +147,7 @@ inference.on(['POST', 'GET'], '/:provider/*', async (c) => { let apiKey: string | undefined; try { - apiKey = await resolveModelProviderEnvValue(provider.envVarName); + apiKey = await resolveModelProviderEnvValue(provider.envVarNames); } catch (error) { console.error( formatSingleLineLog(`${logPrefix} Failed to resolve provider key`, { diff --git a/apps/api/src/handlers/inference/registry.ts b/apps/api/src/handlers/inference/registry.ts index 61e130a9f..bc1b8afd7 100644 --- a/apps/api/src/handlers/inference/registry.ts +++ b/apps/api/src/handlers/inference/registry.ts @@ -13,15 +13,15 @@ interface InferenceProviderDefinition { id: SetupModelProviderId; name: string; upstreamBaseUrl: string; - /** Deployment env var holding the provider API key. */ - envVarName: string; + /** Deployment env vars holding the provider API key, in precedence order. */ + envVarNames: readonly string[]; authHeader: InferenceProviderAuthHeader; /** * Upstream path prefixes the gateway forwards. Everything else is rejected * so a run token can only reach inference endpoints, never the provider's * account, billing, or admin surface. */ - allowedPathPrefixes: readonly string[]; + allowedPaths: readonly string[]; } /** @@ -35,25 +35,30 @@ const INFERENCE_PROVIDERS: readonly InferenceProviderDefinition[] = [ id: 'openrouter', name: 'OpenRouter', upstreamBaseUrl: 'https://openrouter.ai/api', - envVarName: 'OPENROUTER_API_KEY', + envVarNames: ['OPENROUTER_API_KEY'], authHeader: { name: 'authorization', scheme: 'bearer' }, - allowedPathPrefixes: ['/v1'], + allowedPaths: [ + '/v1/chat/completions', + '/v1/completions', + '/v1/embeddings', + '/v1/models', + ], }, { id: 'anthropic', name: 'Anthropic', upstreamBaseUrl: 'https://api.anthropic.com', - envVarName: 'ANTHROPIC_API_KEY', + envVarNames: ['ANTHROPIC_API_KEY'], authHeader: { name: 'x-api-key' }, - allowedPathPrefixes: ['/v1/messages', '/v1/models'], + allowedPaths: ['/v1/messages', '/v1/messages/count_tokens', '/v1/models'], }, { id: 'openai', name: 'OpenAI', upstreamBaseUrl: 'https://api.openai.com', - envVarName: 'OPENAI_API_KEY', + envVarNames: ['OPENAI_API_KEY'], authHeader: { name: 'authorization', scheme: 'bearer' }, - allowedPathPrefixes: [ + allowedPaths: [ '/v1/chat/completions', '/v1/responses', '/v1/embeddings', @@ -64,9 +69,9 @@ const INFERENCE_PROVIDERS: readonly InferenceProviderDefinition[] = [ id: 'google', name: 'Google Gemini', upstreamBaseUrl: 'https://generativelanguage.googleapis.com', - envVarName: 'GEMINI_API_KEY', + envVarNames: ['GEMINI_API_KEY', 'GOOGLE_GENERATIVE_AI_API_KEY'], authHeader: { name: 'x-goog-api-key' }, - allowedPathPrefixes: ['/v1beta/models', '/v1/models'], + allowedPaths: ['/v1beta/models', '/v1/models'], }, ]; @@ -77,19 +82,21 @@ export function getInferenceProvider( } /** - * A path is allowed when it equals an allowed prefix or nests under it with a - * `/` separator, so `/v1/messages-admin` cannot ride on a `/v1/messages` - * allowance while `/v1beta/models/gemini-2.5-pro:streamGenerateContent` - * passes under `/v1beta/models`. + * A path is allowed when it exactly matches an inference endpoint. Google + * model routes are the exception because their model ID and action are part + * of the nested path; their provider API has no account surface below models. */ export function isInferencePathAllowed( provider: InferenceProviderDefinition, upstreamPath: string, ): boolean { - return provider.allowedPathPrefixes.some( - (prefix) => - upstreamPath === prefix || upstreamPath.startsWith(`${prefix}/`), - ); + return provider.allowedPaths.some((path) => { + if (upstreamPath === path) { + return true; + } + + return provider.id === 'google' && upstreamPath.startsWith(`${path}/`); + }); } export function formatProviderAuthHeaderValue( diff --git a/packages/db/src/lib/__tests__/model-runtime-config.test.ts b/packages/db/src/lib/__tests__/model-runtime-config.test.ts new file mode 100644 index 000000000..ea479eac4 --- /dev/null +++ b/packages/db/src/lib/__tests__/model-runtime-config.test.ts @@ -0,0 +1,48 @@ +const { mockResolveDeploymentEnvVar } = vi.hoisted(() => ({ + mockResolveDeploymentEnvVar: vi.fn(), +})); + +vi.mock('../environment-variables', async (importOriginal) => ({ + ...(await importOriginal()), + resolveDeploymentEnvVar: mockResolveDeploymentEnvVar, +})); + +import { resolveModelProviderEnvValue } from '../model-runtime-config'; +import type { DatabaseOrTransaction } from '../../db'; + +describe('resolveModelProviderEnvValue', () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + it('checks all runtime aliases before querying persisted values', async () => { + mockResolveDeploymentEnvVar.mockResolvedValue('persisted-key'); + + const value = await resolveModelProviderEnvValue( + ['GEMINI_API_KEY', 'GOOGLE_GENERATIVE_AI_API_KEY'], + { + runtimeEnv: { GOOGLE_GENERATIVE_AI_API_KEY: 'runtime-alias-key' }, + }, + ); + + expect(value).toBe('runtime-alias-key'); + expect(mockResolveDeploymentEnvVar).not.toHaveBeenCalled(); + }); + + it('looks up only the requested persisted key', async () => { + const executor = {} as DatabaseOrTransaction; + mockResolveDeploymentEnvVar.mockResolvedValue('persisted-key'); + + const value = await resolveModelProviderEnvValue('OPENAI_API_KEY', { + runtimeEnv: {}, + executor, + }); + + expect(value).toBe('persisted-key'); + expect(mockResolveDeploymentEnvVar).toHaveBeenCalledWith( + 'OPENAI_API_KEY', + executor, + {}, + ); + }); +}); diff --git a/packages/db/src/lib/environment-variables.ts b/packages/db/src/lib/environment-variables.ts index 48b433a87..4b27af061 100644 --- a/packages/db/src/lib/environment-variables.ts +++ b/packages/db/src/lib/environment-variables.ts @@ -13,8 +13,9 @@ export function stringifyDecryptedEnvVarValue(value: unknown): string { export async function resolveDeploymentEnvVar( name: string, executor: DatabaseOrTransaction = db, + runtimeEnv: Partial> = process.env, ): Promise { - const processValue = process.env[name]?.trim(); + const processValue = runtimeEnv[name]?.trim(); if (processValue) { return processValue; diff --git a/packages/db/src/lib/model-runtime-config.ts b/packages/db/src/lib/model-runtime-config.ts index 27f1930bf..fa39e7530 100644 --- a/packages/db/src/lib/model-runtime-config.ts +++ b/packages/db/src/lib/model-runtime-config.ts @@ -17,7 +17,10 @@ import { resolveOpenCodeAuthContent } from './chatgpt-subscription'; import { type DatabaseOrTransaction, db } from '../db'; import { deploymentSettings } from '../schema'; -import { stringifyDecryptedEnvVarValue } from './environment-variables'; +import { + resolveDeploymentEnvVar, + stringifyDecryptedEnvVarValue, +} from './environment-variables'; const DEFAULT_DEPLOYMENT_ID = 'default'; @@ -129,24 +132,38 @@ function resolveProviderKeyNames({ * deployment environment variables. */ export async function resolveModelProviderEnvValue( - envVarName: string, + envVarNames: string | readonly string[], options: { runtimeEnv?: Partial>; executor?: DatabaseOrTransaction; } = {}, ): Promise { const runtimeEnv = options.runtimeEnv ?? process.env; - const runtimeValue = normalizeConfiguredValue(runtimeEnv[envVarName]); + const names = typeof envVarNames === 'string' ? [envVarNames] : envVarNames; + + for (const envVarName of names) { + const runtimeValue = normalizeConfiguredValue(runtimeEnv[envVarName]); - if (runtimeValue) { - return runtimeValue; + if (runtimeValue) { + return runtimeValue; + } } - const persistedEnvVars = await loadPersistedDeploymentEnvVars( - options.executor ?? db, - ); + for (const envVarName of names) { + const persistedValue = await resolveDeploymentEnvVar( + envVarName, + options.executor ?? db, + {}, + ); + + const normalizedValue = normalizeConfiguredValue(persistedValue); + + if (normalizedValue) { + return normalizedValue; + } + } - return normalizeConfiguredValue(persistedEnvVars[envVarName]); + return undefined; } export async function resolveEffectiveModelRuntimeEnv( From 9ec8bfbd15bfe0fcb2e3c64b428f6e5bea656d5a Mon Sep 17 00:00:00 2001 From: Matt Rubens Date: Thu, 16 Jul 2026 17:54:27 -0400 Subject: [PATCH 3/7] Route sandbox inference through the gateway behind a feature flag * Route sandbox inference through the gateway behind a deployment flag When R_INFERENCE_GATEWAY is enabled (control-plane env or persisted deployment env var), task sandboxes no longer receive raw provider API keys for gateway-covered providers (OpenRouter, Anthropic, OpenAI, Google Gemini). Instead: - The dequeue/resume env resolver emits R_INFERENCE_GATEWAY_URL and withholds covered keys from the sandbox env (control-plane inference keeps direct keys since it holds no run token). - The worker daemon launch env skips covered operator keys. - The worker's OpenCode config rebases covered providers onto the gateway with the run token as the API key, preserving OpenRouter attribution headers and leaving ChatGPT-subscription OAuth and Bedrock/Vertex untouched. - The API token middleware accepts the run token from x-api-key / x-goog-api-key on the inference gateway surface only, since the Anthropic and Gemini SDKs send their key through those headers. Provider definitions are shared between the gateway and the worker via @roomote/types so path suffixes and env keys cannot drift. * Gate the inference gateway on a feature flag instead of an env var Replaces the R_INFERENCE_GATEWAY env var with an InferenceGateway feature flag stored in deployment metadata (metadata key inference_gateway), so the gateway can be toggled from the admin controls like other deployment-level features. - resolveEffectiveModelRuntimeEnv takes an explicit inferenceGateway option instead of sniffing env vars; the dequeue path evaluates the flag through the shared evaluator and fails closed (direct keys) if flag evaluation is unavailable. - Worker daemon env builders take inferenceGatewayEnabled via BuildWorkerEnvOptions; each controller spawn path evaluates the flag before building the env. - Control-plane inference is unaffected: callers that do not pass the option keep direct provider keys. --------- Co-authored-by: Matt Rubens <2600+mrubens@users.noreply.github.com> --- apps/api/src/handlers/inference/registry.ts | 88 ++---------- .../__tests__/tokenAuthMiddleware.test.ts | 125 ++++++++++++++++ .../api/src/middleware/tokenAuthMiddleware.ts | 36 ++++- .../inference-gateway-flag.ts | 28 ++++ .../compute-providers/spawn-blaxel-worker.ts | 2 + .../compute-providers/spawn-daytona-worker.ts | 2 + .../compute-providers/spawn-docker-worker.ts | 2 + .../src/compute-providers/spawn-e2b-worker.ts | 2 + .../compute-providers/spawn-modal-worker.ts | 2 + apps/worker/src/run-task/agent-home.test.ts | 77 ++++++++++ apps/worker/src/run-task/agent-home.ts | 102 +++++++++++-- apps/worker/src/run-task/env.ts | 1 + .../src/worker-env/__tests__/base.test.ts | 19 +++ .../compute-providers/src/worker-env/base.ts | 18 ++- .../src/worker-env/blaxel.ts | 8 +- .../src/worker-env/daytona.ts | 8 +- .../src/worker-env/docker.ts | 8 +- .../compute-providers/src/worker-env/e2b.ts | 8 +- .../compute-providers/src/worker-env/modal.ts | 8 +- .../compute-providers/src/worker-env/types.ts | 6 + .../db/src/lib/model-runtime-config.test.ts | 62 ++++++++ packages/db/src/lib/model-runtime-config.ts | 33 +++++ packages/feature-flags/src/config.ts | 11 ++ packages/feature-flags/src/types.ts | 1 + .../server/lib/task-runs/dequeue-helpers.ts | 60 +++++++- packages/types/src/index.ts | 1 + packages/types/src/inference-gateway.ts | 136 ++++++++++++++++++ 27 files changed, 748 insertions(+), 106 deletions(-) create mode 100644 apps/api/src/middleware/__tests__/tokenAuthMiddleware.test.ts create mode 100644 apps/controller/src/compute-providers/inference-gateway-flag.ts create mode 100644 packages/types/src/inference-gateway.ts diff --git a/apps/api/src/handlers/inference/registry.ts b/apps/api/src/handlers/inference/registry.ts index bc1b8afd7..2835c9a25 100644 --- a/apps/api/src/handlers/inference/registry.ts +++ b/apps/api/src/handlers/inference/registry.ts @@ -1,84 +1,12 @@ -import type { SetupModelProviderId } from '@roomote/types'; - -/** - * How an upstream provider expects its API key: the header to set and whether - * the value is Bearer-prefixed or sent raw. - */ -interface InferenceProviderAuthHeader { - name: string; - scheme?: 'bearer'; -} - -interface InferenceProviderDefinition { - id: SetupModelProviderId; - name: string; - upstreamBaseUrl: string; - /** Deployment env vars holding the provider API key, in precedence order. */ - envVarNames: readonly string[]; - authHeader: InferenceProviderAuthHeader; - /** - * Upstream path prefixes the gateway forwards. Everything else is rejected - * so a run token can only reach inference endpoints, never the provider's - * account, billing, or admin surface. - */ - allowedPaths: readonly string[]; -} - -/** - * Providers reachable through the inference gateway. Key-authenticated - * HTTP-proxyable providers only: providers that need request signing - * (Vertex service accounts) or client-managed OAuth refresh (ChatGPT - * subscriptions) are handled separately. - */ -const INFERENCE_PROVIDERS: readonly InferenceProviderDefinition[] = [ - { - id: 'openrouter', - name: 'OpenRouter', - upstreamBaseUrl: 'https://openrouter.ai/api', - envVarNames: ['OPENROUTER_API_KEY'], - authHeader: { name: 'authorization', scheme: 'bearer' }, - allowedPaths: [ - '/v1/chat/completions', - '/v1/completions', - '/v1/embeddings', - '/v1/models', - ], - }, - { - id: 'anthropic', - name: 'Anthropic', - upstreamBaseUrl: 'https://api.anthropic.com', - envVarNames: ['ANTHROPIC_API_KEY'], - authHeader: { name: 'x-api-key' }, - allowedPaths: ['/v1/messages', '/v1/messages/count_tokens', '/v1/models'], - }, - { - id: 'openai', - name: 'OpenAI', - upstreamBaseUrl: 'https://api.openai.com', - envVarNames: ['OPENAI_API_KEY'], - authHeader: { name: 'authorization', scheme: 'bearer' }, - allowedPaths: [ - '/v1/chat/completions', - '/v1/responses', - '/v1/embeddings', - '/v1/models', - ], - }, - { - id: 'google', - name: 'Google Gemini', - upstreamBaseUrl: 'https://generativelanguage.googleapis.com', - envVarNames: ['GEMINI_API_KEY', 'GOOGLE_GENERATIVE_AI_API_KEY'], - authHeader: { name: 'x-goog-api-key' }, - allowedPaths: ['/v1beta/models', '/v1/models'], - }, -]; +import { + getInferenceGatewayProvider, + type InferenceGatewayProvider, +} from '@roomote/types'; export function getInferenceProvider( providerId: string, -): InferenceProviderDefinition | undefined { - return INFERENCE_PROVIDERS.find((provider) => provider.id === providerId); +): InferenceGatewayProvider | undefined { + return getInferenceGatewayProvider(providerId); } /** @@ -87,7 +15,7 @@ export function getInferenceProvider( * of the nested path; their provider API has no account surface below models. */ export function isInferencePathAllowed( - provider: InferenceProviderDefinition, + provider: InferenceGatewayProvider, upstreamPath: string, ): boolean { return provider.allowedPaths.some((path) => { @@ -100,7 +28,7 @@ export function isInferencePathAllowed( } export function formatProviderAuthHeaderValue( - provider: InferenceProviderDefinition, + provider: InferenceGatewayProvider, apiKey: string, ): string { return provider.authHeader.scheme === 'bearer' ? `Bearer ${apiKey}` : apiKey; diff --git a/apps/api/src/middleware/__tests__/tokenAuthMiddleware.test.ts b/apps/api/src/middleware/__tests__/tokenAuthMiddleware.test.ts new file mode 100644 index 000000000..a3d2ca221 --- /dev/null +++ b/apps/api/src/middleware/__tests__/tokenAuthMiddleware.test.ts @@ -0,0 +1,125 @@ +import { Hono } from 'hono'; + +import type { Variables } from '../../types'; + +const { mockValidateRunToken, mockValidateAuthToken, mockFindDeployment } = + vi.hoisted(() => ({ + mockValidateRunToken: vi.fn(), + mockValidateAuthToken: vi.fn(), + mockFindDeployment: vi.fn(), + })); + +vi.mock('@roomote/auth', () => ({ + validateRunToken: mockValidateRunToken, + validateAuthToken: mockValidateAuthToken, +})); + +vi.mock('@roomote/db/server', () => ({ + db: { + query: { + deploymentSettings: { findFirst: mockFindDeployment }, + }, + }, + deploymentSettings: { id: 'id' }, + eq: vi.fn(), +})); + +import { tokenAuthMiddleware } from '../tokenAuthMiddleware'; + +const RUN_TOKEN_CONTEXT = { + runId: 42, + userId: null, + principal: 'deployment', + tokenType: 'run', + version: 1, +}; + +function createApp() { + const app = new Hono<{ Variables: Variables }>(); + + app.use('*', tokenAuthMiddleware()); + app.all('*', (c) => c.json({ authContext: c.get('authContext') ?? null })); + + return app; +} + +async function requestAuthContext( + path: string, + headers: Record, +): Promise { + const response = await createApp().request(path, { headers }); + const body = (await response.json()) as { authContext: unknown }; + + return body.authContext; +} + +describe('tokenAuthMiddleware token extraction', () => { + beforeEach(() => { + vi.clearAllMocks(); + mockFindDeployment.mockResolvedValue({ metadata: null }); + mockValidateRunToken.mockImplementation(async (token: string) => { + if (token !== 'valid-run-token') { + throw new Error('invalid token'); + } + + return RUN_TOKEN_CONTEXT; + }); + mockValidateAuthToken.mockRejectedValue(new Error('invalid token')); + }); + + it('accepts a bearer token on any path', async () => { + const authContext = await requestAuthContext('/api/task-runs/1', { + authorization: 'Bearer valid-run-token', + }); + + expect(authContext).toEqual(RUN_TOKEN_CONTEXT); + }); + + it('accepts the run token from x-api-key on the inference gateway', async () => { + const authContext = await requestAuthContext( + '/api/inference/anthropic/v1/messages', + { 'x-api-key': 'valid-run-token' }, + ); + + expect(authContext).toEqual(RUN_TOKEN_CONTEXT); + }); + + it('accepts the run token from x-goog-api-key on the inference gateway', async () => { + const authContext = await requestAuthContext( + '/api/inference/google/v1beta/models/gemini-2.5-pro:generateContent', + { 'x-goog-api-key': 'valid-run-token' }, + ); + + expect(authContext).toEqual(RUN_TOKEN_CONTEXT); + }); + + it('prefers the Authorization header over provider key headers', async () => { + const authContext = await requestAuthContext( + '/api/inference/anthropic/v1/messages', + { + authorization: 'Bearer valid-run-token', + 'x-api-key': 'some-other-value', + }, + ); + + expect(authContext).toEqual(RUN_TOKEN_CONTEXT); + }); + + it('ignores provider key headers outside the inference gateway', async () => { + const authContext = await requestAuthContext('/api/task-runs/1', { + 'x-api-key': 'valid-run-token', + }); + + expect(authContext).toBeNull(); + expect(mockValidateRunToken).not.toHaveBeenCalled(); + }); + + it('does not authenticate an invalid token from provider key headers', async () => { + const authContext = await requestAuthContext( + '/api/inference/anthropic/v1/messages', + { 'x-api-key': 'not-a-token' }, + ); + + expect(authContext).toBeNull(); + }); +}); diff --git a/apps/api/src/middleware/tokenAuthMiddleware.ts b/apps/api/src/middleware/tokenAuthMiddleware.ts index 84daa3394..b1fb8ddbb 100644 --- a/apps/api/src/middleware/tokenAuthMiddleware.ts +++ b/apps/api/src/middleware/tokenAuthMiddleware.ts @@ -26,13 +26,41 @@ async function deploymentAllowsTokenAuth(): Promise { return !isRoomoteDeploymentDisabled(deployment?.metadata); } +const INFERENCE_GATEWAY_PATH_PREFIX = '/api/inference'; + +/** + * Provider SDKs pointed at the inference gateway send their "API key" (the + * run token) through provider-specific headers: `x-api-key` for Anthropic, + * `x-goog-api-key` for Gemini. Accept the token from those headers on the + * inference gateway surface only; everywhere else the Authorization bearer + * header remains the single token transport. + */ +function extractBearerToken( + c: Context<{ Variables: Variables }>, +): string | undefined { + const authHeader = c.req.header('Authorization'); + + if (authHeader?.startsWith('Bearer ')) { + return authHeader.substring(7); + } + + const path = c.req.path; + + if ( + path === INFERENCE_GATEWAY_PATH_PREFIX || + path.startsWith(`${INFERENCE_GATEWAY_PATH_PREFIX}/`) + ) { + return c.req.header('x-api-key') ?? c.req.header('x-goog-api-key'); + } + + return undefined; +} + export const tokenAuthMiddleware = () => createMiddleware(async (c: Context<{ Variables: Variables }>, next: Next) => { - const authHeader = c.req.header('Authorization'); - - if (authHeader?.startsWith('Bearer ')) { - const token = authHeader.substring(7); + const token = extractBearerToken(c); + if (token) { // Try run token first (has more specific claims) let isRunToken = false; diff --git a/apps/controller/src/compute-providers/inference-gateway-flag.ts b/apps/controller/src/compute-providers/inference-gateway-flag.ts new file mode 100644 index 000000000..6775d09ce --- /dev/null +++ b/apps/controller/src/compute-providers/inference-gateway-flag.ts @@ -0,0 +1,28 @@ +import { + FeatureFlag, + getFeatureFlagEvaluator, +} from '@roomote/feature-flags/server'; +import { getRedis } from '@roomote/redis'; + +/** + * Evaluate the deployment-level InferenceGateway flag for worker spawn env + * construction. Fails closed (provider keys ship to the worker daemon as + * before) when flag evaluation is unavailable, so a Redis outage cannot + * break task inference. + */ +export async function isInferenceGatewayEnabledForWorkerEnv(): Promise { + try { + return await getFeatureFlagEvaluator(getRedis()).evaluate( + FeatureFlag.InferenceGateway, + { isDeploymentContext: true }, + ); + } catch (error) { + console.warn( + `[spawn-worker] InferenceGateway flag evaluation failed; forwarding provider keys: ${ + error instanceof Error ? error.message : String(error) + }`, + ); + + return false; + } +} diff --git a/apps/controller/src/compute-providers/spawn-blaxel-worker.ts b/apps/controller/src/compute-providers/spawn-blaxel-worker.ts index 1efeaf209..f40d7ac2e 100644 --- a/apps/controller/src/compute-providers/spawn-blaxel-worker.ts +++ b/apps/controller/src/compute-providers/spawn-blaxel-worker.ts @@ -19,6 +19,7 @@ import { resolveAuthBypassHeaderName, resolveAuthBypassValue, } from '@roomote/compute-providers'; +import { isInferenceGatewayEnabledForWorkerEnv } from './inference-gateway-flag'; import { primeEnvironmentOidcForMachine } from '../sandbox-oidc'; import { @@ -185,6 +186,7 @@ export async function spawnBlaxelWorker( args, env: buildBlaxelWorkerEnv({ authToken, + inferenceGatewayEnabled: await isInferenceGatewayEnabledForWorkerEnv(), sandboxExpiresAtMs: Date.now() + config.blaxelTimeoutMs, deploymentSlug: config.deploymentSlug, environmentId, diff --git a/apps/controller/src/compute-providers/spawn-daytona-worker.ts b/apps/controller/src/compute-providers/spawn-daytona-worker.ts index 055e020aa..8c2c5ea70 100644 --- a/apps/controller/src/compute-providers/spawn-daytona-worker.ts +++ b/apps/controller/src/compute-providers/spawn-daytona-worker.ts @@ -20,6 +20,7 @@ import { resolveAuthBypassHeaderName, resolveAuthBypassValue, } from '@roomote/compute-providers'; +import { isInferenceGatewayEnabledForWorkerEnv } from './inference-gateway-flag'; import { primeEnvironmentOidcForMachine } from '../sandbox-oidc'; import { @@ -330,6 +331,7 @@ export async function spawnDaytonaWorker( args, env: buildDaytonaWorkerEnv({ authToken, + inferenceGatewayEnabled: await isInferenceGatewayEnabledForWorkerEnv(), sandboxExpiresAtMs: Date.now() + daytonaTimeoutMs, deploymentSlug, environmentId, diff --git a/apps/controller/src/compute-providers/spawn-docker-worker.ts b/apps/controller/src/compute-providers/spawn-docker-worker.ts index 2ea6d2c83..a44c6b0a5 100644 --- a/apps/controller/src/compute-providers/spawn-docker-worker.ts +++ b/apps/controller/src/compute-providers/spawn-docker-worker.ts @@ -23,6 +23,7 @@ import { resolveAuthBypassHeaderName, resolveAuthBypassValue, } from '@roomote/compute-providers'; +import { isInferenceGatewayEnabledForWorkerEnv } from './inference-gateway-flag'; import { getNamedPortsForTaskRun, @@ -365,6 +366,7 @@ export async function spawnDockerWorker( const workerEnv = buildDockerWorkerEnv({ authToken, + inferenceGatewayEnabled: await isInferenceGatewayEnabledForWorkerEnv(), sandboxExpiresAtMs: Date.now() + config.dockerTimeoutMs, deploymentSlug: config.deploymentSlug, environmentId: taskRun.payload.environmentId, diff --git a/apps/controller/src/compute-providers/spawn-e2b-worker.ts b/apps/controller/src/compute-providers/spawn-e2b-worker.ts index 65fc9c7bd..893e157b3 100644 --- a/apps/controller/src/compute-providers/spawn-e2b-worker.ts +++ b/apps/controller/src/compute-providers/spawn-e2b-worker.ts @@ -20,6 +20,7 @@ import { resolveAuthBypassHeaderName, resolveAuthBypassValue, } from '@roomote/compute-providers'; +import { isInferenceGatewayEnabledForWorkerEnv } from './inference-gateway-flag'; import { primeEnvironmentOidcForMachine } from '../sandbox-oidc'; import { @@ -326,6 +327,7 @@ export async function spawnE2bWorker( args, env: buildE2bWorkerEnv({ authToken, + inferenceGatewayEnabled: await isInferenceGatewayEnabledForWorkerEnv(), sandboxExpiresAtMs: Date.now() + e2bTimeoutMs, deploymentSlug, environmentId, diff --git a/apps/controller/src/compute-providers/spawn-modal-worker.ts b/apps/controller/src/compute-providers/spawn-modal-worker.ts index 17b64e3e5..47069b5e5 100644 --- a/apps/controller/src/compute-providers/spawn-modal-worker.ts +++ b/apps/controller/src/compute-providers/spawn-modal-worker.ts @@ -22,6 +22,7 @@ import { resolveAuthBypassHeaderName, resolveAuthBypassValue, } from '@roomote/compute-providers'; +import { isInferenceGatewayEnabledForWorkerEnv } from './inference-gateway-flag'; import { primeEnvironmentOidcForMachine } from '../sandbox-oidc'; import { @@ -409,6 +410,7 @@ export async function spawnModalWorker( args, env: buildModalWorkerEnv({ authToken, + inferenceGatewayEnabled: await isInferenceGatewayEnabledForWorkerEnv(), sandboxExpiresAtMs: Date.now() + modalTimeoutMs, deploymentSlug, environmentId, diff --git a/apps/worker/src/run-task/agent-home.test.ts b/apps/worker/src/run-task/agent-home.test.ts index c436443c8..9b6583fd7 100644 --- a/apps/worker/src/run-task/agent-home.test.ts +++ b/apps/worker/src/run-task/agent-home.test.ts @@ -125,6 +125,83 @@ describe('generateOpenCodeConfig provider support', () => { ).toThrow('AWS_REGION must be a valid AWS region'); }); + it('rebases gateway-covered providers onto the inference gateway', () => { + const result = generateOpenCodeConfig({ + homeDir: createHomeDir(), + runtimeEnv: { + R_MODEL: 'anthropic/claude-sonnet-5', + R_SMALL_MODEL: 'google/gemini-3.5-flash', + R_INFERENCE_GATEWAY_URL: 'https://api.example.com/api/inference', + }, + }); + const config = JSON.parse(result.configContent) as { + provider: Record; + }; + + expect(config.provider.anthropic).toMatchObject({ + options: { + baseURL: 'https://api.example.com/api/inference/anthropic/v1', + apiKey: '{env:ROOMOTE_CLOUD_TOKEN}', + }, + }); + expect(config.provider.google).toMatchObject({ + options: { + baseURL: 'https://api.example.com/api/inference/google/v1beta', + apiKey: '{env:ROOMOTE_CLOUD_TOKEN}', + }, + }); + }); + + it('keeps OpenRouter attribution headers when rebasing onto the gateway', () => { + const result = generateOpenCodeConfig({ + homeDir: createHomeDir(), + runtimeEnv: { + R_MODEL: 'openrouter/anthropic/claude-sonnet-5', + R_INFERENCE_GATEWAY_URL: 'https://api.example.com/api/inference', + }, + }); + const config = JSON.parse(result.configContent) as { + provider: { + openrouter: { options: Record }; + }; + }; + + expect(config.provider.openrouter.options).toMatchObject({ + baseURL: 'https://api.example.com/api/inference/openrouter/v1', + apiKey: '{env:ROOMOTE_CLOUD_TOKEN}', + headers: expect.objectContaining({}) as Record, + }); + }); + + it('leaves the openai provider alone when a ChatGPT subscription is present', () => { + const result = generateOpenCodeConfig({ + homeDir: createHomeDir(), + runtimeEnv: { + R_MODEL: 'openai/gpt-5.4-codex', + R_INFERENCE_GATEWAY_URL: 'https://api.example.com/api/inference', + OPENCODE_AUTH_CONTENT: '{"openai":{"type":"oauth"}}', + }, + }); + const config = JSON.parse(result.configContent) as { + provider: Record; + }; + + expect(config.provider.openai).toBeUndefined(); + }); + + it('does not touch provider base URLs when no gateway URL is present', () => { + const result = generateOpenCodeConfig({ + homeDir: createHomeDir(), + runtimeEnv: { + R_MODEL: 'anthropic/claude-sonnet-5', + ANTHROPIC_API_KEY: 'sk-anthropic', + }, + }); + + expect(result.configContent).not.toContain('baseURL'); + expect(result.configContent).not.toContain('ROOMOTE_CLOUD_TOKEN'); + }); + it('materializes inline Google Vertex credentials before OpenCode starts', () => { const homeDir = createHomeDir(); const credentialsJson = JSON.stringify({ diff --git a/apps/worker/src/run-task/agent-home.ts b/apps/worker/src/run-task/agent-home.ts index 09de2f359..a4b5e713b 100644 --- a/apps/worker/src/run-task/agent-home.ts +++ b/apps/worker/src/run-task/agent-home.ts @@ -2,14 +2,19 @@ import * as fs from 'node:fs'; import * as path from 'node:path'; import { + buildInferenceGatewayOpenCodeBaseUrl, buildOpenCodeModelReasoningOptions, + CHATGPT_OPENCODE_PROVIDER_ID, collectOpenRouterVariantModelAlias, + getInferenceGatewayProvider, GOOGLE_APPLICATION_CREDENTIALS_ENV_VAR_NAME, getMcpIntegration, + INFERENCE_GATEWAY_URL_ENV_VAR_NAME, isInlineGoogleCredentialsValue, mergeOpenCodeModelReasoningOptions, mergeOpenRouterVariantAliasModels, normalizeOptionalReasoningEffort, + OPENCODE_AUTH_CONTENT_ENV_VAR_NAME, renderManualSkillMarkdown, resolveOpenRouterVariantModelAlias, OPENCODE_ARCHITECT_AGENT, @@ -618,6 +623,73 @@ function mergeBedrockMantleProviderConfig( }; } +/** + * When the dequeue env carries an inference gateway URL, rebase each + * gateway-covered provider that a selected model uses onto the gateway. The + * SDK authenticates with the run token (already in the harness env), which + * the gateway exchanges for the deployment's provider key server-side, so + * the raw key never enters the sandbox. + */ +function mergeInferenceGatewayProviderConfig( + providerConfig: Record, + runtimeEnv: Record, + modelIds: Array, +): Record { + const gatewayUrl = runtimeEnv[INFERENCE_GATEWAY_URL_ENV_VAR_NAME]?.trim(); + + if (!gatewayUrl) { + return providerConfig; + } + + const providerIds = new Set( + modelIds.flatMap((modelId) => { + const providerId = modelId?.trim().split('/')[0]; + + return providerId ? [providerId] : []; + }), + ); + + let merged = providerConfig; + + for (const providerId of providerIds) { + const gatewayProvider = getInferenceGatewayProvider(providerId); + + if (!gatewayProvider) { + continue; + } + + // ChatGPT-subscription OAuth authenticates through opencode's Codex + // plugin directly; leave the openai provider on its default base URL + // when a subscription record is present. + if ( + providerId === CHATGPT_OPENCODE_PROVIDER_ID && + runtimeEnv[OPENCODE_AUTH_CONTENT_ENV_VAR_NAME] + ) { + continue; + } + + const existingProvider = asRecord(merged[providerId]); + const existingOptions = asRecord(existingProvider.options); + + merged = { + ...merged, + [providerId]: { + ...existingProvider, + options: { + ...existingOptions, + baseURL: buildInferenceGatewayOpenCodeBaseUrl( + gatewayUrl, + gatewayProvider, + ), + apiKey: '{env:ROOMOTE_CLOUD_TOKEN}', + }, + }, + }; + } + + return merged; +} + function normalizeStringList(value: unknown): string[] { if (Array.isArray(value)) { return value.filter((entry): entry is string => typeof entry === 'string'); @@ -1132,18 +1204,26 @@ function resolveModelBackedOpenCodeConfig( ); } - const providerConfig = mergeBedrockMantleProviderConfig( - mergeOpenRouterVariantAliasModels(providerReasoningConfig, variantAliases), + const configuredModelIds = [ + effectiveCodingModel, + model, + smallModel, + visionModel, + codeReviewModel, + exploreModel, + planningModel, + ]; + const providerConfig = mergeInferenceGatewayProviderConfig( + mergeBedrockMantleProviderConfig( + mergeOpenRouterVariantAliasModels( + providerReasoningConfig, + variantAliases, + ), + runtimeEnv, + configuredModelIds, + ), runtimeEnv, - [ - effectiveCodingModel, - model, - smallModel, - visionModel, - codeReviewModel, - exploreModel, - planningModel, - ], + configuredModelIds, ); return { diff --git a/apps/worker/src/run-task/env.ts b/apps/worker/src/run-task/env.ts index 8c14aa17e..689060e79 100644 --- a/apps/worker/src/run-task/env.ts +++ b/apps/worker/src/run-task/env.ts @@ -19,6 +19,7 @@ const BLOCKED_HARNESS_ENV_KEYS = new Set([ 'PREVIEW_PROXY_SUBDOMAIN_SUFFIX', ]); const MODEL_RUNTIME_ENV_KEYS = [ + 'R_INFERENCE_GATEWAY_URL', 'R_MODEL', 'R_SMALL_MODEL', 'R_VISION_MODEL', diff --git a/packages/compute-providers/src/worker-env/__tests__/base.test.ts b/packages/compute-providers/src/worker-env/__tests__/base.test.ts index 3ea70d4af..ce670be9f 100644 --- a/packages/compute-providers/src/worker-env/__tests__/base.test.ts +++ b/packages/compute-providers/src/worker-env/__tests__/base.test.ts @@ -142,4 +142,23 @@ describe('buildBaseWorkerEnv', () => { expect(env.R_MODEL).toBe('vercel/openai/gpt-5.4'); expect(env.AI_GATEWAY_API_KEY).toBe('vercel-key'); }); + + it('holds back gateway-covered provider keys when the inference gateway is enabled', () => { + process.env.R_MODEL = 'anthropic/claude-sonnet-5'; + process.env.ANTHROPIC_API_KEY = 'anthropic-key'; + process.env.OPENROUTER_API_KEY = 'openrouter-key'; + process.env.AWS_BEARER_TOKEN_BEDROCK = 'bedrock-key'; + + const env = buildBaseWorkerEnv({ + authToken: 'auth-token', + extraEnv: {}, + inferenceGatewayEnabled: true, + }); + + expect(env.R_MODEL).toBe('anthropic/claude-sonnet-5'); + expect(env.ANTHROPIC_API_KEY).toBeUndefined(); + expect(env.OPENROUTER_API_KEY).toBeUndefined(); + // The gateway cannot serve Bedrock yet, so its key still ships. + expect(env.AWS_BEARER_TOKEN_BEDROCK).toBe('bedrock-key'); + }); }); diff --git a/packages/compute-providers/src/worker-env/base.ts b/packages/compute-providers/src/worker-env/base.ts index a79a564d7..da42b0582 100644 --- a/packages/compute-providers/src/worker-env/base.ts +++ b/packages/compute-providers/src/worker-env/base.ts @@ -1,6 +1,7 @@ import { Env } from '@roomote/env'; import { DEFAULT_MODEL_PROVIDER_ENV_KEYS, + INFERENCE_GATEWAY_PROVIDER_ENV_VAR_NAMES, parseModelProviderEnvKeys, } from '@roomote/types'; @@ -28,7 +29,9 @@ function getOperatorModelProviderEnvKeys(): string[] { return [...new Set([...DEFAULT_MODEL_PROVIDER_ENV_KEYS, ...configured])]; } -function buildOperatorModelProviderEnv(): Record { +function buildOperatorModelProviderEnv( + inferenceGatewayEnabled: boolean, +): Record { const env: Record = {}; const model = process.env.R_MODEL?.trim(); const smallModel = process.env.R_SMALL_MODEL?.trim(); @@ -80,11 +83,21 @@ function buildOperatorModelProviderEnv(): Record { env.R_MODEL_ENV_KEYS = process.env.R_MODEL_ENV_KEYS; } + // When the inference gateway is enabled, gateway-covered provider keys + // stay on the control plane: the per-task dequeue env routes the harness + // through the gateway instead. Keys the gateway cannot serve yet (Bedrock, + // Vertex) still ship with the worker daemon env. + const gatewayCoveredKeys = new Set(INFERENCE_GATEWAY_PROVIDER_ENV_VAR_NAMES); + for (const key of getOperatorModelProviderEnvKeys()) { if (BLOCKED_WORKER_ENV_KEYS.has(key)) { continue; } + if (inferenceGatewayEnabled && gatewayCoveredKeys.has(key)) { + continue; + } + const value = process.env[key]; if (value !== undefined) { @@ -99,6 +112,7 @@ export function buildBaseWorkerEnv({ authToken, sandboxExpiresAtMs, extraEnv, + inferenceGatewayEnabled = false, }: BuildWorkerEnvOptions): Record { const previewProxyBaseUrl = process.env.PREVIEW_PROXY_BASE_URL; @@ -147,7 +161,7 @@ export function buildBaseWorkerEnv({ ...(process.env.PREVIEW_AUTH_COOKIE_NAME && { PREVIEW_AUTH_COOKIE_NAME: process.env.PREVIEW_AUTH_COOKIE_NAME, }), - ...buildOperatorModelProviderEnv(), + ...buildOperatorModelProviderEnv(inferenceGatewayEnabled), ...filterWorkerExtraEnv(extraEnv), }; } diff --git a/packages/compute-providers/src/worker-env/blaxel.ts b/packages/compute-providers/src/worker-env/blaxel.ts index 370bae37f..0874dc7c7 100644 --- a/packages/compute-providers/src/worker-env/blaxel.ts +++ b/packages/compute-providers/src/worker-env/blaxel.ts @@ -6,6 +6,7 @@ export function buildBlaxelWorkerEnv({ authToken, sandboxExpiresAtMs, extraEnv, + inferenceGatewayEnabled, deploymentSlug, environmentId, image, @@ -15,7 +16,12 @@ export function buildBlaxelWorkerEnv({ image: string; }): Record { return { - ...buildBaseWorkerEnv({ authToken, sandboxExpiresAtMs, extraEnv }), + ...buildBaseWorkerEnv({ + authToken, + sandboxExpiresAtMs, + extraEnv, + inferenceGatewayEnabled, + }), ...buildWorkerContextEnv({ provider: 'blaxel', fingerprint: image, diff --git a/packages/compute-providers/src/worker-env/daytona.ts b/packages/compute-providers/src/worker-env/daytona.ts index df84778c8..b4ad2377c 100644 --- a/packages/compute-providers/src/worker-env/daytona.ts +++ b/packages/compute-providers/src/worker-env/daytona.ts @@ -7,6 +7,7 @@ export function buildDaytonaWorkerEnv({ authToken, sandboxExpiresAtMs, extraEnv, + inferenceGatewayEnabled, deploymentSlug, environmentId, snapshotName, @@ -16,7 +17,12 @@ export function buildDaytonaWorkerEnv({ snapshotName: string; }): Record { return { - ...buildBaseWorkerEnv({ authToken, sandboxExpiresAtMs, extraEnv }), + ...buildBaseWorkerEnv({ + authToken, + sandboxExpiresAtMs, + extraEnv, + inferenceGatewayEnabled, + }), ...buildWorkerContextEnv({ provider: 'daytona', fingerprint: snapshotName, diff --git a/packages/compute-providers/src/worker-env/docker.ts b/packages/compute-providers/src/worker-env/docker.ts index cb8fc2ca4..6f05fcc77 100644 --- a/packages/compute-providers/src/worker-env/docker.ts +++ b/packages/compute-providers/src/worker-env/docker.ts @@ -6,6 +6,7 @@ export function buildDockerWorkerEnv({ authToken, sandboxExpiresAtMs, extraEnv, + inferenceGatewayEnabled, deploymentSlug, environmentId, image, @@ -15,7 +16,12 @@ export function buildDockerWorkerEnv({ image?: string; }): Record { return { - ...buildBaseWorkerEnv({ authToken, sandboxExpiresAtMs, extraEnv }), + ...buildBaseWorkerEnv({ + authToken, + sandboxExpiresAtMs, + extraEnv, + inferenceGatewayEnabled, + }), ...buildWorkerContextEnv({ provider: 'docker', fingerprint: image, diff --git a/packages/compute-providers/src/worker-env/e2b.ts b/packages/compute-providers/src/worker-env/e2b.ts index c57911b97..e96b1a1d3 100644 --- a/packages/compute-providers/src/worker-env/e2b.ts +++ b/packages/compute-providers/src/worker-env/e2b.ts @@ -7,6 +7,7 @@ export function buildE2bWorkerEnv({ authToken, sandboxExpiresAtMs, extraEnv, + inferenceGatewayEnabled, deploymentSlug, environmentId, templateId, @@ -16,7 +17,12 @@ export function buildE2bWorkerEnv({ templateId: string; }): Record { return { - ...buildBaseWorkerEnv({ authToken, sandboxExpiresAtMs, extraEnv }), + ...buildBaseWorkerEnv({ + authToken, + sandboxExpiresAtMs, + extraEnv, + inferenceGatewayEnabled, + }), ...buildWorkerContextEnv({ provider: 'e2b', fingerprint: templateId, diff --git a/packages/compute-providers/src/worker-env/modal.ts b/packages/compute-providers/src/worker-env/modal.ts index 7714bee3e..6bec6834e 100644 --- a/packages/compute-providers/src/worker-env/modal.ts +++ b/packages/compute-providers/src/worker-env/modal.ts @@ -7,6 +7,7 @@ export function buildModalWorkerEnv({ authToken, sandboxExpiresAtMs, extraEnv, + inferenceGatewayEnabled, deploymentSlug, environmentId, baseImageRef, @@ -16,7 +17,12 @@ export function buildModalWorkerEnv({ baseImageRef: string; }): Record { return { - ...buildBaseWorkerEnv({ authToken, sandboxExpiresAtMs, extraEnv }), + ...buildBaseWorkerEnv({ + authToken, + sandboxExpiresAtMs, + extraEnv, + inferenceGatewayEnabled, + }), ...buildWorkerContextEnv({ provider: 'modal', fingerprint: baseImageRef, diff --git a/packages/compute-providers/src/worker-env/types.ts b/packages/compute-providers/src/worker-env/types.ts index 2985c7775..ce56fccda 100644 --- a/packages/compute-providers/src/worker-env/types.ts +++ b/packages/compute-providers/src/worker-env/types.ts @@ -2,4 +2,10 @@ export type BuildWorkerEnvOptions = { authToken: string; sandboxExpiresAtMs?: number; extraEnv?: Record; + /** + * Evaluated InferenceGateway feature flag. When true, gateway-covered + * provider keys are withheld from the worker daemon env; the per-task + * dequeue env routes the harness through the gateway instead. + */ + inferenceGatewayEnabled?: boolean; }; diff --git a/packages/db/src/lib/model-runtime-config.test.ts b/packages/db/src/lib/model-runtime-config.test.ts index 28b088063..8e1407085 100644 --- a/packages/db/src/lib/model-runtime-config.test.ts +++ b/packages/db/src/lib/model-runtime-config.test.ts @@ -620,4 +620,66 @@ describe('resolveEffectiveModelRuntimeEnv', () => { expect(env).not.toHaveProperty('OPENCODE_AUTH_CONTENT'); }); + + describe('inference gateway', () => { + beforeEach(() => { + mockDeploymentSettingsFindFirst.mockResolvedValue({ + runtimeModelConfig: { roomoteModel: 'anthropic/claude-sonnet-4' }, + }); + }); + + it('replaces gateway-covered keys with the gateway URL when enabled', async () => { + const env = await resolveEffectiveModelRuntimeEnv({ + inferenceGateway: true, + runtimeEnv: { TRPC_URL: 'https://api.roomote.example.com' }, + deploymentEnvVars: { ANTHROPIC_API_KEY: 'sk-anthropic' }, + }); + + expect(env).not.toHaveProperty('ANTHROPIC_API_KEY'); + expect(env.R_INFERENCE_GATEWAY_URL).toBe( + 'https://api.roomote.example.com/api/inference', + ); + }); + + it('keeps raw keys when the gateway option is not passed', async () => { + const env = await resolveEffectiveModelRuntimeEnv({ + runtimeEnv: { TRPC_URL: 'https://api.roomote.example.com' }, + deploymentEnvVars: { ANTHROPIC_API_KEY: 'sk-anthropic' }, + }); + + expect(env.ANTHROPIC_API_KEY).toBe('sk-anthropic'); + expect(env).not.toHaveProperty('R_INFERENCE_GATEWAY_URL'); + }); + + it('falls back to raw keys when no platform API URL is available', async () => { + const env = await resolveEffectiveModelRuntimeEnv({ + inferenceGateway: true, + runtimeEnv: {}, + deploymentEnvVars: { ANTHROPIC_API_KEY: 'sk-anthropic' }, + }); + + expect(env.ANTHROPIC_API_KEY).toBe('sk-anthropic'); + expect(env).not.toHaveProperty('R_INFERENCE_GATEWAY_URL'); + }); + + it('keeps keys the gateway does not cover', async () => { + const env = await resolveEffectiveModelRuntimeEnv({ + inferenceGateway: true, + runtimeEnv: { + TRPC_URL: 'https://api.roomote.example.com', + R_MODEL_ENV_KEYS: 'ANTHROPIC_API_KEY,AWS_BEARER_TOKEN_BEDROCK', + }, + deploymentEnvVars: { + ANTHROPIC_API_KEY: 'sk-anthropic', + AWS_BEARER_TOKEN_BEDROCK: 'bedrock-key', + }, + }); + + expect(env).not.toHaveProperty('ANTHROPIC_API_KEY'); + expect(env.AWS_BEARER_TOKEN_BEDROCK).toBe('bedrock-key'); + expect(env.R_INFERENCE_GATEWAY_URL).toBe( + 'https://api.roomote.example.com/api/inference', + ); + }); + }); }); diff --git a/packages/db/src/lib/model-runtime-config.ts b/packages/db/src/lib/model-runtime-config.ts index fa39e7530..b992caacb 100644 --- a/packages/db/src/lib/model-runtime-config.ts +++ b/packages/db/src/lib/model-runtime-config.ts @@ -1,9 +1,12 @@ import { eq } from 'drizzle-orm'; import { + buildInferenceGatewayUrl, CHATGPT_OPENCODE_PROVIDER_ID, DEFAULT_MODEL_ROLE_REASONING_EFFORTS, getModelProviderEnvKeyCandidates, getTaskModelCatalog, + INFERENCE_GATEWAY_PROVIDER_ENV_VAR_NAMES, + INFERENCE_GATEWAY_URL_ENV_VAR_NAME, isConfiguredEnvValue, normalizeDeploymentModelConfig, normalizeOptionalReasoningEffort, @@ -171,6 +174,15 @@ export async function resolveEffectiveModelRuntimeEnv( runtimeEnv?: Partial>; deploymentEnvVars?: Record; executor?: DatabaseOrTransaction; + /** + * Route sandbox inference through the gateway: provider keys the gateway + * covers stay on the control plane and `R_INFERENCE_GATEWAY_URL` is + * emitted instead. Callers resolving env for the sandbox pass the + * evaluated InferenceGateway feature flag; control-plane callers + * (routing, titles, summaries) omit it because that inference holds no + * run token to present to the gateway. + */ + inferenceGateway?: boolean; } = {}, ): Promise> { const runtimeEnv = options.runtimeEnv ?? process.env; @@ -283,8 +295,26 @@ export async function resolveEffectiveModelRuntimeEnv( resolvedRoomotePlanningModel, ], }); + // Gateway-covered provider keys stay on the control plane, and the sandbox + // gets the gateway URL instead. Keys the gateway cannot serve yet (Bedrock, + // Vertex) still flow through. Requires a sandbox-reachable platform API + // URL; when none is configured, fall back to direct keys rather than + // breaking inference. + const gatewayPlatformApiUrl = normalizeConfiguredValue(runtimeEnv.TRPC_URL); + const inferenceGatewayUrl = + options.inferenceGateway && gatewayPlatformApiUrl + ? buildInferenceGatewayUrl(gatewayPlatformApiUrl) + : undefined; + const gatewayCoveredKeyNames = new Set( + INFERENCE_GATEWAY_PROVIDER_ENV_VAR_NAMES, + ); + const resolvedProviderKeyValues = Object.fromEntries( providerKeyNames.flatMap((envVarName) => { + if (inferenceGatewayUrl && gatewayCoveredKeyNames.has(envVarName)) { + return []; + } + const value = normalizeConfiguredValue(runtimeEnv[envVarName]) ?? normalizeConfiguredValue(persistedEnvVars[envVarName]); @@ -364,6 +394,9 @@ export async function resolveEffectiveModelRuntimeEnv( R_MODEL_ENV_KEYS: providerKeyNames.join(','), }), ...resolvedProviderKeyValues, + ...(inferenceGatewayUrl && { + [INFERENCE_GATEWAY_URL_ENV_VAR_NAME]: inferenceGatewayUrl, + }), ...(injectedOpenCodeAuthContent && { OPENCODE_AUTH_CONTENT: injectedOpenCodeAuthContent, }), diff --git a/packages/feature-flags/src/config.ts b/packages/feature-flags/src/config.ts index 27f56f626..28ad4cb98 100644 --- a/packages/feature-flags/src/config.ts +++ b/packages/feature-flags/src/config.ts @@ -64,6 +64,17 @@ export const FEATURE_FLAG_CONFIG: FeatureFlagConfigMap = { 'Enable background subagents so the Task tool can launch subagents asynchronously via its background flag, and standard-task delivery ships the PR before visual proof instead of blocking on it. Off by default: proof runs foreground, before delivery.', }, + [FeatureFlag.InferenceGateway]: { + // Off by default: sandboxes receive raw provider API keys as before. + // When enabled, gateway-covered provider keys stay on the control plane + // and task sandboxes route inference through /api/inference with their + // run-scoped token. + defaultValue: false, + metadataKey: 'inference_gateway', + description: + 'Route sandbox inference through the platform inference gateway so provider API keys for covered providers (OpenRouter, Anthropic, OpenAI, Google Gemini) stay on the control plane instead of entering task sandboxes', + }, + [FeatureFlag.CodeMode]: { // Off by default: keeps every MCP tool schema in the agent tool list. When // enabled, the worker injects OPENCODE_EXPERIMENTAL_CODE_MODE so OpenCode diff --git a/packages/feature-flags/src/types.ts b/packages/feature-flags/src/types.ts index 1d8f8190f..73e3ee91d 100644 --- a/packages/feature-flags/src/types.ts +++ b/packages/feature-flags/src/types.ts @@ -11,6 +11,7 @@ export enum FeatureFlag { AuthorshipRules = 'AuthorshipRules', BackgroundSubagents = 'BackgroundSubagents', CodeMode = 'CodeMode', + InferenceGateway = 'InferenceGateway', } export type FeatureFlagValue = diff --git a/packages/sdk/src/server/lib/task-runs/dequeue-helpers.ts b/packages/sdk/src/server/lib/task-runs/dequeue-helpers.ts index 2a7b19353..0875fbadb 100644 --- a/packages/sdk/src/server/lib/task-runs/dequeue-helpers.ts +++ b/packages/sdk/src/server/lib/task-runs/dequeue-helpers.ts @@ -1,5 +1,7 @@ import { CONTROL_PLANE_ENV_VAR_NAMES, + INFERENCE_GATEWAY_PROVIDER_ENV_VAR_NAMES, + INFERENCE_GATEWAY_URL_ENV_VAR_NAME, RunStatus, buildSourceControlTokenMetadata, getSourceControlProviderLabel, @@ -22,6 +24,11 @@ import { sql, } from '@roomote/db/server'; import { decryptSecrets } from '@roomote/db/encryption'; +import { + FeatureFlag, + getFeatureFlagEvaluator, +} from '@roomote/feature-flags/server'; +import { getRedis } from '@roomote/redis'; import { createTaskRunWorkerGitHubToken } from '@roomote/github'; import { createTaskRunScopedGitLabTokens } from '@roomote/gitlab'; import { createTaskRunBitbucketCredentials } from '@roomote/bitbucket'; @@ -208,19 +215,64 @@ export async function fetchResolvedRuntimeEnvVars( deploymentEnvVars ?? (await loadPersistedDeploymentEnvVarsFromDb()); const resolvedModelRuntimeEnv = await resolveEffectiveModelRuntimeEnv({ deploymentEnvVars: envVars, + inferenceGateway: await isInferenceGatewayFlagEnabled(), }); return redactControlPlaneEnvVars( redactSourceControlProviderEnvVars( - withLegacySnapshotModelEnvAliases({ - ...envVars, - ...resolvedModelRuntimeEnv, - }), + redactInferenceGatewayProviderKeys( + withLegacySnapshotModelEnvAliases({ + ...envVars, + ...resolvedModelRuntimeEnv, + }), + ), options?.sourceControlProvider, ), ); } +/** + * Evaluate the deployment-level InferenceGateway feature flag. Fails closed + * (direct provider keys, the long-standing behavior) if flag evaluation is + * unavailable, so a Redis outage cannot break task inference. + */ +async function isInferenceGatewayFlagEnabled(): Promise { + try { + return await getFeatureFlagEvaluator(getRedis()).evaluate( + FeatureFlag.InferenceGateway, + { isDeploymentContext: true }, + ); + } catch (error) { + console.warn( + `[fetchResolvedRuntimeEnvVars] InferenceGateway flag evaluation failed; using direct provider keys: ${ + error instanceof Error ? error.message : String(error) + }`, + ); + + return false; + } +} + +/** + * When the inference gateway is active (the resolver emitted a gateway URL), + * provider keys the gateway covers stay on the control plane. The raw + * deployment env vars are spread into the sandbox env above, so the keys + * must be stripped from the merged result, not just left unresolved. + */ +function redactInferenceGatewayProviderKeys( + envVars: Record, +): Record { + if (!envVars[INFERENCE_GATEWAY_URL_ENV_VAR_NAME]) { + return envVars; + } + + const coveredKeyNames = new Set(INFERENCE_GATEWAY_PROVIDER_ENV_VAR_NAMES); + + return Object.fromEntries( + Object.entries(envVars).filter(([key]) => !coveredKeyNames.has(key)), + ); +} + async function loadPersistedDeploymentEnvVarsFromDb(): Promise< Record > { diff --git a/packages/types/src/index.ts b/packages/types/src/index.ts index 673aa3079..4e8d143b7 100644 --- a/packages/types/src/index.ts +++ b/packages/types/src/index.ts @@ -28,6 +28,7 @@ export * from './linear'; export * from './llm-citation-artifacts'; export * from './live-previews'; export * from './logging'; +export * from './inference-gateway'; export * from './model-provider-config'; export * from './recommended-task-models'; export * from './opencode-openrouter-variants'; diff --git a/packages/types/src/inference-gateway.ts b/packages/types/src/inference-gateway.ts new file mode 100644 index 000000000..30a7e2c9e --- /dev/null +++ b/packages/types/src/inference-gateway.ts @@ -0,0 +1,136 @@ +import type { SetupModelProviderId } from './model-provider-config'; + +/** + * Sandbox-facing env var carrying the gateway base URL (e.g. + * `https://api.example.com/api/inference`). Emitted by the control plane at + * dequeue when the InferenceGateway feature flag is enabled; its presence is + * what switches the worker's OpenCode config onto the gateway. + */ +export const INFERENCE_GATEWAY_URL_ENV_VAR_NAME = 'R_INFERENCE_GATEWAY_URL'; + +interface InferenceGatewayAuthHeader { + name: string; + scheme?: 'bearer'; +} + +export interface InferenceGatewayProvider { + /** + * Provider id, used both as the gateway path segment and as the OpenCode + * provider id (they match for every supported provider). + */ + id: SetupModelProviderId; + name: string; + /** + * Deployment env vars holding the provider API key the gateway injects, + * in precedence order. + */ + envVarNames: readonly string[]; + upstreamBaseUrl: string; + /** How the upstream expects its API key when the gateway forwards. */ + authHeader: InferenceGatewayAuthHeader; + /** + * Upstream inference endpoints the gateway forwards, matched exactly + * (Google model routes additionally match nested paths because the model + * ID and action live below `/models`). Everything else is rejected so a + * run token can only reach inference endpoints, never the provider's + * account, billing, or admin surface. + */ + allowedPaths: readonly string[]; + /** + * Base-path suffix the provider's SDK expects on a baseURL override, + * mirroring the upstream default base path (the SDK appends endpoint paths + * like `/messages` or `/chat/completions` below it). + */ + openCodeBaseUrlSuffix: string; +} + +/** + * Providers reachable through the inference gateway. Key-authenticated + * HTTP-proxyable providers only: providers that need request signing + * (Vertex service accounts) or client-managed OAuth refresh (ChatGPT + * subscriptions) are handled separately. + */ +export const INFERENCE_GATEWAY_PROVIDERS: readonly InferenceGatewayProvider[] = + [ + { + id: 'openrouter', + name: 'OpenRouter', + envVarNames: ['OPENROUTER_API_KEY'], + upstreamBaseUrl: 'https://openrouter.ai/api', + authHeader: { name: 'authorization', scheme: 'bearer' }, + allowedPaths: [ + '/v1/chat/completions', + '/v1/completions', + '/v1/embeddings', + '/v1/models', + ], + openCodeBaseUrlSuffix: '/v1', + }, + { + id: 'anthropic', + name: 'Anthropic', + envVarNames: ['ANTHROPIC_API_KEY'], + upstreamBaseUrl: 'https://api.anthropic.com', + authHeader: { name: 'x-api-key' }, + allowedPaths: ['/v1/messages', '/v1/messages/count_tokens', '/v1/models'], + openCodeBaseUrlSuffix: '/v1', + }, + { + id: 'openai', + name: 'OpenAI', + envVarNames: ['OPENAI_API_KEY'], + upstreamBaseUrl: 'https://api.openai.com', + authHeader: { name: 'authorization', scheme: 'bearer' }, + allowedPaths: [ + '/v1/chat/completions', + '/v1/responses', + '/v1/embeddings', + '/v1/models', + ], + openCodeBaseUrlSuffix: '/v1', + }, + { + id: 'google', + name: 'Google Gemini', + envVarNames: ['GEMINI_API_KEY', 'GOOGLE_GENERATIVE_AI_API_KEY'], + upstreamBaseUrl: 'https://generativelanguage.googleapis.com', + authHeader: { name: 'x-goog-api-key' }, + allowedPaths: ['/v1beta/models', '/v1/models'], + openCodeBaseUrlSuffix: '/v1beta', + }, + ]; + +/** + * Env var names of provider keys the gateway replaces. When the gateway is + * enabled these keys stay on the control plane instead of entering sandbox + * env vars. + */ +export const INFERENCE_GATEWAY_PROVIDER_ENV_VAR_NAMES: readonly string[] = + INFERENCE_GATEWAY_PROVIDERS.flatMap((provider) => provider.envVarNames); + +export function getInferenceGatewayProvider( + providerId: string, +): InferenceGatewayProvider | undefined { + return INFERENCE_GATEWAY_PROVIDERS.find( + (provider) => provider.id === providerId, + ); +} + +/** + * The OpenCode `baseURL` override pointing a provider's SDK at the gateway: + * `/`. + */ +export function buildInferenceGatewayOpenCodeBaseUrl( + gatewayUrl: string, + provider: InferenceGatewayProvider, +): string { + return `${gatewayUrl.replace(/\/+$/, '')}/${provider.id}${provider.openCodeBaseUrlSuffix}`; +} + +/** + * The gateway base URL for a deployment, derived from the sandbox-reachable + * platform API URL (the same base workers use for tRPC). + */ +export function buildInferenceGatewayUrl(platformApiUrl: string): string { + return `${platformApiUrl.replace(/\/+$/, '')}/api/inference`; +} From 17ceb384cffdb50ee15615e3b0df688397b3b362 Mon Sep 17 00:00:00 2001 From: Matt Rubens <2600+mrubens@users.noreply.github.com> Date: Thu, 16 Jul 2026 18:06:14 -0400 Subject: [PATCH 4/7] Cover all key-authenticated model providers in the inference gateway Extends the gateway registry from four providers to every provider the setup catalog supports with API-key auth: Vercel AI Gateway, Requesty, Baseten, Together AI, Moonshot AI, MiniMax, OpenCode Zen, xAI, and Amazon Bedrock (Mantle). Upstream bases follow models.dev, the registry OpenCode itself resolves providers from. - MiniMax and Bedrock Mantle proxy through their Anthropic-compatible endpoints (x-api-key); the rest are OpenAI-compatible bearer-auth. - Bedrock upstreams are region-templated: the gateway resolves AWS_REGION (deployment env, default us-east-1) per request and validates it before substitution. The region var still flows to sandboxes; only the bearer token is withheld. - Nested-path allowance is now a per-provider registry field instead of a hardcoded Google check, reused for the Vercel AI Gateway protocol. - The worker's gateway rebase needs no changes: it is model-prefix driven, and the Bedrock Mantle merge runs before the gateway merge so its baseURL/apiKey are overridden while its model config is kept. Still direct: Google Vertex (request signing) and ChatGPT subscriptions (client-managed OAuth refresh). --- .../__tests__/inference-gateway.test.ts | 138 ++++++++++++++- apps/api/src/handlers/inference/index.ts | 9 +- apps/api/src/handlers/inference/registry.ts | 38 ++++- apps/worker/src/run-task/agent-home.test.ts | 25 +++ .../src/worker-env/__tests__/base.test.ts | 9 +- .../db/src/lib/model-runtime-config.test.ts | 11 +- packages/feature-flags/src/config.ts | 2 +- packages/types/src/inference-gateway.ts | 160 ++++++++++++++++-- 8 files changed, 365 insertions(+), 27 deletions(-) diff --git a/apps/api/src/handlers/inference/__tests__/inference-gateway.test.ts b/apps/api/src/handlers/inference/__tests__/inference-gateway.test.ts index 3c4ea6bba..e6af3f245 100644 --- a/apps/api/src/handlers/inference/__tests__/inference-gateway.test.ts +++ b/apps/api/src/handlers/inference/__tests__/inference-gateway.test.ts @@ -86,7 +86,16 @@ describe('inference gateway', () => { vi.clearAllMocks(); vi.unstubAllGlobals(); mockFindTaskRun.mockResolvedValue({ id: 42 }); - mockResolveModelProviderEnvValue.mockResolvedValue('provider-secret-key'); + mockResolveModelProviderEnvValue.mockImplementation( + async (names: string | readonly string[]) => { + const nameList = typeof names === 'string' ? [names] : names; + + // Region lookups resolve separately from API keys. + return nameList.includes('AWS_REGION') + ? undefined + : 'provider-secret-key'; + }, + ); }); it('rejects user auth tokens', async () => { @@ -162,6 +171,133 @@ describe('inference gateway', () => { expect(fetchMock).not.toHaveBeenCalled(); }); + it('proxies OpenAI-compatible aggregator providers with bearer auth', async () => { + const app = createApp(createRunToken()); + + const cases: Array<[string, string]> = [ + ['requesty', 'https://router.requesty.ai/v1/chat/completions'], + ['baseten', 'https://inference.baseten.co/v1/chat/completions'], + ['togetherai', 'https://api.together.xyz/v1/chat/completions'], + ['moonshotai', 'https://api.moonshot.ai/v1/chat/completions'], + ['opencode', 'https://opencode.ai/zen/v1/chat/completions'], + ['xai', 'https://api.x.ai/v1/chat/completions'], + ]; + + for (const [providerId, expectedUrl] of cases) { + // A Response body can only be consumed once, so re-stub per provider. + const fetchMock = stubUpstreamFetch(); + + const response = await postMessages( + app, + `/api/inference/${providerId}/v1/chat/completions`, + ); + + expect(response.status).toBe(200); + + const [url, init] = fetchMock.mock.calls[0] as [string, RequestInit]; + expect(url).toBe(expectedUrl); + expect(new Headers(init.headers).get('authorization')).toBe( + 'Bearer provider-secret-key', + ); + } + }); + + it('proxies MiniMax through its Anthropic-compatible endpoint', async () => { + const fetchMock = stubUpstreamFetch(); + const response = await postMessages( + createApp(createRunToken()), + '/api/inference/minimax/v1/messages', + ); + + expect(response.status).toBe(200); + + const [url, init] = fetchMock.mock.calls[0] as [string, RequestInit]; + expect(url).toBe('https://api.minimax.io/anthropic/v1/messages'); + expect(new Headers(init.headers).get('x-api-key')).toBe( + 'provider-secret-key', + ); + }); + + it('allows nested paths under the Vercel AI Gateway protocol base', async () => { + const fetchMock = stubUpstreamFetch(); + const response = await postMessages( + createApp(createRunToken()), + '/api/inference/vercel/v1/ai/language-model', + ); + + expect(response.status).toBe(200); + + const [url] = fetchMock.mock.calls[0] as [string]; + expect(url).toBe('https://ai-gateway.vercel.sh/v1/ai/language-model'); + }); + + it('substitutes the default region for Bedrock upstreams', async () => { + const fetchMock = stubUpstreamFetch(); + const response = await postMessages( + createApp(createRunToken()), + '/api/inference/bedrock-mantle/v1/messages', + ); + + expect(response.status).toBe(200); + expect(mockResolveModelProviderEnvValue).toHaveBeenCalledWith([ + 'AWS_BEARER_TOKEN_BEDROCK', + ]); + + const [url, init] = fetchMock.mock.calls[0] as [string, RequestInit]; + expect(url).toBe( + 'https://bedrock-mantle.us-east-1.api.aws/anthropic/v1/messages', + ); + expect(new Headers(init.headers).get('x-api-key')).toBe( + 'provider-secret-key', + ); + }); + + it('substitutes a configured region for Bedrock upstreams', async () => { + const fetchMock = stubUpstreamFetch(); + mockResolveModelProviderEnvValue.mockImplementation( + async (names: string | readonly string[]) => { + const nameList = typeof names === 'string' ? [names] : names; + + return nameList.includes('AWS_REGION') + ? 'eu-west-1' + : 'provider-secret-key'; + }, + ); + + const response = await postMessages( + createApp(createRunToken()), + '/api/inference/bedrock-mantle/v1/messages', + ); + + expect(response.status).toBe(200); + + const [url] = fetchMock.mock.calls[0] as [string]; + expect(url).toBe( + 'https://bedrock-mantle.eu-west-1.api.aws/anthropic/v1/messages', + ); + }); + + it('rejects invalid Bedrock regions before building the upstream URL', async () => { + const fetchMock = stubUpstreamFetch(); + mockResolveModelProviderEnvValue.mockImplementation( + async (names: string | readonly string[]) => { + const nameList = typeof names === 'string' ? [names] : names; + + return nameList.includes('AWS_REGION') + ? 'https://evil.example.com/' + : 'provider-secret-key'; + }, + ); + + const response = await postMessages( + createApp(createRunToken()), + '/api/inference/bedrock-mantle/v1/messages', + ); + + expect(response.status).toBe(500); + expect(fetchMock).not.toHaveBeenCalled(); + }); + it('returns 404 when the provider key is not configured', async () => { const fetchMock = stubUpstreamFetch(); mockResolveModelProviderEnvValue.mockResolvedValue(undefined); diff --git a/apps/api/src/handlers/inference/index.ts b/apps/api/src/handlers/inference/index.ts index e31f32380..2ed8525df 100644 --- a/apps/api/src/handlers/inference/index.ts +++ b/apps/api/src/handlers/inference/index.ts @@ -19,6 +19,7 @@ import { formatProviderAuthHeaderValue, getInferenceProvider, isInferencePathAllowed, + resolveProviderUpstreamBaseUrl, } from './registry'; /** @@ -145,12 +146,14 @@ inference.on(['POST', 'GET'], '/:provider/*', async (c) => { } let apiKey: string | undefined; + let upstreamBaseUrl: string; try { apiKey = await resolveModelProviderEnvValue(provider.envVarNames); + upstreamBaseUrl = await resolveProviderUpstreamBaseUrl(provider); } catch (error) { console.error( - formatSingleLineLog(`${logPrefix} Failed to resolve provider key`, { + formatSingleLineLog(`${logPrefix} Failed to resolve provider config`, { requestId, method, runId: auth.runId, @@ -159,7 +162,7 @@ inference.on(['POST', 'GET'], '/:provider/*', async (c) => { ); return c.json( - { error: `Failed to resolve the ${provider.name} API key` }, + { error: `Failed to resolve the ${provider.name} configuration` }, 500, ); } @@ -174,7 +177,7 @@ inference.on(['POST', 'GET'], '/:provider/*', async (c) => { } const search = new URL(c.req.url).search; - const upstreamUrl = `${provider.upstreamBaseUrl}${upstreamPath}${search}`; + const upstreamUrl = `${upstreamBaseUrl}${upstreamPath}${search}`; try { const upstreamResponse = await fetchWithLongLivedStreamDispatcher( diff --git a/apps/api/src/handlers/inference/registry.ts b/apps/api/src/handlers/inference/registry.ts index 2835c9a25..b98324380 100644 --- a/apps/api/src/handlers/inference/registry.ts +++ b/apps/api/src/handlers/inference/registry.ts @@ -1,7 +1,9 @@ import { getInferenceGatewayProvider, + INFERENCE_GATEWAY_REGION_PATTERN, type InferenceGatewayProvider, } from '@roomote/types'; +import { resolveModelProviderEnvValue } from '@roomote/db/server'; export function getInferenceProvider( providerId: string, @@ -10,9 +12,10 @@ export function getInferenceProvider( } /** - * A path is allowed when it exactly matches an inference endpoint. Google - * model routes are the exception because their model ID and action are part - * of the nested path; their provider API has no account surface below models. + * A path is allowed when it exactly matches an inference endpoint. Providers + * whose route shape requires it (Google model routes, Vercel's AI Gateway + * protocol) additionally allow nested paths; their upstreams serve no + * account surface below the allowed prefixes. */ export function isInferencePathAllowed( provider: InferenceGatewayProvider, @@ -23,10 +26,37 @@ export function isInferencePathAllowed( return true; } - return provider.id === 'google' && upstreamPath.startsWith(`${path}/`); + return ( + provider.allowNestedPaths === true && upstreamPath.startsWith(`${path}/`) + ); }); } +/** + * Resolve the provider's upstream base, substituting the `{region}` + * placeholder from the deployment's region env var (falling back to the + * provider default) for region-templated upstreams like Bedrock. + */ +export async function resolveProviderUpstreamBaseUrl( + provider: InferenceGatewayProvider, +): Promise { + if (!provider.region) { + return provider.upstreamBaseUrl; + } + + const region = + (await resolveModelProviderEnvValue([provider.region.envVarName])) ?? + provider.region.default; + + if (!INFERENCE_GATEWAY_REGION_PATTERN.test(region)) { + throw new Error( + `${provider.region.envVarName} must be a valid region for ${provider.name}. Received "${region}".`, + ); + } + + return provider.upstreamBaseUrl.replace('{region}', region); +} + export function formatProviderAuthHeaderValue( provider: InferenceGatewayProvider, apiKey: string, diff --git a/apps/worker/src/run-task/agent-home.test.ts b/apps/worker/src/run-task/agent-home.test.ts index 9b6583fd7..7dd297a2d 100644 --- a/apps/worker/src/run-task/agent-home.test.ts +++ b/apps/worker/src/run-task/agent-home.test.ts @@ -173,6 +173,31 @@ describe('generateOpenCodeConfig provider support', () => { }); }); + it('rebases Bedrock Mantle onto the gateway while keeping its model config', () => { + const result = generateOpenCodeConfig({ + homeDir: createHomeDir(), + runtimeEnv: { + R_MODEL: 'bedrock-mantle/anthropic.claude-sonnet-5', + AWS_REGION: 'us-west-2', + R_INFERENCE_GATEWAY_URL: 'https://api.example.com/api/inference', + }, + }); + const config = JSON.parse(result.configContent) as { + provider: Record; + }; + + expect(config.provider['bedrock-mantle']).toMatchObject({ + npm: '@ai-sdk/anthropic', + options: { + baseURL: 'https://api.example.com/api/inference/bedrock-mantle/v1', + apiKey: '{env:ROOMOTE_CLOUD_TOKEN}', + }, + models: { + 'anthropic.claude-sonnet-5': {}, + }, + }); + }); + it('leaves the openai provider alone when a ChatGPT subscription is present', () => { const result = generateOpenCodeConfig({ homeDir: createHomeDir(), diff --git a/packages/compute-providers/src/worker-env/__tests__/base.test.ts b/packages/compute-providers/src/worker-env/__tests__/base.test.ts index ce670be9f..3f82f7f6a 100644 --- a/packages/compute-providers/src/worker-env/__tests__/base.test.ts +++ b/packages/compute-providers/src/worker-env/__tests__/base.test.ts @@ -148,6 +148,8 @@ describe('buildBaseWorkerEnv', () => { process.env.ANTHROPIC_API_KEY = 'anthropic-key'; process.env.OPENROUTER_API_KEY = 'openrouter-key'; process.env.AWS_BEARER_TOKEN_BEDROCK = 'bedrock-key'; + process.env.XAI_API_KEY = 'xai-key'; + process.env.AWS_REGION = 'us-west-2'; const env = buildBaseWorkerEnv({ authToken: 'auth-token', @@ -158,7 +160,10 @@ describe('buildBaseWorkerEnv', () => { expect(env.R_MODEL).toBe('anthropic/claude-sonnet-5'); expect(env.ANTHROPIC_API_KEY).toBeUndefined(); expect(env.OPENROUTER_API_KEY).toBeUndefined(); - // The gateway cannot serve Bedrock yet, so its key still ships. - expect(env.AWS_BEARER_TOKEN_BEDROCK).toBe('bedrock-key'); + expect(env.AWS_BEARER_TOKEN_BEDROCK).toBeUndefined(); + expect(env.XAI_API_KEY).toBeUndefined(); + // Region config is not a secret and Bedrock's Mantle merge still + // validates it sandbox-side. + expect(env.AWS_REGION).toBe('us-west-2'); }); }); diff --git a/packages/db/src/lib/model-runtime-config.test.ts b/packages/db/src/lib/model-runtime-config.test.ts index 8e1407085..7506f5c3a 100644 --- a/packages/db/src/lib/model-runtime-config.test.ts +++ b/packages/db/src/lib/model-runtime-config.test.ts @@ -667,16 +667,23 @@ describe('resolveEffectiveModelRuntimeEnv', () => { inferenceGateway: true, runtimeEnv: { TRPC_URL: 'https://api.roomote.example.com', - R_MODEL_ENV_KEYS: 'ANTHROPIC_API_KEY,AWS_BEARER_TOKEN_BEDROCK', + R_MODEL_ENV_KEYS: + 'ANTHROPIC_API_KEY,AWS_BEARER_TOKEN_BEDROCK,GOOGLE_APPLICATION_CREDENTIALS', }, deploymentEnvVars: { ANTHROPIC_API_KEY: 'sk-anthropic', AWS_BEARER_TOKEN_BEDROCK: 'bedrock-key', + GOOGLE_APPLICATION_CREDENTIALS: '{"type":"service_account"}', }, }); expect(env).not.toHaveProperty('ANTHROPIC_API_KEY'); - expect(env.AWS_BEARER_TOKEN_BEDROCK).toBe('bedrock-key'); + expect(env).not.toHaveProperty('AWS_BEARER_TOKEN_BEDROCK'); + // Vertex needs request signing, not header injection, so its + // credentials still flow until the gateway can serve it. + expect(env.GOOGLE_APPLICATION_CREDENTIALS).toBe( + '{"type":"service_account"}', + ); expect(env.R_INFERENCE_GATEWAY_URL).toBe( 'https://api.roomote.example.com/api/inference', ); diff --git a/packages/feature-flags/src/config.ts b/packages/feature-flags/src/config.ts index 28ad4cb98..8f17d4112 100644 --- a/packages/feature-flags/src/config.ts +++ b/packages/feature-flags/src/config.ts @@ -72,7 +72,7 @@ export const FEATURE_FLAG_CONFIG: FeatureFlagConfigMap = { defaultValue: false, metadataKey: 'inference_gateway', description: - 'Route sandbox inference through the platform inference gateway so provider API keys for covered providers (OpenRouter, Anthropic, OpenAI, Google Gemini) stay on the control plane instead of entering task sandboxes', + 'Route sandbox inference through the platform inference gateway so API keys for key-authenticated model providers stay on the control plane instead of entering task sandboxes (Vertex and ChatGPT-subscription auth still flow directly)', }, [FeatureFlag.CodeMode]: { diff --git a/packages/types/src/inference-gateway.ts b/packages/types/src/inference-gateway.ts index 30a7e2c9e..32685479b 100644 --- a/packages/types/src/inference-gateway.ts +++ b/packages/types/src/inference-gateway.ts @@ -8,6 +8,17 @@ import type { SetupModelProviderId } from './model-provider-config'; */ export const INFERENCE_GATEWAY_URL_ENV_VAR_NAME = 'R_INFERENCE_GATEWAY_URL'; +/** + * OpenCode registers Bedrock models under this provider id (an + * Anthropic-compatible endpoint), distinct from the `amazon-bedrock` setup + * provider id. + */ +export const BEDROCK_MANTLE_OPENCODE_PROVIDER_ID = 'bedrock-mantle'; + +/** Matches valid cloud regions like `us-east-1`. */ +export const INFERENCE_GATEWAY_REGION_PATTERN = + /^[a-z]{2}(?:-[a-z0-9]+)+-\d+$/u; + interface InferenceGatewayAuthHeader { name: string; scheme?: 'bearer'; @@ -18,24 +29,38 @@ export interface InferenceGatewayProvider { * Provider id, used both as the gateway path segment and as the OpenCode * provider id (they match for every supported provider). */ - id: SetupModelProviderId; + id: SetupModelProviderId | typeof BEDROCK_MANTLE_OPENCODE_PROVIDER_ID; name: string; /** * Deployment env vars holding the provider API key the gateway injects, * in precedence order. */ envVarNames: readonly string[]; + /** + * Upstream API base. May contain a `{region}` placeholder resolved + * per-request from `region` below. + */ upstreamBaseUrl: string; + /** + * Region resolution for `{region}`-templated upstreams: the deployment env + * var to read and the fallback when it is unset. + */ + region?: { envVarName: string; default: string }; /** How the upstream expects its API key when the gateway forwards. */ authHeader: InferenceGatewayAuthHeader; /** - * Upstream inference endpoints the gateway forwards, matched exactly - * (Google model routes additionally match nested paths because the model - * ID and action live below `/models`). Everything else is rejected so a - * run token can only reach inference endpoints, never the provider's - * account, billing, or admin surface. + * Upstream inference endpoints the gateway forwards, matched exactly. + * Everything else is rejected so a run token can only reach inference + * endpoints, never the provider's account, billing, or admin surface. */ allowedPaths: readonly string[]; + /** + * Also allow paths nested below each `allowedPaths` entry. Only for + * upstreams whose route shape requires it (Google puts the model ID and + * action below `/models`; Vercel's AI Gateway protocol lives below + * `/v1/ai`) and whose API has no account surface under those prefixes. + */ + allowNestedPaths?: boolean; /** * Base-path suffix the provider's SDK expects on a baseURL override, * mirroring the upstream default base path (the SDK appends endpoint paths @@ -44,11 +69,30 @@ export interface InferenceGatewayProvider { openCodeBaseUrlSuffix: string; } +/** + * Inference endpoints for OpenAI-compatible upstreams + * (`@ai-sdk/openai-compatible` and SDKs that share its route shape). + */ +const OPENAI_COMPATIBLE_INFERENCE_PATHS: readonly string[] = [ + '/v1/chat/completions', + '/v1/completions', + '/v1/embeddings', + '/v1/models', +]; + +/** Inference endpoints for Anthropic-compatible upstreams. */ +const ANTHROPIC_COMPATIBLE_INFERENCE_PATHS: readonly string[] = [ + '/v1/messages', + '/v1/messages/count_tokens', + '/v1/models', +]; + /** * Providers reachable through the inference gateway. Key-authenticated * HTTP-proxyable providers only: providers that need request signing * (Vertex service accounts) or client-managed OAuth refresh (ChatGPT - * subscriptions) are handled separately. + * subscriptions) are handled separately. Upstream bases follow models.dev, + * the registry OpenCode itself resolves providers from. */ export const INFERENCE_GATEWAY_PROVIDERS: readonly InferenceGatewayProvider[] = [ @@ -58,12 +102,7 @@ export const INFERENCE_GATEWAY_PROVIDERS: readonly InferenceGatewayProvider[] = envVarNames: ['OPENROUTER_API_KEY'], upstreamBaseUrl: 'https://openrouter.ai/api', authHeader: { name: 'authorization', scheme: 'bearer' }, - allowedPaths: [ - '/v1/chat/completions', - '/v1/completions', - '/v1/embeddings', - '/v1/models', - ], + allowedPaths: OPENAI_COMPATIBLE_INFERENCE_PATHS, openCodeBaseUrlSuffix: '/v1', }, { @@ -72,7 +111,7 @@ export const INFERENCE_GATEWAY_PROVIDERS: readonly InferenceGatewayProvider[] = envVarNames: ['ANTHROPIC_API_KEY'], upstreamBaseUrl: 'https://api.anthropic.com', authHeader: { name: 'x-api-key' }, - allowedPaths: ['/v1/messages', '/v1/messages/count_tokens', '/v1/models'], + allowedPaths: ANTHROPIC_COMPATIBLE_INFERENCE_PATHS, openCodeBaseUrlSuffix: '/v1', }, { @@ -96,8 +135,101 @@ export const INFERENCE_GATEWAY_PROVIDERS: readonly InferenceGatewayProvider[] = upstreamBaseUrl: 'https://generativelanguage.googleapis.com', authHeader: { name: 'x-goog-api-key' }, allowedPaths: ['/v1beta/models', '/v1/models'], + allowNestedPaths: true, openCodeBaseUrlSuffix: '/v1beta', }, + { + id: 'vercel', + name: 'Vercel AI Gateway', + envVarNames: ['AI_GATEWAY_API_KEY'], + upstreamBaseUrl: 'https://ai-gateway.vercel.sh', + authHeader: { name: 'authorization', scheme: 'bearer' }, + // The @ai-sdk/gateway protocol (config, language-model, embedding-model) + // lives below /v1/ai; the host serves inference only. + allowedPaths: ['/v1/ai'], + allowNestedPaths: true, + openCodeBaseUrlSuffix: '/v1/ai', + }, + { + id: 'requesty', + name: 'Requesty', + envVarNames: ['REQUESTY_API_KEY'], + upstreamBaseUrl: 'https://router.requesty.ai', + authHeader: { name: 'authorization', scheme: 'bearer' }, + allowedPaths: OPENAI_COMPATIBLE_INFERENCE_PATHS, + openCodeBaseUrlSuffix: '/v1', + }, + { + id: 'baseten', + name: 'Baseten', + envVarNames: ['BASETEN_API_KEY'], + upstreamBaseUrl: 'https://inference.baseten.co', + authHeader: { name: 'authorization', scheme: 'bearer' }, + allowedPaths: OPENAI_COMPATIBLE_INFERENCE_PATHS, + openCodeBaseUrlSuffix: '/v1', + }, + { + id: 'togetherai', + name: 'Together AI', + envVarNames: ['TOGETHER_API_KEY'], + upstreamBaseUrl: 'https://api.together.xyz', + authHeader: { name: 'authorization', scheme: 'bearer' }, + allowedPaths: OPENAI_COMPATIBLE_INFERENCE_PATHS, + openCodeBaseUrlSuffix: '/v1', + }, + { + id: 'moonshotai', + name: 'Moonshot AI', + envVarNames: ['MOONSHOT_API_KEY'], + upstreamBaseUrl: 'https://api.moonshot.ai', + authHeader: { name: 'authorization', scheme: 'bearer' }, + allowedPaths: OPENAI_COMPATIBLE_INFERENCE_PATHS, + openCodeBaseUrlSuffix: '/v1', + }, + { + id: 'minimax', + name: 'MiniMax', + envVarNames: ['MINIMAX_API_KEY'], + // models.dev registers MiniMax through its Anthropic-compatible + // endpoint (@ai-sdk/anthropic with base /anthropic/v1). + upstreamBaseUrl: 'https://api.minimax.io/anthropic', + authHeader: { name: 'x-api-key' }, + allowedPaths: ANTHROPIC_COMPATIBLE_INFERENCE_PATHS, + openCodeBaseUrlSuffix: '/v1', + }, + { + id: 'opencode', + name: 'OpenCode Zen', + envVarNames: ['OPENCODE_API_KEY'], + upstreamBaseUrl: 'https://opencode.ai/zen', + authHeader: { name: 'authorization', scheme: 'bearer' }, + allowedPaths: OPENAI_COMPATIBLE_INFERENCE_PATHS, + openCodeBaseUrlSuffix: '/v1', + }, + { + id: 'xai', + name: 'xAI', + envVarNames: ['XAI_API_KEY'], + upstreamBaseUrl: 'https://api.x.ai', + authHeader: { name: 'authorization', scheme: 'bearer' }, + allowedPaths: [ + '/v1/chat/completions', + '/v1/responses', + '/v1/embeddings', + '/v1/models', + ], + openCodeBaseUrlSuffix: '/v1', + }, + { + id: BEDROCK_MANTLE_OPENCODE_PROVIDER_ID, + name: 'Amazon Bedrock', + envVarNames: ['AWS_BEARER_TOKEN_BEDROCK'], + upstreamBaseUrl: 'https://bedrock-mantle.{region}.api.aws/anthropic', + region: { envVarName: 'AWS_REGION', default: 'us-east-1' }, + authHeader: { name: 'x-api-key' }, + allowedPaths: ANTHROPIC_COMPATIBLE_INFERENCE_PATHS, + openCodeBaseUrlSuffix: '/v1', + }, ]; /** From fcf06bef9e460efcabfada2450adfed24321002a Mon Sep 17 00:00:00 2001 From: Matt Rubens <2600+mrubens@users.noreply.github.com> Date: Thu, 16 Jul 2026 19:14:19 -0400 Subject: [PATCH 5/7] Harden the inference gateway from the code review Merges develop and addresses the /code-review findings, consolidating the gateway decision to a single per-run signal. - Gate the /api/inference endpoint on the InferenceGateway flag: it was live on every deployment, so a run token could pull a provider key its sandbox never held. Fails closed on evaluation error. - Reject encoded-slash / dot-segment traversal on nested-path providers (google, vercel) so a run token cannot ride the allowlist to a non-inference upstream endpoint. - Build the sandbox gateway URL worker-side from the container-reachable platform URL instead of dequeue's raw TRPC_URL (fixes a Docker/ self-host outage where the emitted localhost URL was unreachable). - Thread one authoritative served-keys list (R_INFERENCE_GATEWAY_KEYS): the resolver withholds and advertises exactly the configured, gateway-covered keys; dequeue redaction strips exactly that set (so a provider key the deployment set for its own code survives); the worker strips exactly that set from the harness env (defeating daemon-env re-injection on pre-flag sandboxes) and rebases exactly those providers (so a repo-pinned provider is not withheld-without-rebase). - Apply the same gateway withholding to the snapshot env path. - Prefer GOOGLE_GENERATIVE_AI_API_KEY over GEMINI_API_KEY to match the SDK's pre-gateway precedence when both are set. - Share one fail-closed flag evaluator across dequeue, spawn, and the gateway; share the Bedrock region/id constants via @roomote/types. - Fix stale comments that claimed Bedrock keys still flow directly. --- .../__tests__/inference-gateway.test.ts | 59 ++++++++++++++++--- apps/api/src/handlers/inference/index.ts | 19 ++++++ apps/api/src/handlers/inference/registry.ts | 25 ++++++++ .../inference-gateway-flag.ts | 25 ++------ apps/web/src/lib/mock-utils.ts | 1 + apps/worker/src/run-task/agent-home.test.ts | 4 ++ apps/worker/src/run-task/agent-home.ts | 50 ++++++++-------- apps/worker/src/run-task/env.ts | 6 +- apps/worker/src/run-task/run-task.ts | 32 ++++++++++ .../compute-providers/src/worker-env/base.ts | 8 ++- .../db/src/lib/model-runtime-config.test.ts | 39 +++++++----- packages/db/src/lib/model-runtime-config.ts | 49 ++++++++------- packages/feature-flags/src/evaluator.ts | 27 +++++++++ packages/feature-flags/src/server/index.ts | 1 + .../__tests__/dequeue-helpers.test.ts | 48 +++++++++++++++ .../__tests__/fetch-snapshot-env.test.ts | 45 ++++++-------- .../server/lib/task-runs/dequeue-helpers.ts | 56 +++++++----------- .../lib/task-runs/fetch-snapshot-env.ts | 16 ++--- packages/types/src/inference-gateway.ts | 55 ++++++++++++++++- 19 files changed, 395 insertions(+), 170 deletions(-) diff --git a/apps/api/src/handlers/inference/__tests__/inference-gateway.test.ts b/apps/api/src/handlers/inference/__tests__/inference-gateway.test.ts index e6af3f245..128fa6694 100644 --- a/apps/api/src/handlers/inference/__tests__/inference-gateway.test.ts +++ b/apps/api/src/handlers/inference/__tests__/inference-gateway.test.ts @@ -3,12 +3,17 @@ import type { AuthTokenContext, RunTokenContext } from '@roomote/types'; import type { Variables } from '../../../types'; -const { mockFindTaskRun, mockResolveModelProviderEnvValue } = vi.hoisted( - () => ({ - mockFindTaskRun: vi.fn(), - mockResolveModelProviderEnvValue: vi.fn(), - }), -); +const { + mockFindTaskRun, + mockResolveModelProviderEnvValue, + mockIsInferenceGatewayEnabledForDeployment, +} = vi.hoisted(() => ({ + mockFindTaskRun: vi.fn(), + mockResolveModelProviderEnvValue: vi.fn(), + mockIsInferenceGatewayEnabledForDeployment: vi.fn( + async (_redis?: unknown, _prefix?: string) => true, + ), +})); vi.mock('@roomote/db/server', () => ({ db: { @@ -21,6 +26,15 @@ vi.mock('@roomote/db/server', () => ({ resolveModelProviderEnvValue: mockResolveModelProviderEnvValue, })); +vi.mock('@roomote/feature-flags/server', () => ({ + isInferenceGatewayEnabledForDeployment: (redis: unknown, prefix?: string) => + mockIsInferenceGatewayEnabledForDeployment(redis, prefix), +})); + +vi.mock('@roomote/redis', () => ({ + getRedis: vi.fn(() => ({})), +})); + import { inference } from '../index'; function createApp(authContext: Variables['authContext']) { @@ -85,6 +99,7 @@ describe('inference gateway', () => { beforeEach(() => { vi.clearAllMocks(); vi.unstubAllGlobals(); + mockIsInferenceGatewayEnabledForDeployment.mockResolvedValue(true); mockFindTaskRun.mockResolvedValue({ id: 42 }); mockResolveModelProviderEnvValue.mockImplementation( async (names: string | readonly string[]) => { @@ -114,6 +129,17 @@ describe('inference gateway', () => { expect(fetchMock).not.toHaveBeenCalled(); }); + it('rejects requests when the gateway flag is disabled for the deployment', async () => { + const fetchMock = stubUpstreamFetch(); + mockIsInferenceGatewayEnabledForDeployment.mockResolvedValue(false); + + const response = await postMessages(createApp(createRunToken())); + + expect(response.status).toBe(403); + expect(mockFindTaskRun).not.toHaveBeenCalled(); + expect(fetchMock).not.toHaveBeenCalled(); + }); + it('rejects run tokens whose task run no longer exists', async () => { const fetchMock = stubUpstreamFetch(); mockFindTaskRun.mockResolvedValue(undefined); @@ -171,6 +197,25 @@ describe('inference gateway', () => { expect(fetchMock).not.toHaveBeenCalled(); }); + it('rejects encoded-slash and dot-segment traversal under nested-path providers', async () => { + const fetchMock = stubUpstreamFetch(); + const app = createApp(createRunToken()); + + const [googleEncoded, googleDots, vercelEncoded] = await Promise.all([ + postMessages( + app, + '/api/inference/google/v1beta/models/..%2F..%2Fv1internal', + ), + postMessages(app, '/api/inference/google/v1beta/models/../admin'), + postMessages(app, '/api/inference/vercel/v1/ai/..%2Fadmin'), + ]); + + expect(googleEncoded.status).toBe(403); + expect(googleDots.status).toBe(403); + expect(vercelEncoded.status).toBe(403); + expect(fetchMock).not.toHaveBeenCalled(); + }); + it('proxies OpenAI-compatible aggregator providers with bearer auth', async () => { const app = createApp(createRunToken()); @@ -384,8 +429,8 @@ describe('inference gateway', () => { const headers = new Headers(init.headers); expect(headers.get('x-goog-api-key')).toBe('provider-secret-key'); expect(mockResolveModelProviderEnvValue).toHaveBeenCalledWith([ - 'GEMINI_API_KEY', 'GOOGLE_GENERATIVE_AI_API_KEY', + 'GEMINI_API_KEY', ]); }); diff --git a/apps/api/src/handlers/inference/index.ts b/apps/api/src/handlers/inference/index.ts index 2ed8525df..c15689d39 100644 --- a/apps/api/src/handlers/inference/index.ts +++ b/apps/api/src/handlers/inference/index.ts @@ -7,6 +7,8 @@ import { resolveModelProviderEnvValue, taskRuns, } from '@roomote/db/server'; +import { isInferenceGatewayEnabledForDeployment } from '@roomote/feature-flags/server'; +import { getRedis } from '@roomote/redis'; import type { Variables } from '../../types'; import { fetchWithLongLivedStreamDispatcher } from '../long-lived-fetch'; @@ -110,6 +112,23 @@ inference.on(['POST', 'GET'], '/:provider/*', async (c) => { ); } + // The gateway is only open when the deployment has opted into it. Without + // this check the endpoint would serve provider keys on every deployment + // regardless of the flag, so a run token could reach a provider key its + // sandbox never held. Fails closed (gateway closed) when evaluation is + // unavailable, matching the withholding side. + const gatewayEnabled = await isInferenceGatewayEnabledForDeployment( + getRedis(), + logPrefix, + ); + + if (!gatewayEnabled) { + return c.json( + { error: 'The inference gateway is not enabled for this deployment' }, + 403, + ); + } + const taskRun = await db.query.taskRuns.findFirst({ columns: { id: true }, where: eq(taskRuns.id, auth.runId), diff --git a/apps/api/src/handlers/inference/registry.ts b/apps/api/src/handlers/inference/registry.ts index b98324380..2804f0766 100644 --- a/apps/api/src/handlers/inference/registry.ts +++ b/apps/api/src/handlers/inference/registry.ts @@ -16,11 +16,21 @@ export function getInferenceProvider( * whose route shape requires it (Google model routes, Vercel's AI Gateway * protocol) additionally allow nested paths; their upstreams serve no * account surface below the allowed prefixes. + * + * Paths containing dot segments or percent-encoded slashes are always rejected + * first: the match runs on the WHATWG-parsed pathname (which preserves `%2F` + * and does not collapse `%2e%2e`), so a nested-path provider could otherwise + * pass `/v1beta/models/..%2F..%2Fadmin` through the `startsWith` check and + * have the upstream normalize it back out of the inference surface. */ export function isInferencePathAllowed( provider: InferenceGatewayProvider, upstreamPath: string, ): boolean { + if (hasTraversalOrEncodedSlash(upstreamPath)) { + return false; + } + return provider.allowedPaths.some((path) => { if (upstreamPath === path) { return true; @@ -32,6 +42,21 @@ export function isInferencePathAllowed( }); } +function hasTraversalOrEncodedSlash(upstreamPath: string): boolean { + const lowered = upstreamPath.toLowerCase(); + + if (lowered.includes('%2f') || lowered.includes('%5c')) { + return true; + } + + // Any `.` or `..` path segment (literal or percent-encoded dots) is out of + // bounds for an inference endpoint allowlist. + return lowered + .replace(/%2e/g, '.') + .split('/') + .some((segment) => segment === '.' || segment === '..'); +} + /** * Resolve the provider's upstream base, substituting the `{region}` * placeholder from the deployment's region env var (falling back to the diff --git a/apps/controller/src/compute-providers/inference-gateway-flag.ts b/apps/controller/src/compute-providers/inference-gateway-flag.ts index 6775d09ce..b44e3047b 100644 --- a/apps/controller/src/compute-providers/inference-gateway-flag.ts +++ b/apps/controller/src/compute-providers/inference-gateway-flag.ts @@ -1,28 +1,11 @@ -import { - FeatureFlag, - getFeatureFlagEvaluator, -} from '@roomote/feature-flags/server'; +import { isInferenceGatewayEnabledForDeployment } from '@roomote/feature-flags/server'; import { getRedis } from '@roomote/redis'; /** * Evaluate the deployment-level InferenceGateway flag for worker spawn env - * construction. Fails closed (provider keys ship to the worker daemon as - * before) when flag evaluation is unavailable, so a Redis outage cannot - * break task inference. + * construction. Delegates to the shared fail-closed evaluator so the spawn + * env and the dequeue env are gated by identical logic. */ export async function isInferenceGatewayEnabledForWorkerEnv(): Promise { - try { - return await getFeatureFlagEvaluator(getRedis()).evaluate( - FeatureFlag.InferenceGateway, - { isDeploymentContext: true }, - ); - } catch (error) { - console.warn( - `[spawn-worker] InferenceGateway flag evaluation failed; forwarding provider keys: ${ - error instanceof Error ? error.message : String(error) - }`, - ); - - return false; - } + return isInferenceGatewayEnabledForDeployment(getRedis(), '[spawn-worker]'); } diff --git a/apps/web/src/lib/mock-utils.ts b/apps/web/src/lib/mock-utils.ts index e2d61aee1..a879975bf 100644 --- a/apps/web/src/lib/mock-utils.ts +++ b/apps/web/src/lib/mock-utils.ts @@ -11,6 +11,7 @@ export const mockFeatureFlags: Record = { [FeatureFlag.AuthorshipRules]: false, [FeatureFlag.BackgroundSubagents]: false, [FeatureFlag.CodeMode]: false, + [FeatureFlag.InferenceGateway]: false, }; export const mockUserResource: UserResource = { diff --git a/apps/worker/src/run-task/agent-home.test.ts b/apps/worker/src/run-task/agent-home.test.ts index ece8db532..d85094025 100644 --- a/apps/worker/src/run-task/agent-home.test.ts +++ b/apps/worker/src/run-task/agent-home.test.ts @@ -143,6 +143,7 @@ describe('generateOpenCodeConfig provider support', () => { R_MODEL: 'anthropic/claude-sonnet-5', R_SMALL_MODEL: 'google/gemini-3.5-flash', R_INFERENCE_GATEWAY_URL: 'https://api.example.com/api/inference', + R_INFERENCE_GATEWAY_KEYS: 'ANTHROPIC_API_KEY,GEMINI_API_KEY', }, }); const config = JSON.parse(result.configContent) as { @@ -169,6 +170,7 @@ describe('generateOpenCodeConfig provider support', () => { runtimeEnv: { R_MODEL: 'openrouter/anthropic/claude-sonnet-5', R_INFERENCE_GATEWAY_URL: 'https://api.example.com/api/inference', + R_INFERENCE_GATEWAY_KEYS: 'OPENROUTER_API_KEY', }, }); const config = JSON.parse(result.configContent) as { @@ -191,6 +193,7 @@ describe('generateOpenCodeConfig provider support', () => { R_MODEL: 'bedrock-mantle/anthropic.claude-sonnet-5', AWS_REGION: 'us-west-2', R_INFERENCE_GATEWAY_URL: 'https://api.example.com/api/inference', + R_INFERENCE_GATEWAY_KEYS: 'AWS_BEARER_TOKEN_BEDROCK', }, }); const config = JSON.parse(result.configContent) as { @@ -215,6 +218,7 @@ describe('generateOpenCodeConfig provider support', () => { runtimeEnv: { R_MODEL: 'openai/gpt-5.4-codex', R_INFERENCE_GATEWAY_URL: 'https://api.example.com/api/inference', + R_INFERENCE_GATEWAY_KEYS: 'OPENAI_API_KEY', OPENCODE_AUTH_CONTENT: '{"openai":{"type":"oauth"}}', }, }); diff --git a/apps/worker/src/run-task/agent-home.ts b/apps/worker/src/run-task/agent-home.ts index f69e08662..ee8865d3f 100644 --- a/apps/worker/src/run-task/agent-home.ts +++ b/apps/worker/src/run-task/agent-home.ts @@ -2,19 +2,23 @@ import * as fs from 'node:fs'; import * as path from 'node:path'; import { + BEDROCK_MANTLE_OPENCODE_PROVIDER_ID, buildInferenceGatewayOpenCodeBaseUrl, buildOpenCodeModelReasoningOptions, CHATGPT_OPENCODE_PROVIDER_ID, collectOpenRouterVariantModelAlias, - getInferenceGatewayProvider, + DEFAULT_BEDROCK_MANTLE_REGION, + getInferenceGatewayProviderByEnvVarName, GOOGLE_APPLICATION_CREDENTIALS_ENV_VAR_NAME, getMcpIntegration, + INFERENCE_GATEWAY_KEYS_ENV_VAR_NAME, + INFERENCE_GATEWAY_REGION_PATTERN, INFERENCE_GATEWAY_URL_ENV_VAR_NAME, isInlineGoogleCredentialsValue, mergeOpenCodeModelReasoningOptions, mergeOpenRouterVariantAliasModels, normalizeOptionalReasoningEffort, - OPENCODE_AUTH_CONTENT_ENV_VAR_NAME, + parseInferenceGatewayKeys, renderManualSkillMarkdown, resolveOpenRouterVariantModelAlias, OPENCODE_ARCHITECT_AGENT, @@ -58,12 +62,6 @@ export const OPENCODE_AUTH_FILE_NAME = 'auth.json'; const OPENROUTER_PROVIDER_ID = 'openrouter'; -const BEDROCK_MANTLE_PROVIDER_ID = 'bedrock-mantle'; - -const DEFAULT_BEDROCK_MANTLE_REGION = 'us-east-1'; - -const AWS_REGION_PATTERN = /^[a-z]{2}(?:-[a-z0-9]+)+-\d+$/u; - /** * OpenRouter identifies the calling application through the `HTTP-Referer` * and `X-Title` request headers rather than the standard `User-Agent`. @@ -572,7 +570,7 @@ function mergeBedrockMantleProviderConfig( const mantleModelIds = [ ...new Set( modelIds.flatMap((modelId) => { - const prefix = `${BEDROCK_MANTLE_PROVIDER_ID}/`; + const prefix = `${BEDROCK_MANTLE_OPENCODE_PROVIDER_ID}/`; const normalized = modelId?.trim(); return normalized?.startsWith(prefix) @@ -588,13 +586,15 @@ function mergeBedrockMantleProviderConfig( const region = runtimeEnv.AWS_REGION?.trim() || DEFAULT_BEDROCK_MANTLE_REGION; - if (!AWS_REGION_PATTERN.test(region)) { + if (!INFERENCE_GATEWAY_REGION_PATTERN.test(region)) { throw new Error( `AWS_REGION must be a valid AWS region for Amazon Bedrock. Received "${region}".`, ); } - const existingProvider = asRecord(providerConfig[BEDROCK_MANTLE_PROVIDER_ID]); + const existingProvider = asRecord( + providerConfig[BEDROCK_MANTLE_OPENCODE_PROVIDER_ID], + ); const existingOptions = asRecord(existingProvider.options); const existingModels = asRecord(existingProvider.models); const models = Object.fromEntries( @@ -609,7 +609,7 @@ function mergeBedrockMantleProviderConfig( return { ...providerConfig, - [BEDROCK_MANTLE_PROVIDER_ID]: { + [BEDROCK_MANTLE_OPENCODE_PROVIDER_ID]: { ...existingProvider, npm: '@ai-sdk/anthropic', name: 'Amazon Bedrock', @@ -636,31 +636,30 @@ function mergeBedrockMantleProviderConfig( function mergeInferenceGatewayProviderConfig( providerConfig: Record, runtimeEnv: Record, - modelIds: Array, ): Record { const gatewayUrl = runtimeEnv[INFERENCE_GATEWAY_URL_ENV_VAR_NAME]?.trim(); + const servedKeyNames = parseInferenceGatewayKeys( + runtimeEnv[INFERENCE_GATEWAY_KEYS_ENV_VAR_NAME], + ); - if (!gatewayUrl) { + if (!gatewayUrl || servedKeyNames.length === 0) { return providerConfig; } - const providerIds = new Set( - modelIds.flatMap((modelId) => { - const providerId = modelId?.trim().split('/')[0]; + // Rebase exactly the providers whose keys the gateway is serving (the same + // set dequeue withheld from the sandbox), so the withheld set and the + // rebased set stay identical regardless of which models reference them. + const gatewayProviders = new Map( + servedKeyNames.flatMap((keyName) => { + const provider = getInferenceGatewayProviderByEnvVarName(keyName); - return providerId ? [providerId] : []; + return provider ? [[provider.id, provider] as const] : []; }), ); let merged = providerConfig; - for (const providerId of providerIds) { - const gatewayProvider = getInferenceGatewayProvider(providerId); - - if (!gatewayProvider) { - continue; - } - + for (const [providerId, gatewayProvider] of gatewayProviders) { // ChatGPT-subscription OAuth authenticates through opencode's Codex // plugin directly; leave the openai provider on its default base URL // when a subscription record is present. @@ -1226,7 +1225,6 @@ function resolveModelBackedOpenCodeConfig( configuredModelIds, ), runtimeEnv, - configuredModelIds, ); return { diff --git a/apps/worker/src/run-task/env.ts b/apps/worker/src/run-task/env.ts index 689060e79..45e5d5bde 100644 --- a/apps/worker/src/run-task/env.ts +++ b/apps/worker/src/run-task/env.ts @@ -19,7 +19,11 @@ const BLOCKED_HARNESS_ENV_KEYS = new Set([ 'PREVIEW_PROXY_SUBDOMAIN_SUFFIX', ]); const MODEL_RUNTIME_ENV_KEYS = [ - 'R_INFERENCE_GATEWAY_URL', + // The gateway URL is built worker-side (run-task) from the container- + // reachable platform URL, not delivered by dequeue. The served-keys list is + // dequeue-delivered and read from the raw dequeue env in run-task; it is + // allowlisted here so it survives into the harness env for config rebasing. + 'R_INFERENCE_GATEWAY_KEYS', 'R_MODEL', 'R_SMALL_MODEL', 'R_VISION_MODEL', diff --git a/apps/worker/src/run-task/run-task.ts b/apps/worker/src/run-task/run-task.ts index 248f362df..e587da4e7 100644 --- a/apps/worker/src/run-task/run-task.ts +++ b/apps/worker/src/run-task/run-task.ts @@ -1,6 +1,10 @@ import { type CommunicationProvider, type AcpRequestUserInputAnswers, + buildInferenceGatewayUrl, + INFERENCE_GATEWAY_KEYS_ENV_VAR_NAME, + INFERENCE_GATEWAY_URL_ENV_VAR_NAME, + parseInferenceGatewayKeys, RunStatus, TaskPayloadKind, type QueuedCommunicationMessage, @@ -640,6 +644,34 @@ export const runTask = async ({ ROOMOTE_PROOF_BROWSER_TARGET: environmentConfig.initialUrl, }), }; + // Inference gateway: dequeue advertises the served provider keys by name + // (R_INFERENCE_GATEWAY_KEYS) rather than the gateway URL, so the URL is + // built here from the worker's own platform URL — already rewritten to be + // container-reachable per compute provider (Docker host.docker.internal, + // etc.). The served keys are stripped from the harness env even though + // dequeue withheld them, because buildOpenCodeHarnessEnv re-adds provider + // keys from the worker daemon's process env (present on sandboxes spawned + // before the flag was enabled); this keeps the withheld set out of the + // harness env and env.sh regardless of daemon state. + const inferenceGatewayServedKeys = parseInferenceGatewayKeys( + unsanitizedEnv[INFERENCE_GATEWAY_KEYS_ENV_VAR_NAME], + ); + + if (inferenceGatewayServedKeys.length > 0) { + for (const servedKey of inferenceGatewayServedKeys) { + delete runtimeEnv[servedKey]; + } + + runtimeEnv[INFERENCE_GATEWAY_KEYS_ENV_VAR_NAME] = + inferenceGatewayServedKeys.join(','); + runtimeEnv[INFERENCE_GATEWAY_URL_ENV_VAR_NAME] = buildInferenceGatewayUrl( + workerEnv.trpcUrl, + ); + } else { + delete runtimeEnv[INFERENCE_GATEWAY_KEYS_ENV_VAR_NAME]; + delete runtimeEnv[INFERENCE_GATEWAY_URL_ENV_VAR_NAME]; + } + const workerHomeDir = runtimeEnv.HOME ?? sanitizedEnv.HOME ?? ''; if (workerHomeDir) { diff --git a/packages/compute-providers/src/worker-env/base.ts b/packages/compute-providers/src/worker-env/base.ts index da42b0582..40267d6c3 100644 --- a/packages/compute-providers/src/worker-env/base.ts +++ b/packages/compute-providers/src/worker-env/base.ts @@ -84,9 +84,11 @@ function buildOperatorModelProviderEnv( } // When the inference gateway is enabled, gateway-covered provider keys - // stay on the control plane: the per-task dequeue env routes the harness - // through the gateway instead. Keys the gateway cannot serve yet (Bedrock, - // Vertex) still ship with the worker daemon env. + // (OpenRouter, Anthropic, OpenAI, Gemini, the aggregators, Bedrock) stay off + // the worker daemon process env so they never appear in the sandbox's + // /proc; the per-task dequeue env routes the harness through the gateway + // instead. Keys the gateway cannot serve (Vertex signing, ChatGPT + // subscription OAuth) are not covered and still ship with the daemon env. const gatewayCoveredKeys = new Set(INFERENCE_GATEWAY_PROVIDER_ENV_VAR_NAMES); for (const key of getOperatorModelProviderEnvKeys()) { diff --git a/packages/db/src/lib/model-runtime-config.test.ts b/packages/db/src/lib/model-runtime-config.test.ts index 7506f5c3a..ab5d34478 100644 --- a/packages/db/src/lib/model-runtime-config.test.ts +++ b/packages/db/src/lib/model-runtime-config.test.ts @@ -628,45 +628,52 @@ describe('resolveEffectiveModelRuntimeEnv', () => { }); }); - it('replaces gateway-covered keys with the gateway URL when enabled', async () => { + it('withholds gateway-served keys and advertises them by name when enabled', async () => { const env = await resolveEffectiveModelRuntimeEnv({ inferenceGateway: true, - runtimeEnv: { TRPC_URL: 'https://api.roomote.example.com' }, + runtimeEnv: {}, deploymentEnvVars: { ANTHROPIC_API_KEY: 'sk-anthropic' }, }); expect(env).not.toHaveProperty('ANTHROPIC_API_KEY'); - expect(env.R_INFERENCE_GATEWAY_URL).toBe( - 'https://api.roomote.example.com/api/inference', - ); + // The worker builds the gateway URL from its own platform URL; the + // resolver only advertises which keys it is serving. + expect(env).not.toHaveProperty('R_INFERENCE_GATEWAY_URL'); + expect(env.R_INFERENCE_GATEWAY_KEYS).toBe('ANTHROPIC_API_KEY'); }); it('keeps raw keys when the gateway option is not passed', async () => { const env = await resolveEffectiveModelRuntimeEnv({ - runtimeEnv: { TRPC_URL: 'https://api.roomote.example.com' }, + runtimeEnv: {}, deploymentEnvVars: { ANTHROPIC_API_KEY: 'sk-anthropic' }, }); expect(env.ANTHROPIC_API_KEY).toBe('sk-anthropic'); - expect(env).not.toHaveProperty('R_INFERENCE_GATEWAY_URL'); + expect(env).not.toHaveProperty('R_INFERENCE_GATEWAY_KEYS'); }); - it('falls back to raw keys when no platform API URL is available', async () => { + it('advertises only configured covered keys, not every set covered key', async () => { + // OpenAI is not a configured role model, so even though OPENAI_API_KEY is + // set it is neither served nor advertised — the dequeue redaction keys + // off the advertised set, so a stray OPENAI_API_KEY the deployment set + // for its own code survives into the sandbox. const env = await resolveEffectiveModelRuntimeEnv({ inferenceGateway: true, runtimeEnv: {}, - deploymentEnvVars: { ANTHROPIC_API_KEY: 'sk-anthropic' }, + deploymentEnvVars: { + ANTHROPIC_API_KEY: 'sk-anthropic', + OPENAI_API_KEY: 'sk-openai-for-user-code', + }, }); - expect(env.ANTHROPIC_API_KEY).toBe('sk-anthropic'); - expect(env).not.toHaveProperty('R_INFERENCE_GATEWAY_URL'); + expect(env).not.toHaveProperty('ANTHROPIC_API_KEY'); + expect(env.R_INFERENCE_GATEWAY_KEYS).toBe('ANTHROPIC_API_KEY'); }); - it('keeps keys the gateway does not cover', async () => { + it('withholds Bedrock but keeps Vertex, which the gateway cannot serve', async () => { const env = await resolveEffectiveModelRuntimeEnv({ inferenceGateway: true, runtimeEnv: { - TRPC_URL: 'https://api.roomote.example.com', R_MODEL_ENV_KEYS: 'ANTHROPIC_API_KEY,AWS_BEARER_TOKEN_BEDROCK,GOOGLE_APPLICATION_CREDENTIALS', }, @@ -680,12 +687,12 @@ describe('resolveEffectiveModelRuntimeEnv', () => { expect(env).not.toHaveProperty('ANTHROPIC_API_KEY'); expect(env).not.toHaveProperty('AWS_BEARER_TOKEN_BEDROCK'); // Vertex needs request signing, not header injection, so its - // credentials still flow until the gateway can serve it. + // credentials still flow. expect(env.GOOGLE_APPLICATION_CREDENTIALS).toBe( '{"type":"service_account"}', ); - expect(env.R_INFERENCE_GATEWAY_URL).toBe( - 'https://api.roomote.example.com/api/inference', + expect(env.R_INFERENCE_GATEWAY_KEYS).toBe( + 'ANTHROPIC_API_KEY,AWS_BEARER_TOKEN_BEDROCK', ); }); }); diff --git a/packages/db/src/lib/model-runtime-config.ts b/packages/db/src/lib/model-runtime-config.ts index b992caacb..876eb8939 100644 --- a/packages/db/src/lib/model-runtime-config.ts +++ b/packages/db/src/lib/model-runtime-config.ts @@ -1,13 +1,12 @@ import { eq } from 'drizzle-orm'; import { - buildInferenceGatewayUrl, CHATGPT_OPENCODE_PROVIDER_ID, DEFAULT_MODEL_ROLE_REASONING_EFFORTS, getModelProviderEnvKeyCandidates, getTaskModelCatalog, - INFERENCE_GATEWAY_PROVIDER_ENV_VAR_NAMES, - INFERENCE_GATEWAY_URL_ENV_VAR_NAME, + INFERENCE_GATEWAY_KEYS_ENV_VAR_NAME, isConfiguredEnvValue, + isInferenceGatewayCoveredEnvVar, normalizeDeploymentModelConfig, normalizeOptionalReasoningEffort, parseModelProviderEnvKeys, @@ -175,12 +174,12 @@ export async function resolveEffectiveModelRuntimeEnv( deploymentEnvVars?: Record; executor?: DatabaseOrTransaction; /** - * Route sandbox inference through the gateway: provider keys the gateway - * covers stay on the control plane and `R_INFERENCE_GATEWAY_URL` is - * emitted instead. Callers resolving env for the sandbox pass the - * evaluated InferenceGateway feature flag; control-plane callers - * (routing, titles, summaries) omit it because that inference holds no - * run token to present to the gateway. + * Route sandbox inference through the gateway: the configured provider + * keys the gateway can serve stay on the control plane and their names + * are advertised via `R_INFERENCE_GATEWAY_KEYS` instead. Callers + * resolving env for the sandbox pass the evaluated InferenceGateway + * feature flag; control-plane callers (routing, titles, summaries) omit it + * because that inference holds no run token to present to the gateway. */ inferenceGateway?: boolean; } = {}, @@ -295,23 +294,23 @@ export async function resolveEffectiveModelRuntimeEnv( resolvedRoomotePlanningModel, ], }); - // Gateway-covered provider keys stay on the control plane, and the sandbox - // gets the gateway URL instead. Keys the gateway cannot serve yet (Bedrock, - // Vertex) still flow through. Requires a sandbox-reachable platform API - // URL; when none is configured, fall back to direct keys rather than - // breaking inference. - const gatewayPlatformApiUrl = normalizeConfiguredValue(runtimeEnv.TRPC_URL); - const inferenceGatewayUrl = - options.inferenceGateway && gatewayPlatformApiUrl - ? buildInferenceGatewayUrl(gatewayPlatformApiUrl) - : undefined; - const gatewayCoveredKeyNames = new Set( - INFERENCE_GATEWAY_PROVIDER_ENV_VAR_NAMES, - ); + // When the gateway is active, the configured provider keys it can serve + // (OpenRouter, Anthropic, OpenAI, Gemini, the aggregators, Bedrock) stay on + // the control plane and are advertised to the worker by name via + // R_INFERENCE_GATEWAY_KEYS; the worker builds the (container-reachable) + // gateway URL from its own platform URL and rebases exactly these providers. + // Only configured keys are withheld, so a provider key the deployment set + // for its own code (not used by any configured model) still reaches the + // sandbox. Providers the gateway cannot serve (Vertex signing, ChatGPT + // subscription OAuth) are never covered and always flow through. + const gatewayServedKeyNames = options.inferenceGateway + ? providerKeyNames.filter((name) => isInferenceGatewayCoveredEnvVar(name)) + : []; + const gatewayServedKeyNameSet = new Set(gatewayServedKeyNames); const resolvedProviderKeyValues = Object.fromEntries( providerKeyNames.flatMap((envVarName) => { - if (inferenceGatewayUrl && gatewayCoveredKeyNames.has(envVarName)) { + if (gatewayServedKeyNameSet.has(envVarName)) { return []; } @@ -394,8 +393,8 @@ export async function resolveEffectiveModelRuntimeEnv( R_MODEL_ENV_KEYS: providerKeyNames.join(','), }), ...resolvedProviderKeyValues, - ...(inferenceGatewayUrl && { - [INFERENCE_GATEWAY_URL_ENV_VAR_NAME]: inferenceGatewayUrl, + ...(gatewayServedKeyNames.length > 0 && { + [INFERENCE_GATEWAY_KEYS_ENV_VAR_NAME]: gatewayServedKeyNames.join(','), }), ...(injectedOpenCodeAuthContent && { OPENCODE_AUTH_CONTENT: injectedOpenCodeAuthContent, diff --git a/packages/feature-flags/src/evaluator.ts b/packages/feature-flags/src/evaluator.ts index 65c18ad75..0875ea3be 100644 --- a/packages/feature-flags/src/evaluator.ts +++ b/packages/feature-flags/src/evaluator.ts @@ -124,3 +124,30 @@ export function getFeatureFlagEvaluator(redis: Redis): FeatureFlagEvaluator { export function resetFeatureFlagEvaluatorForTests(): void { evaluatorInstance = null; } + +/** + * Evaluate the deployment-level InferenceGateway flag, failing closed + * (returns false: direct provider keys, the long-standing behavior) when flag + * evaluation is unavailable so a Redis outage cannot break task inference. + * Shared by the dequeue env resolver and every worker-spawn env builder so + * the two halves of the sandbox env are never gated by divergent logic. + */ +export async function isInferenceGatewayEnabledForDeployment( + redis: Redis, + logPrefix = '[inference-gateway]', +): Promise { + try { + return await getFeatureFlagEvaluator(redis).evaluate( + FeatureFlag.InferenceGateway, + { isDeploymentContext: true }, + ); + } catch (error) { + console.warn( + `${logPrefix} InferenceGateway flag evaluation failed; using direct provider keys: ${ + error instanceof Error ? error.message : String(error) + }`, + ); + + return false; + } +} diff --git a/packages/feature-flags/src/server/index.ts b/packages/feature-flags/src/server/index.ts index b07e1ba31..b9f73e106 100644 --- a/packages/feature-flags/src/server/index.ts +++ b/packages/feature-flags/src/server/index.ts @@ -8,6 +8,7 @@ export { FeatureFlagEvaluator, getFeatureFlagEvaluator, + isInferenceGatewayEnabledForDeployment, resetFeatureFlagEvaluatorForTests, } from '../evaluator'; export { MetadataCache } from '../cache'; diff --git a/packages/sdk/src/server/lib/task-runs/__tests__/dequeue-helpers.test.ts b/packages/sdk/src/server/lib/task-runs/__tests__/dequeue-helpers.test.ts index 3048b6d78..6def4b6e3 100644 --- a/packages/sdk/src/server/lib/task-runs/__tests__/dequeue-helpers.test.ts +++ b/packages/sdk/src/server/lib/task-runs/__tests__/dequeue-helpers.test.ts @@ -11,6 +11,7 @@ const { mockResolveEffectiveModelRuntimeEnv, mockTaskRunsFindFirst, mockNotifySourceRunOnSettle, + mockIsInferenceGatewayEnabledForDeployment, } = vi.hoisted(() => ({ mockDecryptSecrets: vi.fn(), mockEnvironmentVariablesFindMany: vi.fn(), @@ -21,6 +22,9 @@ const { mockResolveEffectiveModelRuntimeEnv: vi.fn(), mockTaskRunsFindFirst: vi.fn(), mockNotifySourceRunOnSettle: vi.fn(), + mockIsInferenceGatewayEnabledForDeployment: vi.fn( + async (_redis?: unknown, _prefix?: string) => false, + ), })); vi.mock('@roomote/db/encryption', () => ({ @@ -88,6 +92,15 @@ vi.mock('@roomote/cloud-agents/server', () => ({ releaseTaskRun: vi.fn(), })); +vi.mock('@roomote/feature-flags/server', () => ({ + isInferenceGatewayEnabledForDeployment: (redis: unknown, prefix?: string) => + mockIsInferenceGatewayEnabledForDeployment(redis, prefix), +})); + +vi.mock('@roomote/redis', () => ({ + getRedis: vi.fn(() => ({})), +})); + vi.mock('../notify-source-run-on-settle', () => ({ notifySourceRunOnSettle: (...args: unknown[]) => mockNotifySourceRunOnSettle(...args), @@ -580,4 +593,39 @@ describe('fetchResolvedRuntimeEnvVars', () => { expect(envVars.ROOMOTE_MODEL).toBe('operator/explicit'); expect(envVars.R_MODEL).toBe('anthropic/claude-test'); }); + + it('strips only the gateway-advertised keys from the merged sandbox env', async () => { + // The resolver advertises exactly ANTHROPIC_API_KEY as served; the raw + // deployment env vars are spread in first, so redaction must remove that + // key while leaving an unrelated provider key the user set for their code. + mockResolveEffectiveModelRuntimeEnv.mockResolvedValueOnce({ + R_MODEL: 'anthropic/claude-test', + R_INFERENCE_GATEWAY_KEYS: 'ANTHROPIC_API_KEY', + }); + + const envVars = await fetchResolvedRuntimeEnvVars({ + ANTHROPIC_API_KEY: 'sk-ant', + OPENAI_API_KEY: 'sk-openai-for-user-code', + MY_APP_CONFIG: 'value', + }); + + expect(envVars).not.toHaveProperty('ANTHROPIC_API_KEY'); + expect(envVars.OPENAI_API_KEY).toBe('sk-openai-for-user-code'); + expect(envVars.MY_APP_CONFIG).toBe('value'); + expect(envVars.R_INFERENCE_GATEWAY_KEYS).toBe('ANTHROPIC_API_KEY'); + }); + + it('leaves the merged env untouched when the gateway is off', async () => { + mockResolveEffectiveModelRuntimeEnv.mockResolvedValueOnce({ + R_MODEL: 'anthropic/claude-test', + ANTHROPIC_API_KEY: 'sk-ant', + }); + + const envVars = await fetchResolvedRuntimeEnvVars({ + ANTHROPIC_API_KEY: 'sk-ant', + }); + + expect(envVars.ANTHROPIC_API_KEY).toBe('sk-ant'); + expect(envVars).not.toHaveProperty('R_INFERENCE_GATEWAY_KEYS'); + }); }); diff --git a/packages/sdk/src/server/lib/task-runs/__tests__/fetch-snapshot-env.test.ts b/packages/sdk/src/server/lib/task-runs/__tests__/fetch-snapshot-env.test.ts index b899e295f..b48f67639 100644 --- a/packages/sdk/src/server/lib/task-runs/__tests__/fetch-snapshot-env.test.ts +++ b/packages/sdk/src/server/lib/task-runs/__tests__/fetch-snapshot-env.test.ts @@ -4,13 +4,11 @@ import type { AuthTokenContext, RunTokenContext } from '@roomote/types'; // (vi.mock calls are hoisted above all other code by vitest). const { mockFindFirst, - mockTransaction, - mockFetchEnvVars, + mockFetchResolvedRuntimeEnvVars, mockCreateSourceControlTokenForTaskRun, } = vi.hoisted(() => ({ mockFindFirst: vi.fn(), - mockTransaction: vi.fn(), - mockFetchEnvVars: vi.fn(), + mockFetchResolvedRuntimeEnvVars: vi.fn(), mockCreateSourceControlTokenForTaskRun: vi.fn(), })); @@ -21,7 +19,7 @@ vi.mock('@roomote/db/server', () => ({ findFirst: mockFindFirst, }, }, - transaction: mockTransaction, + update: () => ({ set: () => ({ where: async () => undefined }) }), }, taskRuns: { id: 'taskRuns.id', @@ -31,7 +29,7 @@ vi.mock('@roomote/db/server', () => ({ })); vi.mock('../dequeue-helpers', () => ({ - fetchEnvVars: mockFetchEnvVars, + fetchResolvedRuntimeEnvVars: mockFetchResolvedRuntimeEnvVars, createSourceControlTokenForTaskRun: mockCreateSourceControlTokenForTaskRun, })); @@ -64,11 +62,6 @@ function makeGitHubToken(token: string) { describe('fetchSnapshotEnv', () => { beforeEach(() => { vi.clearAllMocks(); - - // By default, db.transaction calls the callback with a fake tx. - mockTransaction.mockImplementation(async (cb: (tx: unknown) => unknown) => - cb({ fake: 'tx' }), - ); }); // ── Happy path: deployment-scoped env vars ─────────────────────────── @@ -82,7 +75,9 @@ describe('fetchSnapshotEnv', () => { const taskRun = makeTaskRun(); mockFindFirst.mockResolvedValue(taskRun); - mockFetchEnvVars.mockResolvedValue({ MY_SECRET: 'value123' }); + mockFetchResolvedRuntimeEnvVars.mockResolvedValue({ + MY_SECRET: 'value123', + }); const token = makeGitHubToken('ghs_token_abc'); mockCreateSourceControlTokenForTaskRun.mockResolvedValue(token); @@ -98,12 +93,11 @@ describe('fetchSnapshotEnv', () => { // Verify db.query was called. expect(mockFindFirst).toHaveBeenCalledOnce(); - // Verify fetchEnvVars was called inside a transaction. - expect(mockTransaction).toHaveBeenCalledOnce(); - expect(mockFetchEnvVars).toHaveBeenCalledWith( - { fake: 'tx' }, - { sourceControlProvider: 'github' }, - ); + // Verify the gateway-aware resolution was used (so snapshot env withholds + // gateway-served provider keys, like the task dequeue path). + expect(mockFetchResolvedRuntimeEnvVars).toHaveBeenCalledWith(undefined, { + sourceControlProvider: 'github', + }); // Verify createSourceControlTokenForTaskRun was called with the task run. expect(mockCreateSourceControlTokenForTaskRun).toHaveBeenCalledWith( @@ -125,7 +119,7 @@ describe('fetchSnapshotEnv', () => { const taskRun = makeTaskRun(); mockFindFirst.mockResolvedValue(taskRun); - mockFetchEnvVars.mockResolvedValue({}); + mockFetchResolvedRuntimeEnvVars.mockResolvedValue({}); const token = makeGitHubToken('ghs_job_token'); mockCreateSourceControlTokenForTaskRun.mockResolvedValue(token); @@ -138,10 +132,9 @@ describe('fetchSnapshotEnv', () => { taskId: 'task_123', }); - expect(mockFetchEnvVars).toHaveBeenCalledWith( - { fake: 'tx' }, - { sourceControlProvider: 'github' }, - ); + expect(mockFetchResolvedRuntimeEnvVars).toHaveBeenCalledWith(undefined, { + sourceControlProvider: 'github', + }); }); // ── Task run not found ────────────────────────────────────────────── @@ -160,7 +153,7 @@ describe('fetchSnapshotEnv', () => { ); // Should not have called downstream helpers. - expect(mockTransaction).not.toHaveBeenCalled(); + expect(mockFetchResolvedRuntimeEnvVars).not.toHaveBeenCalled(); expect(mockCreateSourceControlTokenForTaskRun).not.toHaveBeenCalled(); }); @@ -174,7 +167,7 @@ describe('fetchSnapshotEnv', () => { }; mockFindFirst.mockResolvedValue(makeTaskRun()); - mockFetchEnvVars.mockResolvedValue({}); + mockFetchResolvedRuntimeEnvVars.mockResolvedValue({}); mockCreateSourceControlTokenForTaskRun.mockResolvedValue( makeGitHubToken('ghs_token_xyz'), ); @@ -196,7 +189,7 @@ describe('fetchSnapshotEnv', () => { }; mockFindFirst.mockResolvedValue(makeTaskRun()); - mockFetchEnvVars.mockResolvedValue({ KEY: 'val' }); + mockFetchResolvedRuntimeEnvVars.mockResolvedValue({ KEY: 'val' }); mockCreateSourceControlTokenForTaskRun.mockResolvedValue(null); await expect(fetchSnapshotEnv(auth, { runId: 42 })).rejects.toThrow( diff --git a/packages/sdk/src/server/lib/task-runs/dequeue-helpers.ts b/packages/sdk/src/server/lib/task-runs/dequeue-helpers.ts index 0875fbadb..23d8fb290 100644 --- a/packages/sdk/src/server/lib/task-runs/dequeue-helpers.ts +++ b/packages/sdk/src/server/lib/task-runs/dequeue-helpers.ts @@ -1,7 +1,7 @@ import { CONTROL_PLANE_ENV_VAR_NAMES, - INFERENCE_GATEWAY_PROVIDER_ENV_VAR_NAMES, - INFERENCE_GATEWAY_URL_ENV_VAR_NAME, + INFERENCE_GATEWAY_KEYS_ENV_VAR_NAME, + parseInferenceGatewayKeys, RunStatus, buildSourceControlTokenMetadata, getSourceControlProviderLabel, @@ -24,10 +24,7 @@ import { sql, } from '@roomote/db/server'; import { decryptSecrets } from '@roomote/db/encryption'; -import { - FeatureFlag, - getFeatureFlagEvaluator, -} from '@roomote/feature-flags/server'; +import { isInferenceGatewayEnabledForDeployment } from '@roomote/feature-flags/server'; import { getRedis } from '@roomote/redis'; import { createTaskRunWorkerGitHubToken } from '@roomote/github'; import { createTaskRunScopedGitLabTokens } from '@roomote/gitlab'; @@ -215,7 +212,10 @@ export async function fetchResolvedRuntimeEnvVars( deploymentEnvVars ?? (await loadPersistedDeploymentEnvVarsFromDb()); const resolvedModelRuntimeEnv = await resolveEffectiveModelRuntimeEnv({ deploymentEnvVars: envVars, - inferenceGateway: await isInferenceGatewayFlagEnabled(), + inferenceGateway: await isInferenceGatewayEnabledForDeployment( + getRedis(), + '[fetchResolvedRuntimeEnvVars]', + ), }); return redactControlPlaneEnvVars( @@ -232,44 +232,28 @@ export async function fetchResolvedRuntimeEnvVars( } /** - * Evaluate the deployment-level InferenceGateway feature flag. Fails closed - * (direct provider keys, the long-standing behavior) if flag evaluation is - * unavailable, so a Redis outage cannot break task inference. - */ -async function isInferenceGatewayFlagEnabled(): Promise { - try { - return await getFeatureFlagEvaluator(getRedis()).evaluate( - FeatureFlag.InferenceGateway, - { isDeploymentContext: true }, - ); - } catch (error) { - console.warn( - `[fetchResolvedRuntimeEnvVars] InferenceGateway flag evaluation failed; using direct provider keys: ${ - error instanceof Error ? error.message : String(error) - }`, - ); - - return false; - } -} - -/** - * When the inference gateway is active (the resolver emitted a gateway URL), - * provider keys the gateway covers stay on the control plane. The raw - * deployment env vars are spread into the sandbox env above, so the keys - * must be stripped from the merged result, not just left unresolved. + * When the gateway is active the resolver emits R_INFERENCE_GATEWAY_KEYS + * listing exactly the configured provider keys it is serving. Those raw keys + * are spread into the sandbox env above (from the deployment env vars), so + * strip precisely that set from the merged result — never a broader set, so a + * provider key the deployment set for its own code (not gateway-served) still + * reaches the sandbox. */ function redactInferenceGatewayProviderKeys( envVars: Record, ): Record { - if (!envVars[INFERENCE_GATEWAY_URL_ENV_VAR_NAME]) { + const servedKeyNames = parseInferenceGatewayKeys( + envVars[INFERENCE_GATEWAY_KEYS_ENV_VAR_NAME], + ); + + if (servedKeyNames.length === 0) { return envVars; } - const coveredKeyNames = new Set(INFERENCE_GATEWAY_PROVIDER_ENV_VAR_NAMES); + const servedKeyNameSet = new Set(servedKeyNames); return Object.fromEntries( - Object.entries(envVars).filter(([key]) => !coveredKeyNames.has(key)), + Object.entries(envVars).filter(([key]) => !servedKeyNameSet.has(key)), ); } diff --git a/packages/sdk/src/server/lib/task-runs/fetch-snapshot-env.ts b/packages/sdk/src/server/lib/task-runs/fetch-snapshot-env.ts index 64371d157..f2c6186df 100644 --- a/packages/sdk/src/server/lib/task-runs/fetch-snapshot-env.ts +++ b/packages/sdk/src/server/lib/task-runs/fetch-snapshot-env.ts @@ -7,7 +7,7 @@ import { import { db, taskRuns, eq } from '@roomote/db/server'; import { - fetchEnvVars, + fetchResolvedRuntimeEnvVars, createSourceControlTokenForTaskRun, } from './dequeue-helpers'; @@ -36,12 +36,14 @@ export async function fetchSnapshotEnv( throw new Error(`${tag} Task run not found: ${runId}`); } - const envVars = await db.transaction(async (tx) => { - return fetchEnvVars(tx, { - sourceControlProvider: resolveSourceControlProviderFromPayload( - taskRun.payload, - ), - }); + // Resolve the snapshot env through the same gateway-aware path as task + // dequeue so gateway-covered provider keys are withheld here too; otherwise + // a snapshot taken with the flag on would bake raw provider keys into the + // snapshot's shell env and the persisted image. + const envVars = await fetchResolvedRuntimeEnvVars(undefined, { + sourceControlProvider: resolveSourceControlProviderFromPayload( + taskRun.payload, + ), }); const sourceControlToken = await createSourceControlTokenForTaskRun( diff --git a/packages/types/src/inference-gateway.ts b/packages/types/src/inference-gateway.ts index 32685479b..317c78d0e 100644 --- a/packages/types/src/inference-gateway.ts +++ b/packages/types/src/inference-gateway.ts @@ -8,6 +8,16 @@ import type { SetupModelProviderId } from './model-provider-config'; */ export const INFERENCE_GATEWAY_URL_ENV_VAR_NAME = 'R_INFERENCE_GATEWAY_URL'; +/** + * Sandbox-facing env var carrying the comma-separated provider key env-var + * names the gateway is serving for this run (the configured, gateway-covered + * keys withheld from the sandbox at dequeue). It is the single authoritative + * signal the worker uses to (1) strip exactly those keys from the harness + * env and (2) rebase exactly those providers onto the gateway, keeping the + * withheld set and the rebased set identical. + */ +export const INFERENCE_GATEWAY_KEYS_ENV_VAR_NAME = 'R_INFERENCE_GATEWAY_KEYS'; + /** * OpenCode registers Bedrock models under this provider id (an * Anthropic-compatible endpoint), distinct from the `amazon-bedrock` setup @@ -19,6 +29,9 @@ export const BEDROCK_MANTLE_OPENCODE_PROVIDER_ID = 'bedrock-mantle'; export const INFERENCE_GATEWAY_REGION_PATTERN = /^[a-z]{2}(?:-[a-z0-9]+)+-\d+$/u; +/** Default AWS region for the Bedrock Mantle Anthropic-compatible endpoint. */ +export const DEFAULT_BEDROCK_MANTLE_REGION = 'us-east-1'; + interface InferenceGatewayAuthHeader { name: string; scheme?: 'bearer'; @@ -131,7 +144,10 @@ export const INFERENCE_GATEWAY_PROVIDERS: readonly InferenceGatewayProvider[] = { id: 'google', name: 'Google Gemini', - envVarNames: ['GEMINI_API_KEY', 'GOOGLE_GENERATIVE_AI_API_KEY'], + // GOOGLE_GENERATIVE_AI_API_KEY first: it is @ai-sdk/google's native env + // var, so when both are set the pre-gateway sandbox used it. GEMINI_API_KEY + // is Roomote's setup-catalog alias and the common single-key case. + envVarNames: ['GOOGLE_GENERATIVE_AI_API_KEY', 'GEMINI_API_KEY'], upstreamBaseUrl: 'https://generativelanguage.googleapis.com', authHeader: { name: 'x-goog-api-key' }, allowedPaths: ['/v1beta/models', '/v1/models'], @@ -225,7 +241,10 @@ export const INFERENCE_GATEWAY_PROVIDERS: readonly InferenceGatewayProvider[] = name: 'Amazon Bedrock', envVarNames: ['AWS_BEARER_TOKEN_BEDROCK'], upstreamBaseUrl: 'https://bedrock-mantle.{region}.api.aws/anthropic', - region: { envVarName: 'AWS_REGION', default: 'us-east-1' }, + region: { + envVarName: 'AWS_REGION', + default: DEFAULT_BEDROCK_MANTLE_REGION, + }, authHeader: { name: 'x-api-key' }, allowedPaths: ANTHROPIC_COMPATIBLE_INFERENCE_PATHS, openCodeBaseUrlSuffix: '/v1', @@ -248,6 +267,38 @@ export function getInferenceGatewayProvider( ); } +const INFERENCE_GATEWAY_PROVIDER_ENV_VAR_NAME_SET = new Set( + INFERENCE_GATEWAY_PROVIDER_ENV_VAR_NAMES, +); + +/** True when `envVarName` is a provider key the gateway can serve. */ +export function isInferenceGatewayCoveredEnvVar(envVarName: string): boolean { + return INFERENCE_GATEWAY_PROVIDER_ENV_VAR_NAME_SET.has(envVarName); +} + +/** + * Resolve the gateway provider that owns a given key env-var name. Used by the + * worker to map the withheld/served key list back to the providers it must + * rebase onto the gateway. + */ +export function getInferenceGatewayProviderByEnvVarName( + envVarName: string, +): InferenceGatewayProvider | undefined { + return INFERENCE_GATEWAY_PROVIDERS.find((provider) => + provider.envVarNames.includes(envVarName), + ); +} + +/** Parse the comma-separated `R_INFERENCE_GATEWAY_KEYS` value into a list. */ +export function parseInferenceGatewayKeys( + value: string | undefined | null, +): string[] { + return (value ?? '') + .split(',') + .map((entry) => entry.trim()) + .filter((entry) => entry.length > 0); +} + /** * The OpenCode `baseURL` override pointing a provider's SDK at the gateway: * `/`. From 91d8b2eddafe20b3ebf13cb50c5a17bff22b1043 Mon Sep 17 00:00:00 2001 From: Matt Rubens <2600+mrubens@users.noreply.github.com> Date: Thu, 16 Jul 2026 19:26:08 -0400 Subject: [PATCH 6/7] Mock the inference-gateway flag helper in controller spawn tests The helper now calls getRedis(); without Redis (CI) the real connection floods ioredis unhandled-error events and fails the modal/daytona/blaxel spawn suites. Mock it to a fixed value like the other spawn dependencies. --- .../compute-providers/__tests__/spawn-blaxel-worker.test.ts | 4 ++++ .../compute-providers/__tests__/spawn-daytona-worker.test.ts | 4 ++++ .../compute-providers/__tests__/spawn-modal-worker.test.ts | 4 ++++ 3 files changed, 12 insertions(+) diff --git a/apps/controller/src/compute-providers/__tests__/spawn-blaxel-worker.test.ts b/apps/controller/src/compute-providers/__tests__/spawn-blaxel-worker.test.ts index 27236df6d..09089342b 100644 --- a/apps/controller/src/compute-providers/__tests__/spawn-blaxel-worker.test.ts +++ b/apps/controller/src/compute-providers/__tests__/spawn-blaxel-worker.test.ts @@ -16,6 +16,10 @@ const mockCreateComputeProviderClient = vi.fn((_arg?: unknown) => ({ })); const mockFindTask = vi.fn(); +vi.mock('../inference-gateway-flag', () => ({ + isInferenceGatewayEnabledForWorkerEnv: vi.fn(async () => false), +})); + vi.mock('@roomote/db/server', async (importOriginal) => { const actual = await importOriginal(); return { diff --git a/apps/controller/src/compute-providers/__tests__/spawn-daytona-worker.test.ts b/apps/controller/src/compute-providers/__tests__/spawn-daytona-worker.test.ts index ef5ba0def..9f5251f9f 100644 --- a/apps/controller/src/compute-providers/__tests__/spawn-daytona-worker.test.ts +++ b/apps/controller/src/compute-providers/__tests__/spawn-daytona-worker.test.ts @@ -15,6 +15,10 @@ const mockGetNamedPortsForTaskRun = vi.fn(); const mockPrimeEnvironmentOidcForMachine = vi.fn(); const mockFindTask = vi.fn(); +vi.mock('../inference-gateway-flag', () => ({ + isInferenceGatewayEnabledForWorkerEnv: vi.fn(async () => false), +})); + vi.mock('@roomote/db/server', async (importOriginal) => { const actual = await importOriginal(); diff --git a/apps/controller/src/compute-providers/__tests__/spawn-modal-worker.test.ts b/apps/controller/src/compute-providers/__tests__/spawn-modal-worker.test.ts index 0363b9105..31580c007 100644 --- a/apps/controller/src/compute-providers/__tests__/spawn-modal-worker.test.ts +++ b/apps/controller/src/compute-providers/__tests__/spawn-modal-worker.test.ts @@ -35,6 +35,10 @@ function mockTaskRun( } as unknown as TaskRun; } +vi.mock('../inference-gateway-flag', () => ({ + isInferenceGatewayEnabledForWorkerEnv: vi.fn(async () => false), +})); + vi.mock('@roomote/db/server', async (importOriginal) => { const actual = await importOriginal(); From 401bb0882684570f0a7f13ae4f4a6d15e5daf1b1 Mon Sep 17 00:00:00 2001 From: Matt Rubens <2600+mrubens@users.noreply.github.com> Date: Thu, 16 Jul 2026 19:44:33 -0400 Subject: [PATCH 7/7] Replace trailing-slash regex with a linear strip (CodeQL ReDoS) CodeQL flagged the anchored `/\/+$/` in the gateway URL builders as a polynomial regular expression on uncontrolled input. Strip trailing slashes with a bounded char-code loop instead; add unit coverage for the URL builders and key lookups. --- .../src/__tests__/inference-gateway.test.ts | 69 +++++++++++++++++++ packages/types/src/inference-gateway.ts | 18 ++++- 2 files changed, 85 insertions(+), 2 deletions(-) create mode 100644 packages/types/src/__tests__/inference-gateway.test.ts diff --git a/packages/types/src/__tests__/inference-gateway.test.ts b/packages/types/src/__tests__/inference-gateway.test.ts new file mode 100644 index 000000000..740404108 --- /dev/null +++ b/packages/types/src/__tests__/inference-gateway.test.ts @@ -0,0 +1,69 @@ +import { + buildInferenceGatewayOpenCodeBaseUrl, + buildInferenceGatewayUrl, + getInferenceGatewayProvider, + getInferenceGatewayProviderByEnvVarName, + INFERENCE_GATEWAY_PROVIDER_ENV_VAR_NAMES, + isInferenceGatewayCoveredEnvVar, + parseInferenceGatewayKeys, +} from '../inference-gateway'; + +describe('inference gateway URL builders', () => { + it('appends the gateway path to a platform URL', () => { + expect(buildInferenceGatewayUrl('https://api.example.com')).toBe( + 'https://api.example.com/api/inference', + ); + }); + + it('strips trailing slashes before appending the gateway path', () => { + expect(buildInferenceGatewayUrl('https://api.example.com///')).toBe( + 'https://api.example.com/api/inference', + ); + }); + + it('builds the OpenCode provider base URL with the provider id and suffix', () => { + const provider = getInferenceGatewayProvider('anthropic'); + expect(provider).toBeDefined(); + expect( + buildInferenceGatewayOpenCodeBaseUrl( + 'https://api.example.com/api/inference/', + provider!, + ), + ).toBe('https://api.example.com/api/inference/anthropic/v1'); + }); +}); + +describe('inference gateway key lookups', () => { + it('maps each covered env var name back to its provider', () => { + for (const envVarName of INFERENCE_GATEWAY_PROVIDER_ENV_VAR_NAMES) { + const provider = getInferenceGatewayProviderByEnvVarName(envVarName); + expect(provider?.envVarNames).toContain(envVarName); + expect(isInferenceGatewayCoveredEnvVar(envVarName)).toBe(true); + } + }); + + it('resolves the Gemini alias to the google provider', () => { + expect(getInferenceGatewayProviderByEnvVarName('GEMINI_API_KEY')?.id).toBe( + 'google', + ); + expect( + getInferenceGatewayProviderByEnvVarName('GOOGLE_GENERATIVE_AI_API_KEY') + ?.id, + ).toBe('google'); + }); + + it('does not cover Vertex or unknown keys', () => { + expect( + isInferenceGatewayCoveredEnvVar('GOOGLE_APPLICATION_CREDENTIALS'), + ).toBe(false); + expect(isInferenceGatewayCoveredEnvVar('SOME_OTHER_KEY')).toBe(false); + }); + + it('parses a comma-separated served-keys value', () => { + expect( + parseInferenceGatewayKeys('ANTHROPIC_API_KEY, OPENROUTER_API_KEY'), + ).toEqual(['ANTHROPIC_API_KEY', 'OPENROUTER_API_KEY']); + expect(parseInferenceGatewayKeys('')).toEqual([]); + expect(parseInferenceGatewayKeys(undefined)).toEqual([]); + }); +}); diff --git a/packages/types/src/inference-gateway.ts b/packages/types/src/inference-gateway.ts index 317c78d0e..a1aa4a75a 100644 --- a/packages/types/src/inference-gateway.ts +++ b/packages/types/src/inference-gateway.ts @@ -299,6 +299,20 @@ export function parseInferenceGatewayKeys( .filter((entry) => entry.length > 0); } +/** + * Strip trailing slashes without a backtracking-prone anchored regex + * (`/\/+$/` trips CodeQL's polynomial-ReDoS heuristic on untrusted input). + */ +function stripTrailingSlashes(value: string): string { + let end = value.length; + + while (end > 0 && value.charCodeAt(end - 1) === 47 /* '/' */) { + end -= 1; + } + + return value.slice(0, end); +} + /** * The OpenCode `baseURL` override pointing a provider's SDK at the gateway: * `/`. @@ -307,7 +321,7 @@ export function buildInferenceGatewayOpenCodeBaseUrl( gatewayUrl: string, provider: InferenceGatewayProvider, ): string { - return `${gatewayUrl.replace(/\/+$/, '')}/${provider.id}${provider.openCodeBaseUrlSuffix}`; + return `${stripTrailingSlashes(gatewayUrl)}/${provider.id}${provider.openCodeBaseUrlSuffix}`; } /** @@ -315,5 +329,5 @@ export function buildInferenceGatewayOpenCodeBaseUrl( * platform API URL (the same base workers use for tRPC). */ export function buildInferenceGatewayUrl(platformApiUrl: string): string { - return `${platformApiUrl.replace(/\/+$/, '')}/api/inference`; + return `${stripTrailingSlashes(platformApiUrl)}/api/inference`; }