diff --git a/src/main/agent-manager.ts b/src/main/agent-manager.ts index f153409..462cbac 100644 --- a/src/main/agent-manager.ts +++ b/src/main/agent-manager.ts @@ -18,6 +18,7 @@ import { SessionStatusType, MessagePartType, MessageRole } from './adapters/codi import { getTaskApiPort, waitForTaskApiServer } from './task-api-server' import { randomUUID } from 'crypto' import { registerSecretSession, unregisterSecretSession, getSecretBrokerPort, writeSecretShellWrapper } from './secret-broker' +import { registerMcpProxyTarget, getMcpAuthProxyPort } from './mcp-auth-proxy' // Coding agent backend type enum enum CodingAgentType { @@ -336,19 +337,31 @@ export class AgentManager extends EventEmitter { } } - // Inject enterprise JWT for Workflo MCP Dev Server + // Inject enterprise JWT for Workflo MCP Dev Server — route through auth proxy + let finalUrl = mcpServer.url if (mcpServer.name === '[Workflo] MCP Dev Server' && this.enterpriseAuth) { - try { - const jwt = await this.enterpriseAuth.getJwt() - finalHeaders = { ...finalHeaders, Authorization: `Bearer ${jwt}` } - } catch (err) { - console.warn('[AgentManager] Failed to inject enterprise JWT for MCP Dev Server:', err) + const proxyPort = getMcpAuthProxyPort() + const proxyUrl = proxyPort ? registerMcpProxyTarget(mcpServer.url) : null + if (proxyUrl) { + finalUrl = proxyUrl + // Don't send static Authorization — proxy injects fresh JWT per request + delete finalHeaders['Authorization'] + console.log(`[AgentManager] MCP Dev Server routed through auth proxy: ${proxyUrl}`) + } else { + // Fallback: static JWT (proxy not running) + try { + const jwt = await this.enterpriseAuth.getJwt() + finalHeaders = { ...finalHeaders, Authorization: `Bearer ${jwt}` } + console.log('[AgentManager] MCP Dev Server using static JWT (proxy not available)') + } catch (err) { + console.warn('[AgentManager] Failed to inject enterprise JWT for MCP Dev Server:', err) + } } } result[mcpServer.name] = { type: 'http', - url: mcpServer.url, + url: finalUrl, headers: finalHeaders } } @@ -380,7 +393,6 @@ export class AgentManager extends EventEmitter { } } - // result keys logged at debug level only return result } diff --git a/src/main/index.ts b/src/main/index.ts index 34021ea..578cfea 100644 --- a/src/main/index.ts +++ b/src/main/index.ts @@ -29,6 +29,7 @@ import { EnterpriseHeartbeat } from './enterprise-heartbeat' import { EnterpriseStateSync } from './enterprise-state-sync' import { setTaskApiNotifier, setTranscriptProvider, stopTaskApiServer } from './task-api-server' import { startSecretBroker, stopSecretBroker, writeSecretShellWrapper } from './secret-broker' +import { startMcpAuthProxy, stopMcpAuthProxy } from './mcp-auth-proxy' import { startMobileApiServer, stopMobileApiServer, broadcastToMobileClients, setMobileApiNotifier } from './mobile-api-server' import { registerUpdaterIpc, initAutoUpdater, isUpdateDownloaded, getPendingVersion } from './auto-updater' import { initCrashLogger } from './crash-logger' @@ -66,6 +67,7 @@ async function shutdownAppServices(): Promise { mcpToolCaller?.destroy() oauthManager?.destroy() stopSecretBroker() + stopMcpAuthProxy() stopMobileApiServer() stopTaskApiServer() @@ -734,6 +736,16 @@ app.whenReady().then(async () => { // Wire enterprise auth into agent manager so it can inject JWT into MCP Dev Server requests agentManager.setEnterpriseAuth(enterpriseAuth) + // Start MCP auth proxy so agent sessions get auto-refreshing JWT + // on every MCP tool call (ACP/OpenCode/Claude Code never refresh + // MCP headers mid-session — the proxy solves this transparently). + try { + const proxyPort = await startMcpAuthProxy(enterpriseAuth) + console.log(`[Main] MCP auth proxy started on port ${proxyPort}`) + } catch (proxyErr) { + console.warn('[Main] MCP auth proxy failed to start (MCP JWT will use static headers):', proxyErr) + } + syncManager.setEnterpriseConnection(apiClient, enterpriseSyncMgr, session.userId, enterpriseStateSyncInstance) console.log('[Main] Enterprise connection restored on startup (with heartbeat)') diff --git a/src/main/ipc-handlers.ts b/src/main/ipc-handlers.ts index 382f95e..a6a6ddc 100644 --- a/src/main/ipc-handlers.ts +++ b/src/main/ipc-handlers.ts @@ -875,6 +875,15 @@ export function registerIpcHandlers( // Pass enterprise auth to agent manager so it can inject JWT into MCP Dev Server requests agentManager.setEnterpriseAuth(enterpriseAuth!) + // Start MCP auth proxy so agent sessions get auto-refreshing JWT + try { + const { startMcpAuthProxy } = await import('./mcp-auth-proxy') + const proxyPort = await startMcpAuthProxy(enterpriseAuth!) + console.log(`[enterprise] MCP auth proxy started on port ${proxyPort}`) + } catch (proxyErr) { + console.warn('[enterprise] MCP auth proxy failed to start:', proxyErr) + } + // Run initial sync (agents, skills, MCP servers) const syncStart = Date.now() console.log('[enterprise] Running initial resource sync...') diff --git a/src/main/mcp-auth-proxy.test.ts b/src/main/mcp-auth-proxy.test.ts new file mode 100644 index 0000000..4484a28 --- /dev/null +++ b/src/main/mcp-auth-proxy.test.ts @@ -0,0 +1,263 @@ +import { beforeEach, afterEach, describe, expect, it, vi } from 'vitest' +import { createServer, type Server as HttpServer, type IncomingMessage, type ServerResponse } from 'http' + +// Mock electron before importing the module +vi.mock('electron', () => ({ + app: { isPackaged: false } +})) + +// We test the proxy by starting both it and a mock upstream server +import { + startMcpAuthProxy, + stopMcpAuthProxy, + registerMcpProxyTarget, + unregisterMcpProxyTarget, + getMcpAuthProxyPort +} from './mcp-auth-proxy' + +// ── Mock EnterpriseAuth ─────────────────────────────────────────── + +class MockAuth { + private callCount = 0 + private jwtValue = 'jwt-token-1' + private shouldFail = false + + async getJwt(): Promise { + this.callCount++ + if (this.shouldFail) throw new Error('Auth failed') + return this.jwtValue + } + + setJwt(value: string): void { + this.jwtValue = value + } + + setFail(fail: boolean): void { + this.shouldFail = fail + } + + getCallCount(): number { + return this.callCount + } +} + +// ── Mock upstream MCP server ────────────────────────────────────── + +function createMockUpstream(): Promise<{ server: HttpServer; port: number; lastHeaders: () => Record; lastBody: () => string; lastMethod: () => string; lastUrl: () => string }> { + let capturedHeaders: Record = {} + let capturedBody = '' + let capturedMethod = '' + let capturedUrl = '' + + return new Promise((resolve) => { + const srv = createServer((req: IncomingMessage, res: ServerResponse) => { + capturedHeaders = { ...req.headers } + capturedMethod = req.method || '' + capturedUrl = req.url || '' + + const chunks: Buffer[] = [] + req.on('data', (c: Buffer) => chunks.push(c)) + req.on('end', () => { + capturedBody = Buffer.concat(chunks).toString() + res.writeHead(200, { 'Content-Type': 'application/json' }) + res.end(JSON.stringify({ result: 'ok', method: capturedMethod })) + }) + }) + + srv.listen(0, '127.0.0.1', () => { + const addr = srv.address() + const p = typeof addr === 'object' && addr ? addr.port : 0 + resolve({ + server: srv, + port: p, + lastHeaders: () => capturedHeaders, + lastBody: () => capturedBody, + lastMethod: () => capturedMethod, + lastUrl: () => capturedUrl + }) + }) + }) +} + +// ── Tests ───────────────────────────────────────────────────────── + +describe('MCP Auth Proxy', () => { + let mockAuth: MockAuth + let upstream: Awaited> + + beforeEach(async () => { + mockAuth = new MockAuth() + upstream = await createMockUpstream() + }) + + afterEach(async () => { + stopMcpAuthProxy() + await new Promise((resolve) => upstream.server.close(() => resolve())) + }) + + it('starts on a random port and reports it', async () => { + const port = await startMcpAuthProxy(mockAuth as never) + expect(port).toBeGreaterThan(0) + expect(getMcpAuthProxyPort()).toBe(port) + }) + + it('returns existing port on duplicate start', async () => { + const port1 = await startMcpAuthProxy(mockAuth as never) + const port2 = await startMcpAuthProxy(mockAuth as never) + expect(port1).toBe(port2) + }) + + it('registers a target and returns a proxy URL', async () => { + const port = await startMcpAuthProxy(mockAuth as never) + const proxyUrl = registerMcpProxyTarget(`http://127.0.0.1:${upstream.port}`) + expect(proxyUrl).toBe(`http://127.0.0.1:${port}/1`) + }) + + it('deduplicates repeated registrations for the same target URL', async () => { + const port = await startMcpAuthProxy(mockAuth as never) + const targetUrl = `http://127.0.0.1:${upstream.port}` + + const url1 = registerMcpProxyTarget(targetUrl) + const url2 = registerMcpProxyTarget(targetUrl) + const url3 = registerMcpProxyTarget(targetUrl) + + // All calls should return the same proxy URL (same ID) + expect(url1).toBe(`http://127.0.0.1:${port}/1`) + expect(url2).toBe(url1) + expect(url3).toBe(url1) + + // A different target URL should get a new ID + const url4 = registerMcpProxyTarget('http://example.com:9999') + expect(url4).toBe(`http://127.0.0.1:${port}/2`) + }) + + it('returns null from registerMcpProxyTarget when proxy is not running', () => { + const result = registerMcpProxyTarget('http://example.com') + expect(result).toBeNull() + }) + + it('forwards POST requests with fresh JWT to upstream', async () => { + await startMcpAuthProxy(mockAuth as never) + const proxyUrl = registerMcpProxyTarget(`http://127.0.0.1:${upstream.port}`)! + + const body = JSON.stringify({ jsonrpc: '2.0', method: 'tools/list', id: 1 }) + const response = await fetch(proxyUrl, { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body + }) + + expect(response.status).toBe(200) + const data = await response.json() + expect(data).toEqual({ result: 'ok', method: 'POST' }) + + // Verify JWT was injected + expect(upstream.lastHeaders()['authorization']).toBe('Bearer jwt-token-1') + // Verify body was forwarded + expect(upstream.lastBody()).toBe(body) + }) + + it('injects a FRESH JWT on each request (not cached)', async () => { + await startMcpAuthProxy(mockAuth as never) + const proxyUrl = registerMcpProxyTarget(`http://127.0.0.1:${upstream.port}`)! + + // First request with jwt-token-1 + await fetch(proxyUrl, { method: 'POST', body: '{}' }) + expect(upstream.lastHeaders()['authorization']).toBe('Bearer jwt-token-1') + + // Change JWT (simulates token refresh) + mockAuth.setJwt('jwt-token-REFRESHED') + + // Second request should use the new JWT + await fetch(proxyUrl, { method: 'POST', body: '{}' }) + expect(upstream.lastHeaders()['authorization']).toBe('Bearer jwt-token-REFRESHED') + + // getJwt was called twice (once per request) + expect(mockAuth.getCallCount()).toBe(2) + }) + + it('preserves sub-paths from the request', async () => { + await startMcpAuthProxy(mockAuth as never) + const proxyUrl = registerMcpProxyTarget(`http://127.0.0.1:${upstream.port}`)! + + await fetch(`${proxyUrl}/sse`, { method: 'GET' }) + expect(upstream.lastUrl()).toBe('/sse') + expect(upstream.lastMethod()).toBe('GET') + }) + + it('returns 404 for unknown target IDs', async () => { + await startMcpAuthProxy(mockAuth as never) + + const port = getMcpAuthProxyPort()! + const response = await fetch(`http://127.0.0.1:${port}/999`, { method: 'POST', body: '{}' }) + expect(response.status).toBe(404) + }) + + it('returns 502 when auth fails', async () => { + await startMcpAuthProxy(mockAuth as never) + const proxyUrl = registerMcpProxyTarget(`http://127.0.0.1:${upstream.port}`)! + + mockAuth.setFail(true) + + const response = await fetch(proxyUrl, { method: 'POST', body: '{}' }) + expect(response.status).toBe(502) + const data = await response.json() + expect(data.error).toContain('Failed to obtain auth token') + }) + + it('unregisters targets correctly', async () => { + await startMcpAuthProxy(mockAuth as never) + const proxyUrl = registerMcpProxyTarget(`http://127.0.0.1:${upstream.port}`)! + + // Works before unregister + const before = await fetch(proxyUrl, { method: 'POST', body: '{}' }) + expect(before.status).toBe(200) + + // Unregister + unregisterMcpProxyTarget(proxyUrl) + + // 404 after unregister + const after = await fetch(proxyUrl, { method: 'POST', body: '{}' }) + expect(after.status).toBe(404) + }) + + it('cleans up on stop', async () => { + await startMcpAuthProxy(mockAuth as never) + expect(getMcpAuthProxyPort()).not.toBeNull() + + stopMcpAuthProxy() + expect(getMcpAuthProxyPort()).toBeNull() + }) + + it('handles multiple concurrent requests with fresh JWT each', async () => { + let callNum = 0 + // Override getJwt to return sequential tokens + mockAuth.getJwt = async () => `jwt-call-${++callNum}` + + await startMcpAuthProxy(mockAuth as never) + const proxyUrl = registerMcpProxyTarget(`http://127.0.0.1:${upstream.port}`)! + + // Fire 5 concurrent requests + const responses = await Promise.all( + Array.from({ length: 5 }, () => + fetch(proxyUrl, { method: 'POST', body: '{}' }) + ) + ) + + // All should succeed + for (const r of responses) { + expect(r.status).toBe(200) + } + + // getJwt should have been called 5 times (once per request) + expect(callNum).toBe(5) + }) + + it('forwards query parameters to upstream', async () => { + await startMcpAuthProxy(mockAuth as never) + const proxyUrl = registerMcpProxyTarget(`http://127.0.0.1:${upstream.port}`)! + + await fetch(`${proxyUrl}?foo=bar&baz=1`, { method: 'GET' }) + expect(upstream.lastUrl()).toBe('/?foo=bar&baz=1') + }) +}) diff --git a/src/main/mcp-auth-proxy.ts b/src/main/mcp-auth-proxy.ts new file mode 100644 index 0000000..ba9f03a --- /dev/null +++ b/src/main/mcp-auth-proxy.ts @@ -0,0 +1,231 @@ +/** + * MCP Auth Proxy — local HTTP reverse proxy for transparent JWT refresh. + * + * Runs in the Electron main process, listens on 127.0.0.1:{random port}. + * Agent MCP configs point to this proxy instead of the real remote MCP server. + * On every forwarded request the proxy calls enterpriseAuth.getJwt() to inject + * a fresh Authorization header, so tokens are never stale — even in sessions + * that run for hours. + * + * Flow: + * 1. Agent session starts → registerMcpProxyTarget(targetUrl) + * 2. buildMcpServersForAdapter sets MCP url to http://127.0.0.1://… + * 3. MCP client sends request → proxy gets fresh JWT → forwards to real server + * 4. Real server response is piped back to MCP client unchanged + * + * Why: + * ACP, OpenCode, and Claude Code all bake MCP server headers at session start. + * None refresh them mid-session. A 1-hour JWT expires during long agent runs, + * causing persistent 401 errors on every MCP tool call with no recovery path. + * This proxy decouples token lifetime from session lifetime. + */ + +import { createServer, request as httpRequest, type Server as HttpServer, type IncomingMessage, type ServerResponse } from 'http' +import { request as httpsRequest } from 'https' +import type { EnterpriseAuth } from './enterprise-auth' + +let server: HttpServer | null = null +let port: number | null = null +let authRef: EnterpriseAuth | null = null + +// Registration ID → target URL +const targets = new Map() + +// Reverse lookup: target URL → ID (deduplicates repeated registrations) +const targetsByUrl = new Map() + +// Monotonic counter for short, collision-free IDs +let nextId = 1 + +export function getMcpAuthProxyPort(): number | null { + return port +} + +/** + * Register a remote MCP server URL and return a localhost proxy URL + * that transparently injects fresh enterprise JWT on every request. + * + * If the same targetUrl was already registered, returns the existing + * proxy URL (prevents leaking IDs when buildMcpServersForAdapter is + * called on every prompt for Claude Code adapter). + * + * Returns null if the proxy is not running. + */ +export function registerMcpProxyTarget(targetUrl: string): string | null { + if (!port) return null + + // Reuse existing registration for the same target URL + const existingId = targetsByUrl.get(targetUrl) + if (existingId) { + return `http://127.0.0.1:${port}/${existingId}` + } + + const id = String(nextId++) + targets.set(id, targetUrl) + targetsByUrl.set(targetUrl, id) + console.log(`[McpAuthProxy] Registered target ${id} → ${targetUrl}`) + return `http://127.0.0.1:${port}/${id}` +} + +/** + * Remove a previously registered target (cleanup on session end). + */ +export function unregisterMcpProxyTarget(proxyUrl: string): void { + // Extract ID from proxy URL: http://127.0.0.1:/ + try { + const url = new URL(proxyUrl) + const id = url.pathname.split('/')[1] + if (id && targets.has(id)) { + const targetUrl = targets.get(id)! + targets.delete(id) + targetsByUrl.delete(targetUrl) + console.log(`[McpAuthProxy] Unregistered target ${id}`) + } + } catch { + // Ignore malformed URLs + } +} + +export function startMcpAuthProxy(auth: EnterpriseAuth): Promise { + if (server && port) return Promise.resolve(port) + authRef = auth + + return new Promise((resolve, reject) => { + server = createServer(handleRequest) + + // Listen on random available port, loopback only + server.listen(0, '127.0.0.1', () => { + const addr = server!.address() + if (typeof addr === 'object' && addr) { + port = addr.port + console.log(`[McpAuthProxy] Started on port ${port}`) + resolve(port) + } else { + reject(new Error('Failed to get MCP auth proxy address')) + } + }) + + server.on('error', reject) + }) +} + +export function stopMcpAuthProxy(): void { + if (server) { + server.close() + server = null + port = null + } + targets.clear() + targetsByUrl.clear() + authRef = null + nextId = 1 +} + +// ── Request handler ─────────────────────────────────────────────── + +async function handleRequest(req: IncomingMessage, res: ServerResponse): Promise { + try { + // Parse: //optional/sub/path + const reqUrl = new URL(req.url || '/', 'http://localhost') + const pathParts = reqUrl.pathname.split('/').filter(Boolean) + const id = pathParts[0] + + if (!id || !targets.has(id)) { + console.warn(`[McpAuthProxy] 404 — unknown target ID "${id}"`) + + res.writeHead(404, { 'Content-Type': 'text/plain' }) + res.end('Unknown proxy target') + return + } + + // Build the target URL: base + remaining sub-path + query string + const targetBase = targets.get(id)! + const subPathParts = pathParts.slice(1) + const targetUrl = new URL(targetBase) + if (subPathParts.length > 0) { + // Append sub-path: //sse → targetBase + /sse + targetUrl.pathname = targetUrl.pathname.replace(/\/$/, '') + '/' + subPathParts.join('/') + } + // else: no sub-path — use targetBase as-is (no trailing slash added) + targetUrl.search = reqUrl.search + + // Get fresh JWT + if (!authRef) { + res.writeHead(503, { 'Content-Type': 'text/plain' }) + res.end('Auth not available') + return + } + + let jwt: string + try { + jwt = await authRef.getJwt() + } catch (err) { + console.error('[McpAuthProxy] Failed to get JWT:', err) + res.writeHead(502, { 'Content-Type': 'application/json' }) + res.end(JSON.stringify({ error: 'Failed to obtain auth token' })) + return + } + + // Build forwarded headers — copy everything except host + const forwardHeaders: Record = {} + for (const [key, value] of Object.entries(req.headers)) { + if (key.toLowerCase() === 'host') continue + // Skip connection-specific headers + if (key.toLowerCase() === 'connection') continue + forwardHeaders[key] = value + } + forwardHeaders['authorization'] = `Bearer ${jwt}` + forwardHeaders['host'] = targetUrl.host + + // Choose http or https + const isHttps = targetUrl.protocol === 'https:' + const targetString = targetUrl.toString() + const reqOptions = { method: req.method, headers: forwardHeaders } + + const callback = (proxyRes: IncomingMessage) => { + // Forward status + headers back to MCP client + const responseHeaders = { ...proxyRes.headers } + // Remove transfer-encoding — Node handles chunking on the proxy→client leg. + // But preserve it for SSE (text/event-stream) responses so events flush immediately. + const contentType = String(proxyRes.headers['content-type'] || '') + const isSSE = contentType.includes('text/event-stream') + if (!isSSE) { + delete responseHeaders['transfer-encoding'] + } + + res.writeHead(proxyRes.statusCode || 502, responseHeaders) + + if (isSSE) { + // For SSE, flush each chunk immediately so events aren't buffered + proxyRes.on('data', (chunk: Buffer) => { + res.write(chunk) + }) + proxyRes.on('end', () => res.end()) + proxyRes.on('error', () => res.end()) + } else { + proxyRes.pipe(res) + } + } + + const proxyReq = isHttps + ? httpsRequest(targetString, reqOptions, callback) + : httpRequest(targetString, reqOptions, callback) + + proxyReq.on('error', (err) => { + console.error(`[McpAuthProxy] Upstream error for target ${id}:`, err.message) + if (!res.headersSent) { + res.writeHead(502, { 'Content-Type': 'application/json' }) + res.end(JSON.stringify({ error: `Upstream error: ${err.message}` })) + } + }) + + // Pipe request body to upstream + req.pipe(proxyReq) + } catch (err) { + console.error('[McpAuthProxy] Unexpected error:', err) + if (!res.headersSent) { + res.writeHead(500, { 'Content-Type': 'text/plain' }) + res.end('Internal proxy error') + } + } +}