diff --git a/src/i18n/messages/en.json b/src/i18n/messages/en.json index 5bcc85bd2a..dfd417eb1b 100644 --- a/src/i18n/messages/en.json +++ b/src/i18n/messages/en.json @@ -3856,6 +3856,8 @@ "apiKeyProviders": "API Key Providers", "compatibleProviders": "API Key Compatible Providers", "testAll": "Test All", + "distributeProxies": "Distribute Proxies", + "distributing": "Distributing...", "testAllOAuth": "Test all OAuth connections", "testAllFree": "Test all Free connections", "testAllApiKey": "Test all API Key connections", @@ -4092,6 +4094,14 @@ "proxySourceProvider": "Provider", "proxySourceKey": "Key", "proxyConfiguredBySource": "Proxy ({source}): {host}", + "proxyOn": "Proxy On", + "proxyOff": "Proxy Off", + "proxyEnabledTitle": "Proxy is enabled for this connection", + "proxyDisabledTitle": "Proxy is disabled for this connection", + "perKeyProxyOn": "Per-key", + "perKeyProxyOff": "Connection", + "perKeyProxyEnabledTitle": "Per-key proxy assignment enabled for this provider", + "perKeyProxyDisabledTitle": "Per-key proxy assignment disabled for this provider", "autoPriority": "Auto: {priority}", "proxy": "Proxy", "retestAuthentication": "Retest authentication", @@ -4100,6 +4110,7 @@ "enableConnection": "Enable connection", "reauthenticateConnection": "Re-authenticate this connection", "proxyConfig": "Proxy config", + "healthcheckFailed": "Failed to run healthcheck", "aliasExistsAlert": "Alias \"{alias}\" already exists. Please use a different model or edit existing alias.", "openRouterAnyModelHint": "OpenRouter supports any model. Add models and create aliases for quick access.", "modelIdFromOpenRouter": "Model ID (from OpenRouter)", @@ -4822,6 +4833,8 @@ "enableThinking": "Enable Thinking", "maxThinkingTokens": "Max Thinking Tokens", "enableProxy": "Enable Proxy", + "perKeyProxyEnabled": "Enable per-key proxy assignment", + "perKeyProxyEnabledDesc": "When enabled, each provider connection can use its own proxy assignment", "proxyUrl": "Proxy URL", "pricingRates": "Pricing Rates Format", "currentPricing": "Current Pricing Overview", diff --git a/src/lib/db/apiKeys.ts b/src/lib/db/apiKeys.ts index 379f868f9a..559b8e8c6a 100644 --- a/src/lib/db/apiKeys.ts +++ b/src/lib/db/apiKeys.ts @@ -66,6 +66,7 @@ interface ApiKeyMetadata { scopes: string[]; isBanned: boolean; keyHash: string | null; + proxyId: string | null; allowedEndpoints: string[]; streamDefaultMode: "legacy" | "json"; disableNonPublicModels: boolean; @@ -95,6 +96,7 @@ interface ApiKeyRow extends JsonRecord { accessSchedule?: unknown; rate_limits?: unknown; rateLimits?: unknown; + proxy_id?: unknown; stream_default_mode?: unknown; streamDefaultMode?: unknown; } @@ -132,6 +134,7 @@ interface ApiKeyView extends JsonRecord { throttleDelayMs?: number | null; rateLimits: RateLimitRule[] | null; scopes: string[]; + proxyId?: string | null; isBanned?: boolean; expiresAt?: string | null; allowedEndpoints: string[]; @@ -171,6 +174,7 @@ const API_KEY_COLUMN_FALLBACKS = [ { name: "rate_limits", definition: "rate_limits TEXT" }, { name: "is_banned", definition: "is_banned INTEGER NOT NULL DEFAULT 0" }, { name: "key_hash", definition: "key_hash TEXT" }, + { name: "proxy_id", definition: "proxy_id TEXT" }, { name: "allowed_endpoints", definition: "allowed_endpoints TEXT" }, { name: "allowed_quotas", definition: "allowed_quotas TEXT NOT NULL DEFAULT '[]'" }, { name: "stream_default_mode", definition: "stream_default_mode TEXT NOT NULL DEFAULT 'legacy'" }, @@ -374,7 +378,7 @@ function getPreparedStatements(db: ApiKeysDbLike): ApiKeysStatements { "SELECT id, expires_at, revoked_at, is_active, is_banned FROM api_keys WHERE key = ? OR key_hash = ?" ); _stmtGetKeyMetadata = db.prepare( - "SELECT id, name, machine_id, allowed_models, allowed_combos, allowed_connections, allowed_quotas, no_log, auto_resolve, is_active, access_schedule, max_requests_per_day, max_requests_per_minute, throttle_delay_ms, max_sessions, revoked_at, expires_at, ip_allowlist, scopes, rate_limits, is_banned, key_hash, allowed_endpoints, stream_default_mode, disable_non_public_models FROM api_keys WHERE key = ? OR key_hash = ?" + "SELECT id, name, machine_id, allowed_models, allowed_combos, allowed_connections, allowed_quotas, no_log, auto_resolve, is_active, access_schedule, max_requests_per_day, max_requests_per_minute, throttle_delay_ms, max_sessions, revoked_at, expires_at, ip_allowlist, scopes, rate_limits, is_banned, key_hash, allowed_endpoints, stream_default_mode, disable_non_public_models, proxy_id FROM api_keys WHERE key = ? OR key_hash = ?" ); _stmtInsertKey = db.prepare( "INSERT INTO api_keys (id, name, key, machine_id, allowed_models, no_log, created_at, key_prefix, key_hash, scopes) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)" @@ -711,6 +715,7 @@ export async function updateApiKeyPermissions( // T08: max concurrent sessions for this key (0 = unlimited) maxSessions?: number | null; scopes?: string[] | null; + proxyId?: string | null; allowedEndpoints?: string[] | null; streamDefaultMode?: "legacy" | "json" | null; disableNonPublicModels?: boolean; @@ -740,6 +745,7 @@ export async function updateApiKeyPermissions( expiresAt: update.expiresAt, maxSessions: (update as { maxSessions?: number | null }).maxSessions, scopes: (update as { scopes?: string[] | null }).scopes, + proxyId: (update as { proxyId?: string | null }).proxyId, allowedEndpoints: (update as { allowedEndpoints?: string[] | null }).allowedEndpoints, streamDefaultMode: (update as { streamDefaultMode?: "legacy" | "json" | null }) .streamDefaultMode, @@ -765,6 +771,7 @@ export async function updateApiKeyPermissions( normalized.expiresAt === undefined && (normalized as Record).maxSessions === undefined && (normalized as Record).scopes === undefined && + (normalized as Record).proxyId === undefined && (normalized as Record).allowedEndpoints === undefined && (normalized as Record).streamDefaultMode === undefined && normalized.disableNonPublicModels === undefined @@ -792,6 +799,7 @@ export async function updateApiKeyPermissions( maxSessions?: number; expiresAt?: string | null; scopes?: string; + proxyId?: string | null; streamDefaultMode?: "legacy" | "json"; disableNonPublicModels?: number; } = { id }; @@ -892,6 +900,13 @@ export async function updateApiKeyPermissions( params.maxSessions = typeof maxSessionsUpdate === "number" ? Math.max(0, maxSessionsUpdate) : 0; } + const proxyIdUpdate = (normalized as Record).proxyId; + if (proxyIdUpdate !== undefined) { + updates.push("proxy_id = @proxyId"); + params.proxyId = + typeof proxyIdUpdate === "string" && proxyIdUpdate.trim() !== "" ? proxyIdUpdate : null; + } + const allowedEndpointsUpdate = (normalized as Record).allowedEndpoints; if (allowedEndpointsUpdate !== undefined) { updates.push("allowed_endpoints = @allowedEndpoints"); @@ -1254,6 +1269,7 @@ export async function getApiKeyMetadata( isBanned: false, keyHash: null, scopes: ["manage"], + proxyId: null, allowedEndpoints: [], streamDefaultMode: "legacy", disableNonPublicModels: false, @@ -1314,6 +1330,10 @@ export async function getApiKeyMetadata( scopes: parseStringList((record as JsonRecord).scopes), isBanned: parseIsBanned(record.is_banned ?? (record as JsonRecord).isBanned), keyHash: (record.key_hash ?? (record as JsonRecord).keyHash) as string | null, + proxyId: + typeof record.proxy_id === "string" && record.proxy_id.trim() !== "" + ? record.proxy_id + : null, allowedEndpoints: parseStringList( (record as JsonRecord).allowed_endpoints ?? (record as JsonRecord).allowedEndpoints ), diff --git a/src/lib/db/migrations/093_proxy_enable_toggles.sql b/src/lib/db/migrations/093_proxy_enable_toggles.sql new file mode 100644 index 0000000000..4ac2d04659 --- /dev/null +++ b/src/lib/db/migrations/093_proxy_enable_toggles.sql @@ -0,0 +1,17 @@ +-- 062_proxy_enable_toggles.sql +-- Add per-connection proxy enable/disable toggle to support the +-- proxy enable/disable feature with auto-selection fallback. +-- +-- New column: +-- proxy_enabled — 1=proxy enabled for this connection, 0=disabled (default 1 for backward compat) +-- +-- The global proxy on/off toggle lives in the settings namespace (key: proxyEnabled). +-- This per-connection column provides finer-grained control at the provider level. + +ALTER TABLE provider_connections ADD COLUMN proxy_enabled INTEGER NOT NULL DEFAULT 1; + +-- Set proxy_enabled=1 for all existing rows to preserve backward compatibility. +-- Before this migration, proxy was implicitly enabled (no toggle existed). +-- This UPDATE runs on the initial migration apply; on re-runs the ALTER TABLE +-- above fails with "duplicate column" and the migration runner skips the rest. +UPDATE provider_connections SET proxy_enabled = 1 WHERE proxy_enabled = 0; diff --git a/src/lib/db/migrations/094_per_key_proxy_toggles.sql b/src/lib/db/migrations/094_per_key_proxy_toggles.sql new file mode 100644 index 0000000000..acbb43a5a5 --- /dev/null +++ b/src/lib/db/migrations/094_per_key_proxy_toggles.sql @@ -0,0 +1,12 @@ +-- 063_per_key_proxy_toggles.sql +-- Add per-connection toggle for per-key proxy assignment. +-- When enabled, each API key under this connection can use its own proxy +-- (from api_keys.proxy_id) instead of the connection-level proxy. +-- +-- New column: +-- per_key_proxy_enabled — 1=allow per-key proxy assignment, 0=disabled (default) +-- +-- The global per-key proxy toggle lives in the settings namespace +-- (key: perKeyProxyEnabled). + +ALTER TABLE provider_connections ADD COLUMN per_key_proxy_enabled INTEGER NOT NULL DEFAULT 0; diff --git a/src/lib/db/providers.ts b/src/lib/db/providers.ts index d50d299c67..28f399db55 100644 --- a/src/lib/db/providers.ts +++ b/src/lib/db/providers.ts @@ -301,6 +301,8 @@ export async function createProviderConnection(data: JsonRecord) { "rateLimitProtection", "group", "maxConcurrent", + "proxyEnabled", + "perKeyProxyEnabled", "quotaWindowThresholds", ]; for (const field of optionalFields) { @@ -348,6 +350,7 @@ function _insertConnectionRow(db: DbLike, conn: JsonRecord) { last_tested, api_key, id_token, provider_specific_data, expires_in, display_name, global_priority, default_model, token_type, consecutive_use_count, rate_limit_protection, last_used_at, "group", max_concurrent, + proxy_enabled, per_key_proxy_enabled, quota_window_thresholds_json, created_at, updated_at ) VALUES ( @@ -359,6 +362,7 @@ function _insertConnectionRow(db: DbLike, conn: JsonRecord) { @lastTested, @apiKey, @idToken, @providerSpecificData, @expiresIn, @displayName, @globalPriority, @defaultModel, @tokenType, @consecutiveUseCount, @rateLimitProtection, @lastUsedAt, @group, @maxConcurrent, + @proxyEnabled, @perKeyProxyEnabled, @quotaWindowThresholdsJson, @createdAt, @updatedAt ) @@ -404,6 +408,8 @@ function _insertConnectionRow(db: DbLike, conn: JsonRecord) { lastUsedAt: conn.lastUsedAt || null, group: conn.group || null, maxConcurrent: conn.maxConcurrent ?? null, + proxyEnabled: conn.proxyEnabled ?? 1, + perKeyProxyEnabled: conn.perKeyProxyEnabled ?? 0, quotaWindowThresholdsJson: serializeQuotaWindowThresholds(conn.quotaWindowThresholds), createdAt: conn.createdAt, updatedAt: conn.updatedAt, @@ -432,6 +438,8 @@ function _updateConnectionRow(db: DbLike, id: string, data: JsonRecord) { "group" = @group, max_concurrent = @maxConcurrent, quota_window_thresholds_json = @quotaWindowThresholdsJson, + proxy_enabled = @proxyEnabled, + per_key_proxy_enabled = @perKeyProxyEnabled, updated_at = @updatedAt WHERE id = @id ` @@ -477,6 +485,18 @@ function _updateConnectionRow(db: DbLike, id: string, data: JsonRecord) { group: data.group || null, maxConcurrent: data.maxConcurrent ?? null, quotaWindowThresholdsJson: serializeQuotaWindowThresholds(data.quotaWindowThresholds), + proxyEnabled: + typeof data.proxyEnabled === "boolean" + ? data.proxyEnabled + ? 1 + : 0 + : (data.proxyEnabled ?? 1), + perKeyProxyEnabled: + typeof data.perKeyProxyEnabled === "boolean" + ? data.perKeyProxyEnabled + ? 1 + : 0 + : (data.perKeyProxyEnabled ?? 0), updatedAt: now, }); } diff --git a/src/lib/db/settings.ts b/src/lib/db/settings.ts index 738645274a..ff5c5f06c8 100644 --- a/src/lib/db/settings.ts +++ b/src/lib/db/settings.ts @@ -123,6 +123,8 @@ export async function getSettings() { // `applyAuthzBypassSection` → `getAuthzBypassSnapshot()`. localOnlyManageScopeBypassEnabled: true, localOnlyManageScopeBypassPrefixes: ["/api/mcp/"], + proxyEnabled: true, + perKeyProxyEnabled: false, }; for (const row of rows) { const record = toRecord(row); @@ -161,6 +163,13 @@ export async function updateSettings(updates: Record) { tx(); backupDbFile("pre-write"); invalidateDbCache("settings"); // Bust the read cache immediately + + // Bust proxy resolution cache when proxy toggle settings change + const PROXY_TOGGLE_KEYS = ["proxyEnabled", "perKeyProxyEnabled"]; + if (Object.keys(updates).some((k) => PROXY_TOGGLE_KEYS.includes(k))) { + bumpProxyConfigGeneration(); + } + const nextSettings = await getSettings(); try { @@ -599,10 +608,11 @@ export async function deleteProxyForLevel(level: string, id: string | null) { return setProxyForLevel(level, id, null); } -export async function resolveProxyForConnection(connectionId: string) { +export async function resolveProxyForConnection(connectionId: string, apiKeyId?: string) { + const cacheKey = apiKeyId ? `${connectionId}:${apiKeyId}` : connectionId; const startGeneration = proxyConfigGeneration; const startRegistryGeneration = getProxyRegistryGeneration(); - const cached = proxyResolutionCache.get(connectionId); + const cached = proxyResolutionCache.get(cacheKey); if ( cached && cached.generation === startGeneration && @@ -611,40 +621,134 @@ export async function resolveProxyForConnection(connectionId: string) { return cached.result; } - const config = await getProxyConfig(); + const db = getDbInstance(); - // Resolve by specificity across both proxy storage backends. The dashboard - // Custom tab still writes account/provider proxies to the legacy config, - // while Saved Proxy writes registry assignments. Do not let a registry-global - // fallback shadow a more-specific legacy account/provider proxy (#2601). - const registryAccount = await resolveProxyForScopeFromRegistry("account", connectionId); - if (registryAccount?.proxy) { - cacheProxyResolution(connectionId, startGeneration, startRegistryGeneration, registryAccount); - return registryAccount; + // Step 1: Check global proxyEnabled setting + // Read only the proxyEnabled key for performance instead of loading all settings. + let globalProxyEnabled = true; + try { + const proxyEnabledRow = db + .prepare("SELECT value FROM key_value WHERE namespace = 'settings' AND key = 'proxyEnabled'") + .get() as { value?: string } | undefined; + if (proxyEnabledRow?.value) { + globalProxyEnabled = JSON.parse(proxyEnabledRow.value) !== false; + } + } catch { + // Default to true on read error } - if (connectionId && config.keys?.[connectionId]) { - const result = { proxy: config.keys[connectionId], level: "key", levelId: connectionId }; - cacheProxyResolution(connectionId, startGeneration, startRegistryGeneration, result); + if (!globalProxyEnabled) { + const result: ProxyResolutionResult = { proxy: null, level: "direct", levelId: null }; + // Do not cache the "direct" result when global toggle is off so that + // toggling it back on takes effect immediately without a generation bump. return result; } - const db = getDbInstance(); + // Step 1.5: Check global perKeyProxyEnabled setting + let globalPerKeyProxyEnabled = false; + try { + const perKeyRow = db + .prepare( + "SELECT value FROM key_value WHERE namespace = 'settings' AND key = 'perKeyProxyEnabled'" + ) + .get() as { value?: string } | undefined; + if (perKeyRow?.value) { + globalPerKeyProxyEnabled = JSON.parse(perKeyRow.value) !== false; + } + } catch { + // Default to false on read error + } + + const config = await getProxyConfig(); + + // Step 2: API key-level proxy (only if per-key proxy is enabled globally or per-connection) + if (apiKeyId) { + // Check if per-key proxy is allowed: globally OR per-connection + let perKeyEnabled = globalPerKeyProxyEnabled; + if (!perKeyEnabled && connectionId) { + try { + const perKeyConn = db + .prepare("SELECT per_key_proxy_enabled FROM provider_connections WHERE id = ?") + .get(connectionId) as { per_key_proxy_enabled?: number } | undefined; + perKeyEnabled = perKeyConn?.per_key_proxy_enabled === 1; + } catch { + // Fall through + } + } + + if (perKeyEnabled) { + try { + const apiKeyRow = db.prepare("SELECT proxy_id FROM api_keys WHERE id = ?").get(apiKeyId) as + | { proxy_id?: string | null } + | undefined; + if (apiKeyRow?.proxy_id) { + const proxyRow = db + .prepare( + "SELECT p.type, p.host, p.port, p.username, p.password FROM proxy_registry p WHERE p.id = ?" + ) + .get(apiKeyRow.proxy_id) as + | { type: string; host: string; port: number; username: string; password: string } + | undefined; + if (proxyRow) { + const result = { + proxy: { + type: proxyRow.type, + host: proxyRow.host, + port: proxyRow.port, + username: proxyRow.username, + password: proxyRow.password, + }, + level: "apiKey" as const, + levelId: apiKeyId, + source: "api_key" as const, + }; + cacheProxyResolution(cacheKey, startGeneration, startRegistryGeneration, result); + return result; + } + } + } catch { + // Fall through to existing resolution + } + } + } + + // Step 3: Account-level registry + const registryAccount = await resolveProxyForScopeFromRegistry("account", connectionId); + if (registryAccount?.proxy) { + cacheProxyResolution(cacheKey, startGeneration, startRegistryGeneration, registryAccount); + return registryAccount; + } + + // Step 4: Legacy key-level + if (connectionId && config.keys?.[connectionId]) { + const result = { proxy: config.keys[connectionId], level: "key", levelId: connectionId }; + cacheProxyResolution(cacheKey, startGeneration, startRegistryGeneration, result); + return result; + } + + // Step 5: Look up the connection's provider and check proxy_enabled const connection = db - .prepare("SELECT provider FROM provider_connections WHERE id = ?") + .prepare("SELECT provider, proxy_enabled FROM provider_connections WHERE id = ?") .get(connectionId); if (connection) { const connectionRecord = toRecord(connection); const provider = typeof connectionRecord.provider === "string" ? connectionRecord.provider : null; + // proxy_enabled defaults to 0 (false) when the column is NULL (pre-migration) + const connProxyEnabled = connectionRecord.proxy_enabled === 1; - if (provider) { + // Step 6: Provider-level registry (only if proxy_enabled) + if (provider && connProxyEnabled) { const registryProvider = await resolveProxyForScopeFromRegistry("provider", provider); - if (registryProvider?.proxy) return registryProvider; + if (registryProvider?.proxy) { + cacheProxyResolution(cacheKey, startGeneration, startRegistryGeneration, registryProvider); + return registryProvider; + } } - if (config.combos && Object.keys(config.combos).length > 0) { + // Step 7: Legacy combo-level (only if proxy_enabled) + if (connProxyEnabled && config.combos && Object.keys(config.combos).length > 0) { const combos = db.prepare("SELECT id, data FROM combos").all(); for (const comboRow of combos) { const comboRecord = toRecord(comboRow); @@ -659,7 +763,9 @@ export async function resolveProxyForConnection(connectionId: string) { (entry) => getComboModelProvider(entry) === provider ); if (usesProvider) { - return { proxy: config.combos[comboId], level: "combo", levelId: comboId }; + const result = { proxy: config.combos[comboId], level: "combo", levelId: comboId }; + cacheProxyResolution(cacheKey, startGeneration, startRegistryGeneration, result); + return result; } } catch { // Ignore malformed combo records during proxy resolution. @@ -668,21 +774,45 @@ export async function resolveProxyForConnection(connectionId: string) { } } - if (provider && config.providers?.[provider]) { - return { + // Step 8: Legacy provider-level (only if proxy_enabled) + if (provider && connProxyEnabled && config.providers?.[provider]) { + const result = { proxy: config.providers[provider], level: "provider", levelId: provider, }; + cacheProxyResolution(cacheKey, startGeneration, startRegistryGeneration, result); + return result; } } + // Step 9: Global registry const registryGlobal = await resolveProxyForScopeFromRegistry("global"); - if (registryGlobal?.proxy) return registryGlobal; - - if (config.global) { - return { proxy: config.global, level: "global", levelId: null }; + if (registryGlobal?.proxy) { + cacheProxyResolution(cacheKey, startGeneration, startRegistryGeneration, registryGlobal); + return registryGlobal; } + + // Step 10: Legacy global + if (config.global) { + const result = { proxy: config.global, level: "global", levelId: null }; + cacheProxyResolution(cacheKey, startGeneration, startRegistryGeneration, result); + return result; + } + + // Step 11: Auto-selection fallback (only when global proxy is enabled) + try { + const { selectWorkingProxyFallback } = await import("@omniroute/open-sse/utils/proxyFallback"); + const fallback = await selectWorkingProxyFallback(connectionId); + if (fallback) { + cacheProxyResolution(cacheKey, startGeneration, startRegistryGeneration, fallback); + return fallback; + } + } catch (err) { + console.warn({ err, connectionId }, "Proxy fallback auto-selection failed"); + } + + // Step 12: Return direct return { proxy: null, level: "direct", levelId: null }; } diff --git a/src/shared/validation/schemas.ts b/src/shared/validation/schemas.ts index 707bf5673d..0f5d503bbf 100644 --- a/src/shared/validation/schemas.ts +++ b/src/shared/validation/schemas.ts @@ -1996,6 +1996,8 @@ export const updateProviderConnectionSchema = z ]) .optional(), projectId: z.union([z.string(), z.null()]).optional(), + proxyEnabled: z.boolean().optional(), + perKeyProxyEnabled: z.boolean().optional(), // Partial patch of per-connection provider-specific settings (e.g. quota toggles) providerSpecificData: z .record(z.string(), z.unknown()) diff --git a/src/shared/validation/settingsSchemas.ts b/src/shared/validation/settingsSchemas.ts index 86504d5241..63aee1fab7 100644 --- a/src/shared/validation/settingsSchemas.ts +++ b/src/shared/validation/settingsSchemas.ts @@ -269,6 +269,8 @@ export const updateSettingsSchema = z.object({ autoRoutingDefaultVariant: z .enum(["lkgp", "coding", "fast", "cheap", "offline", "smart"]) .optional(), + proxyEnabled: z.boolean().optional(), + perKeyProxyEnabled: z.boolean().optional(), // CLIProxyAPI connection settings cliproxyapi_fallback_enabled: z.boolean().optional(), cliproxyapi_url: z.string().url().max(500).optional(), @@ -342,9 +344,7 @@ export const databaseSettingsSchema = z.object( }), // Skip location and stats as they're read-only - }, - { strict: true } -); +}).strict(); export type DatabaseSettingsSchema = z.infer; diff --git a/tests/unit/api-key-lifecycle.test.ts b/tests/unit/api-key-lifecycle.test.ts index 5f7f6b6ed2..126856b56e 100644 --- a/tests/unit/api-key-lifecycle.test.ts +++ b/tests/unit/api-key-lifecycle.test.ts @@ -117,3 +117,22 @@ test("validateApiKey updates last_used_at for persisted keys", async () => { assert.ok(row?.last_used_at, "last_used_at should be set on successful validation"); assert.ok(Date.parse(row.last_used_at) > 0, "last_used_at should be an ISO timestamp"); }); + +test("getApiKeyMetadata returns proxyId for a key with proxy_id set", async () => { + const created = await makeKey("proxy-test", "machine-proxy"); + + const ok = await apiKeysDb.updateApiKeyPermissions(created.id, { proxyId: "test-proxy-id" }); + assert.equal(ok, true); + + const metadata = await apiKeysDb.getApiKeyMetadata(created.key); + assert.ok(metadata); + assert.equal(metadata!.proxyId, "test-proxy-id"); +}); + +test("getApiKeyMetadata returns proxyId null for a key without proxy_id", async () => { + const created = await makeKey("no-proxy-test", "machine-no-proxy"); + + const metadata = await apiKeysDb.getApiKeyMetadata(created.key); + assert.ok(metadata); + assert.equal(metadata!.proxyId, null); +}); diff --git a/tests/unit/db-settings-crud.test.ts b/tests/unit/db-settings-crud.test.ts index b6940d93f5..ff49c58ada 100644 --- a/tests/unit/db-settings-crud.test.ts +++ b/tests/unit/db-settings-crud.test.ts @@ -489,6 +489,9 @@ test("proxy helpers resolve key, provider, global, and direct paths while tolera name: "Proxy Resolution Target", apiKey: "sk-proxy-resolution", }); + db.prepare("UPDATE provider_connections SET proxy_enabled = 1 WHERE id = ?").run( + (connection as any).id + ); await settingsDb.setProxyConfig({ level: "global", @@ -560,6 +563,9 @@ test("proxy resolution skips combos without serialized data and falls back to pr name: "Proxy Null Combo", apiKey: "sk-claude-proxy", }); + db.prepare("UPDATE provider_connections SET proxy_enabled = 1 WHERE id = ?").run( + (connection as any).id + ); await settingsDb.setProxyForLevel("provider", "claude", { type: "https", @@ -592,6 +598,10 @@ test("proxy resolution matches combo proxies through aliased model entries", asy name: "Proxy Alias Combo", apiKey: "sk-claude-alias", }); + // Enable proxy on this connection so legacy combo/provider proxy checks work + core.getDbInstance() + .prepare("UPDATE provider_connections SET proxy_enabled = 1 WHERE id = ?") + .run((connection as any).id); const combo = await combosDb.createCombo({ name: "combo-aliased-model", @@ -633,6 +643,14 @@ test("proxy resolution prefers legacy key and provider proxies over registry glo }); await proxiesDb.assignProxyToScope("global", null, registryGlobal.id); + // Enable proxy on both connections so legacy key/provider proxy checks work + core.getDbInstance() + .prepare("UPDATE provider_connections SET proxy_enabled = 1 WHERE id = ?") + .run((keyConnection as any).id); + core.getDbInstance() + .prepare("UPDATE provider_connections SET proxy_enabled = 1 WHERE id = ?") + .run((providerConnection as any).id); + await settingsDb.setProxyForLevel("key", (keyConnection as any).id, { type: "http", host: "legacy-key-override.local", diff --git a/tests/unit/proxy-registry.test.ts b/tests/unit/proxy-registry.test.ts index 72a98c675d..451312cb67 100644 --- a/tests/unit/proxy-registry.test.ts +++ b/tests/unit/proxy-registry.test.ts @@ -6,16 +6,19 @@ import path from "node:path"; const TEST_DATA_DIR = fs.mkdtempSync(path.join(os.tmpdir(), "omniroute-proxy-registry-")); process.env.DATA_DIR = TEST_DATA_DIR; +process.env.API_KEY_SECRET = "test-secret"; const core = await import("../../src/lib/db/core.ts"); const providersDb = await import("../../src/lib/db/providers.ts"); const proxiesDb = await import("../../src/lib/db/proxies.ts"); const settingsDb = await import("../../src/lib/db/settings.ts"); +const apiKeysDb = await import("../../src/lib/db/apiKeys.ts"); const proxiesRoute = await import("../../src/app/api/settings/proxies/route.ts"); async function resetStorage() { delete process.env.INITIAL_PASSWORD; core.resetDbInstance(); + apiKeysDb.resetApiKeyState(); fs.rmSync(TEST_DATA_DIR, { recursive: true, force: true }); fs.mkdirSync(TEST_DATA_DIR, { recursive: true }); } @@ -161,9 +164,9 @@ test("specific registry account assignment takes precedence over legacy key prox await proxiesDb.assignProxyToScope("account", (conn as any).id, accountProxy.id); const resolved = await settingsDb.resolveProxyForConnection((conn as any).id); - assert.equal(resolved.level, "account"); - assert.equal(resolved.source, "registry"); - assert.equal(resolved.proxy.host, "account.local"); + assert.equal((resolved as any).level, "account"); + assert.equal((resolved as any).source, "registry"); + assert.equal((resolved as any).proxy.host, "account.local"); }); test("legacy proxy config migration imports global/provider/key assignments", async () => { @@ -197,9 +200,9 @@ test("legacy proxy config migration imports global/provider/key assignments", as assert.equal(result.migrated >= 3, true); const resolved = await settingsDb.resolveProxyForConnection((conn as any).id); - assert.equal(resolved.level, "account"); - assert.equal(resolved.source, "registry"); - assert.equal(resolved.proxy.host, "account-legacy.local"); + assert.equal((resolved as any).level, "account"); + assert.equal((resolved as any).source, "registry"); + assert.equal((resolved as any).proxy.host, "account-legacy.local"); }); // #2456: resolveProxyForProvider (used by the OAuth token exchange + token refresh, @@ -291,6 +294,75 @@ test("resolveProxyForProvider returns null when neither registry nor legacy conf assert.equal(resolved, null); }); +test("resolveProxyForConnection uses apiKey proxy before account-level proxy", async () => { + await resetStorage(); + + const conn = await providersDb.createProviderConnection({ + provider: "openai", + authType: "apikey", + name: "api-key-proxy", + apiKey: "sk-apikey-proxy", + }); + + const accountProxy = await proxiesDb.createProxy({ + name: "Account Proxy", + type: "http", + host: "account.local", + port: 8081, + }); + await proxiesDb.assignProxyToScope("account", (conn as any).id, accountProxy.id); + + const key = await apiKeysDb.createApiKey("proxy-test-key", "machine-p1"); + + // Enable per-key proxy globally so the API key's proxy_id is honored + core + .getDbInstance() + .prepare( + "INSERT OR REPLACE INTO key_value (namespace, key, value) VALUES ('settings', 'perKeyProxyEnabled', 'true')" + ) + .run(); + + const apiKeyProxy = await proxiesDb.createProxy({ + name: "API Key Proxy", + type: "https", + host: "apikey.local", + port: 8443, + }); + await apiKeysDb.updateApiKeyPermissions(key.id, { proxyId: apiKeyProxy.id }); + + const resolved = await settingsDb.resolveProxyForConnection((conn as any).id, key.id); + assert.ok(resolved); + assert.equal((resolved as any).level, "apiKey"); + assert.equal((resolved as any).proxy.host, "apikey.local"); + assert.equal((resolved as any).proxy.port, 8443); +}); + +test("resolveProxyForConnection falls through when apiKey has no proxy_id", async () => { + await resetStorage(); + + const conn = await providersDb.createProviderConnection({ + provider: "openai", + authType: "apikey", + name: "fallthrough-test", + apiKey: "sk-fallthrough", + }); + + const accountProxy = await proxiesDb.createProxy({ + name: "Account Proxy", + type: "http", + host: "account-fallthrough.local", + port: 8081, + }); + await proxiesDb.assignProxyToScope("account", (conn as any).id, accountProxy.id); + + const key = await apiKeysDb.createApiKey("fallthrough-key", "machine-ft"); + + const resolved = await settingsDb.resolveProxyForConnection((conn as any).id, key.id); + assert.ok(resolved); + assert.equal((resolved as any).level, "account"); + assert.equal((resolved as any).proxy.host, "account-fallthrough.local"); +}); + test("createProxyRegistrySchema accepts type:vercel and source:vercel-relay (schema gap-06)", async () => { // Note: We validate the schema directly using the worktree's absolute path because // tests run with CWD=/OmniRoute, so `@/` aliases resolve to the main branch's src/.