diff --git a/src/lib/usage/providerLimits.ts b/src/lib/usage/providerLimits.ts index ac12edb201..575e8acb50 100644 --- a/src/lib/usage/providerLimits.ts +++ b/src/lib/usage/providerLimits.ts @@ -113,6 +113,7 @@ async function syncToCloudIfEnabled() { async function refreshAndUpdateCredentials(connection: ProviderConnectionLike) { const executor = getExecutor(connection.provider); const credentials = { + connectionId: connection.id, accessToken: connection.accessToken, refreshToken: connection.refreshToken, expiresAt: connection.tokenExpiresAt || connection.expiresAt || null, diff --git a/src/sse/services/tokenRefresh.ts b/src/sse/services/tokenRefresh.ts index 9d35b8dfe7..63415b09b0 100755 --- a/src/sse/services/tokenRefresh.ts +++ b/src/sse/services/tokenRefresh.ts @@ -1,6 +1,10 @@ // Re-export from open-sse with local logger import * as log from "../utils/logger"; -import { updateProviderConnection, resolveProxyForProvider } from "@/lib/localDb"; +import { + updateProviderConnection, + resolveProxyForConnection, + resolveProxyForProvider, +} from "@/lib/localDb"; import { TOKEN_EXPIRY_BUFFER_MS as BUFFER_MS, getRefreshLeadMs as _getRefreshLeadMs, @@ -28,17 +32,28 @@ import { export const TOKEN_EXPIRY_BUFFER_MS = BUFFER_MS; +async function resolveProxyForCredentials(provider: string, credentials?: any) { + if (credentials?.connectionId) { + const resolved = await resolveProxyForConnection(credentials.connectionId); + if (resolved?.proxy) { + return resolved.proxy; + } + } + + return resolveProxyForProvider(provider); +} + export const refreshAccessToken = async ( provider: string, refreshToken: string, credentials: any ) => { - const proxy = await resolveProxyForProvider(provider); + const proxy = await resolveProxyForCredentials(provider, credentials); return _refreshAccessToken(provider, refreshToken, credentials, log, proxy); }; -export const refreshClaudeOAuthToken = async (refreshToken: string) => { - const proxy = await resolveProxyForProvider("claude"); +export const refreshClaudeOAuthToken = async (refreshToken: string, credentials?: any) => { + const proxy = await resolveProxyForCredentials("claude", credentials); return _refreshClaudeOAuthToken(refreshToken, log, proxy); }; @@ -46,34 +61,35 @@ export const refreshGoogleToken = async ( refreshToken: string, clientId: string, clientSecret: string, - provider: string = "gemini" + provider: string = "gemini", + credentials?: any ) => { - const proxy = await resolveProxyForProvider(provider); + const proxy = await resolveProxyForCredentials(provider, credentials); return _refreshGoogleToken(refreshToken, clientId, clientSecret, log, proxy); }; -export const refreshQwenToken = async (refreshToken: string) => { - const proxy = await resolveProxyForProvider("qwen"); +export const refreshQwenToken = async (refreshToken: string, credentials?: any) => { + const proxy = await resolveProxyForCredentials("qwen", credentials); return _refreshQwenToken(refreshToken, log, proxy); }; -export const refreshCodexToken = async (refreshToken: string) => { - const proxy = await resolveProxyForProvider("codex"); +export const refreshCodexToken = async (refreshToken: string, credentials?: any) => { + const proxy = await resolveProxyForCredentials("codex", credentials); return _refreshCodexToken(refreshToken, log, proxy); }; -export const refreshQoderToken = async (refreshToken: string) => { - const proxy = await resolveProxyForProvider("qoder"); +export const refreshQoderToken = async (refreshToken: string, credentials?: any) => { + const proxy = await resolveProxyForCredentials("qoder", credentials); return _refreshQoderToken(refreshToken, log, proxy); }; -export const refreshGitHubToken = async (refreshToken: string) => { - const proxy = await resolveProxyForProvider("github"); +export const refreshGitHubToken = async (refreshToken: string, credentials?: any) => { + const proxy = await resolveProxyForCredentials("github", credentials); return _refreshGitHubToken(refreshToken, log, proxy); }; -export const refreshCopilotToken = async (githubAccessToken: string) => { - const proxy = await resolveProxyForProvider("github"); +export const refreshCopilotToken = async (githubAccessToken: string, credentials?: any) => { + const proxy = await resolveProxyForCredentials("github", credentials); return _refreshCopilotToken(githubAccessToken, log, proxy); }; @@ -82,12 +98,12 @@ export const getAccessToken = async ( credentials: any, onPersist?: (result: any) => Promise ) => { - const proxy = await resolveProxyForProvider(provider); + const proxy = await resolveProxyForCredentials(provider, credentials); return _getAccessToken(provider, credentials, log, proxy, onPersist); }; export const refreshTokenByProvider = async (provider: string, credentials: any) => { - const proxy = await resolveProxyForProvider(provider); + const proxy = await resolveProxyForCredentials(provider, credentials); return _refreshTokenByProvider(provider, credentials, log, proxy); }; @@ -213,7 +229,10 @@ export async function checkAndRefreshToken(provider: string, credentials: any) { expiresIn: Math.round((copilotExpiresAt - now) / 1000), }); - const copilotToken = await refreshCopilotToken(updatedCredentials.accessToken); + const copilotToken = await refreshCopilotToken( + updatedCredentials.accessToken, + updatedCredentials + ); if (copilotToken) { await updateProviderCredentials(updatedCredentials.connectionId, { providerSpecificData: { @@ -239,9 +258,9 @@ export async function checkAndRefreshToken(provider: string, credentials: any) { // Local-specific: Refresh GitHub and Copilot tokens together export async function refreshGitHubAndCopilotTokens(credentials: any) { - const newGitHubCredentials = await refreshGitHubToken(credentials.refreshToken); + const newGitHubCredentials = await refreshGitHubToken(credentials.refreshToken, credentials); if (newGitHubCredentials?.accessToken) { - const copilotToken = await refreshCopilotToken(newGitHubCredentials.accessToken); + const copilotToken = await refreshCopilotToken(newGitHubCredentials.accessToken, credentials); if (copilotToken) { return { ...newGitHubCredentials, diff --git a/tests/unit/token-refresh-route-service.test.ts b/tests/unit/token-refresh-route-service.test.ts index 546f9ca56a..b8d50ad509 100644 --- a/tests/unit/token-refresh-route-service.test.ts +++ b/tests/unit/token-refresh-route-service.test.ts @@ -1,6 +1,8 @@ import test from "node:test"; import assert from "node:assert/strict"; import fs from "node:fs"; +import http from "node:http"; +import net from "node:net"; import os from "node:os"; import path from "node:path"; @@ -10,6 +12,7 @@ process.env.API_KEY_SECRET = "test-api-key-secret"; const core = await import("../../src/lib/db/core.ts"); const providersDb = await import("../../src/lib/db/providers.ts"); +const settingsDb = await import("../../src/lib/db/settings.ts"); const tokenRefresh = await import("../../src/sse/services/tokenRefresh.ts"); const { PROVIDERS, OAUTH_ENDPOINTS } = await import("../../open-sse/config/constants.ts"); @@ -46,6 +49,104 @@ async function withMockedNow(now, fn) { } } +async function withHttpServer(handler, fn) { + const server = http.createServer(handler); + + await new Promise((resolve, reject) => { + server.once("error", reject); + server.listen(0, "127.0.0.1", resolve); + }); + + const address = server.address(); + assert.ok(address && typeof address === "object"); + + try { + return await fn({ + host: "127.0.0.1", + port: address.port, + url: `http://127.0.0.1:${address.port}`, + }); + } finally { + await new Promise((resolve, reject) => { + server.close((error) => { + if (error) reject(error); + else resolve(); + }); + }); + } +} + +async function withConnectProxyServer(fn, options = {}) { + const connectRequests = []; + const server = http.createServer((_req, res) => { + res.writeHead(501); + res.end("CONNECT only"); + }); + + server.on("connect", (req, clientSocket, head) => { + connectRequests.push(String(req.url || "")); + const [host, portText] = String(req.url || "").split(":"); + const targetHost = options.targetHost || host; + const targetPort = Number(options.targetPort || portText || 80); + const upstreamSocket = net.connect(targetPort, targetHost, () => { + clientSocket.write("HTTP/1.1 200 Connection Established\r\n\r\n"); + if (head && head.length > 0) { + upstreamSocket.write(head); + } + upstreamSocket.pipe(clientSocket); + clientSocket.pipe(upstreamSocket); + }); + + const closeSockets = () => { + upstreamSocket.destroy(); + clientSocket.destroy(); + }; + + upstreamSocket.on("error", closeSockets); + clientSocket.on("error", closeSockets); + }); + + await new Promise((resolve, reject) => { + server.once("error", reject); + server.listen(0, "127.0.0.1", resolve); + }); + + const address = server.address(); + assert.ok(address && typeof address === "object"); + + try { + return await fn({ + host: "127.0.0.1", + port: address.port, + url: `http://127.0.0.1:${address.port}`, + connectRequests, + }); + } finally { + await new Promise((resolve, reject) => { + server.close((error) => { + if (error) reject(error); + else resolve(); + }); + }); + } +} + +async function withPatchedProvider(providerId, config, fn) { + const hadOwnConfig = Object.prototype.hasOwnProperty.call(PROVIDERS, providerId); + const previousConfig = hadOwnConfig ? PROVIDERS[providerId] : undefined; + PROVIDERS[providerId] = config; + + try { + return await fn(); + } finally { + if (hadOwnConfig) { + PROVIDERS[providerId] = previousConfig; + } else { + delete PROVIDERS[providerId]; + } + } +} + test.beforeEach(async () => { await resetStorage(); }); @@ -206,6 +307,132 @@ test("updateProviderCredentials persists rotated tokens and returns false for mi assert.equal(missing, false); }); +test("OAuth refresh prefers connection proxy over provider proxy", async () => { + const providerId = "custom-oauth-connection-proxy"; + const refreshRequests = []; + + await withHttpServer( + (req, res) => { + refreshRequests.push({ method: req.method, url: req.url }); + res.writeHead(200, { "content-type": "application/json" }); + res.end( + JSON.stringify({ + access_token: "connection-proxy-access", + refresh_token: "connection-proxy-refresh", + expires_in: 1800, + }) + ); + }, + async (tokenServer) => { + await withConnectProxyServer( + async (accountProxy) => { + await withPatchedProvider( + providerId, + { + tokenUrl: `http://oauth-refresh.example.test:${tokenServer.port}/token`, + clientId: "connection-proxy-client-id", + clientSecret: "connection-proxy-client-secret", + }, + async () => { + const connection = await providersDb.createProviderConnection({ + provider: providerId, + authType: "oauth", + name: "Connection Proxy OAuth", + accessToken: "old-access", + refreshToken: "refresh-via-account-proxy", + }); + + try { + await settingsDb.setProxyForLevel("provider", providerId, { + type: "http", + host: "provider-proxy.invalid", + port: 8080, + }); + await settingsDb.setProxyForLevel("key", (connection as any).id, { + type: "http", + host: accountProxy.host, + port: accountProxy.port, + }); + + const result = await tokenRefresh.getAccessToken(providerId, { + connectionId: (connection as any).id, + refreshToken: "refresh-via-account-proxy", + }); + + assert.equal(result.accessToken, "connection-proxy-access"); + assert.equal(refreshRequests.length, 1); + assert.deepEqual(refreshRequests[0], { method: "POST", url: "/token" }); + assert.deepEqual(accountProxy.connectRequests, [ + `oauth-refresh.example.test:${tokenServer.port}`, + ]); + } finally { + await settingsDb.deleteProxyForLevel("provider", providerId); + await settingsDb.deleteProxyForLevel("key", (connection as any).id); + } + } + ); + }, + { targetHost: tokenServer.host, targetPort: tokenServer.port } + ); + } + ); +}); + +test("provider-specific refresh helper accepts connection proxy context", async () => { + const refreshRequests = []; + + await withHttpServer( + (req, res) => { + refreshRequests.push({ method: req.method, url: req.url }); + res.writeHead(200, { "content-type": "application/json" }); + res.end( + JSON.stringify({ + access_token: "claude-connection-proxy-access", + refresh_token: "claude-connection-proxy-refresh", + expires_in: 1200, + }) + ); + }, + async (tokenServer) => { + await withConnectProxyServer(async (accountProxy) => { + const originalAnthropicTokenUrl = OAUTH_ENDPOINTS.anthropic.token; + OAUTH_ENDPOINTS.anthropic.token = `${tokenServer.url}/token`; + let connectionId: string | undefined; + try { + const connection = await providersDb.createProviderConnection({ + provider: "claude", + authType: "oauth", + name: "Claude Connection Proxy OAuth", + accessToken: "old-access", + refreshToken: "refresh-claude-via-account-proxy", + }); + connectionId = (connection as any).id; + + await settingsDb.setProxyForLevel("key", connectionId, { + type: "http", + host: accountProxy.host, + port: accountProxy.port, + }); + + const result = await tokenRefresh.refreshClaudeOAuthToken( + "refresh-claude-via-account-proxy", + { connectionId } + ); + + assert.equal(result.accessToken, "claude-connection-proxy-access"); + assert.equal(refreshRequests.length, 1); + assert.deepEqual(refreshRequests[0], { method: "POST", url: "/token" }); + } finally { + if (connectionId) { + await settingsDb.deleteProxyForLevel("key", connectionId); + } + OAUTH_ENDPOINTS.anthropic.token = originalAnthropicTokenUrl; + } + }); + } + ); +}); + test("checkAndRefreshToken refreshes expiring OAuth access tokens and updates the connection", async () => { const now = 1_700_000_000_000; const connection = await providersDb.createProviderConnection({