fix(proxy): use connection proxy for OAuth refresh

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,

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<void>
 ) => {
-  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,

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({