feat: disable pkce flow during plugin load (aws#2153)

Author: bywang56

server/aws-lsp-codewhisperer/src/language-server/agenticChat/errors.ts

@@ -11,6 +11,7 @@ type AgenticChatErrorCode =
     | 'MCPServerInitTimeout' // mcp server failed to start within allowed time
     | 'MCPToolExecTimeout' // mcp tool call failed to complete within allowed time
     | 'MCPServerConnectionFailed' // mcp server failed to connect
+    | 'MCPServerAuthFailed' // mcp server failed to complete auth flow
     | 'RequestAborted' // request was aborted by the user
     | 'RequestThrottled' // request was aborted by the user
 

server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/mcp/mcpManager.ts

@@ -42,6 +42,10 @@ import { OAuthClient } from './mcpOauthClient'
 
 export const MCP_SERVER_STATUS_CHANGED = 'mcpServerStatusChanged'
 export const AGENT_TOOLS_CHANGED = 'agentToolsChanged'
+export enum AuthIntent {
+    Interactive = 'interactive',
+    Silent = 'silent',
+}
 
 /**
  * Manages MCP servers and their tools
@@ -166,7 +170,7 @@ export class McpManager {
     /**
      * Load configurations and initialize each enabled server.
      */
-    private async discoverAllServers(): Promise<void> {
+    private async discoverAllServers(authIntent: AuthIntent = AuthIntent.Silent): Promise<void> {
         // Load agent config
         const result = await loadAgentConfig(this.features.workspace, this.features.logging, this.agentPaths)
 
@@ -217,7 +221,7 @@ export class McpManager {
             // Process servers in batches
             for (let i = 0; i < totalServers; i += MAX_CONCURRENT_SERVERS) {
                 const batch = serversToInit.slice(i, i + MAX_CONCURRENT_SERVERS)
-                const batchPromises = batch.map(([name, cfg]) => this.initOneServer(name, cfg))
+                const batchPromises = batch.map(([name, cfg]) => this.initOneServer(name, cfg, authIntent))
 
                 this.features.logging.debug(
                     `MCP: initializing batch of ${batch.length} servers (${i + 1}-${Math.min(i + MAX_CONCURRENT_SERVERS, totalServers)} of ${totalServers})`
@@ -292,7 +296,11 @@ export class McpManager {
      * Start a server process, connect client, and register its tools.
      * Errors are logged but do not stop discovery of other servers.
      */
-    private async initOneServer(serverName: string, cfg: MCPServerConfig): Promise<void> {
+    private async initOneServer(
+        serverName: string,
+        cfg: MCPServerConfig,
+        authIntent: AuthIntent = AuthIntent.Silent
+    ): Promise<void> {
         const DEFAULT_SERVER_INIT_TIMEOUT_MS = 60_000
         this.setState(serverName, McpServerStatus.INITIALIZING, 0)
 
@@ -365,10 +373,19 @@ export class McpManager {
 
                         if (needsOAuth) {
                             OAuthClient.initialize(this.features.workspace, this.features.logging)
-                            const bearer = await OAuthClient.getValidAccessToken(base)
+                            const bearer = await OAuthClient.getValidAccessToken(base, {
+                                interactive: authIntent === 'interactive',
+                            })
                             // add authorization header if we are able to obtain a bearer token
                             if (bearer) {
                                 headers = { ...headers, Authorization: `Bearer ${bearer}` }
+                            } else if (authIntent === 'silent') {
+                                // In silent mode we never launch a browser. If we cannot obtain a token
+                                // from cache/refresh, surface a clear auth-required error and stop here.
+                                throw new AgenticChatError(
+                                    `MCP: server '${serverName}' requires OAuth. Open "Edit MCP Server" and save to sign in.`,
+                                    'MCPServerAuthFailed'
+                                )
                             }
                         }
 
@@ -707,7 +724,7 @@ export class McpManager {
             await saveAgentConfig(this.features.workspace, this.features.logging, this.agentConfig, agentPath)
 
             // Add server tools to tools list after initialization
-            await this.initOneServer(sanitizedName, newCfg)
+            await this.initOneServer(sanitizedName, newCfg, AuthIntent.Interactive)
         } catch (err) {
             this.features.logging.error(
                 `Failed to add MCP server '${serverName}': ${err instanceof Error ? err.message : String(err)}`
@@ -872,7 +889,7 @@ export class McpManager {
                 this.setState(serverName, McpServerStatus.DISABLED, 0)
                 this.emitToolsChanged(serverName)
             } else {
-                await this.initOneServer(serverName, newCfg)
+                await this.initOneServer(serverName, newCfg, AuthIntent.Interactive)
             }
         } catch (err) {
             this.handleError(serverName, err)
@@ -1086,7 +1103,7 @@ export class McpManager {
                 this.setState(serverName, McpServerStatus.DISABLED, 0)
             } else {
                 if (!this.clients.has(serverName) && serverName !== 'Built-in') {
-                    await this.initOneServer(serverName, this.mcpServers.get(serverName)!)
+                    await this.initOneServer(serverName, this.mcpServers.get(serverName)!, AuthIntent.Silent)
                 }
             }
 

server/aws-lsp-codewhisperer/src/language-server/agenticChat/tools/mcp/mcpOauthClient.ts

@@ -42,15 +42,55 @@ export class OAuthClient {
 
     /**
      * Return a valid Bearer token, reusing cache or refresh-token if possible,
-     * otherwise driving one interactive PKCE flow.
+     * otherwise (when interactive) driving one PKCE flow that may launch a browser.
      */
-    public static async getValidAccessToken(mcpBase: URL): Promise<string> {
+    public static async getValidAccessToken(
+        mcpBase: URL,
+        opts: { interactive?: boolean } = { interactive: true }
+    ): Promise<string | undefined> {
+        const interactive = opts?.interactive !== false
         const key = this.computeKey(mcpBase)
         const regPath = path.join(this.cacheDir, `${key}.registration.json`)
         const tokPath = path.join(this.cacheDir, `${key}.token.json`)
 
+        // ===== Silent branch: try cached token, then refresh, never opens a browser =====
+        if (!interactive) {
+            // 1) cached access token
+            const cachedTok = await this.read<Token>(tokPath)
+            if (cachedTok) {
+                const expiry = cachedTok.obtained_at + cachedTok.expires_in * 1000
+                if (Date.now() < expiry) {
+                    this.logger.info(`OAuth: using still-valid cached token (silent)`)
+                    return cachedTok.access_token
+                }
+                this.logger.info(`OAuth: cached token expired → try refresh (silent)`)
+            }
+
+            // 2) refresh-token grant (if we have registration and refresh token)
+            const savedReg = await this.read<Registration>(regPath)
+            if (cachedTok?.refresh_token && savedReg) {
+                try {
+                    const meta = await this.discoverAS(mcpBase)
+                    const refreshed = await this.refreshGrant(meta, savedReg, mcpBase, cachedTok.refresh_token)
+                    if (refreshed) {
+                        await this.write(tokPath, refreshed)
+                        this.logger.info(`OAuth: refresh grant succeeded (silent)`)
+                        return refreshed.access_token
+                    }
+                    this.logger.info(`OAuth: refresh grant did not succeed (silent)`)
+                } catch (e) {
+                    this.logger.warn(`OAuth: silent refresh failed — ${e instanceof Error ? e.message : String(e)}`)
+                }
+            }
+
+            // 3) no token in silent mode → caller should surface auth-required UI
+            return undefined
+        }
+
+        // ===== Interactive branch: may open a browser (PKCE) =====
         // 1) Spin up (or reuse) loopback server + redirect URI
-        let server: http.Server, redirectUri: string
+        let server: http.Server | null = null
+        let redirectUri: string
         const savedReg = await this.read<Registration>(regPath)
         if (savedReg) {
             const port = Number(new URL(savedReg.redirect_uri).port)
@@ -75,7 +115,9 @@ export class OAuthClient {
                 }
             }
         } else {
-            ;({ server, redirectUri } = await this.buildCallbackServer())
+            const created = await this.buildCallbackServer()
+            server = created.server
+            redirectUri = created.redirectUri
             this.logger.info(`OAuth: new redirect URI ${redirectUri}`)
         }
 
@@ -85,7 +127,7 @@ export class OAuthClient {
             if (cached) {
                 const expiry = cached.obtained_at + cached.expires_in * 1000
                 if (Date.now() < expiry) {
-                    this.logger.info(`OAuth: using still‑valid cached token`)
+                    this.logger.info(`OAuth: using still-valid cached token`)
                     return cached.access_token
                 }
                 this.logger.info(`OAuth: cached token expired → try refresh`)
@@ -98,6 +140,7 @@ export class OAuthClient {
             } catch (e: any) {
                 throw new Error(`OAuth discovery failed: ${e?.message ?? String(e)}`)
             }
+
             // 4) Register (or reuse) a dynamic client
             const scopes = ['openid', 'offline_access']
             let reg: Registration
@@ -107,7 +150,7 @@ export class OAuthClient {
                 throw new Error(`OAuth client registration failed: ${e?.message ?? String(e)}`)
             }
 
-            // 5) Refresh‑token grant (one shot)
+            // 5) Refresh-token grant (one shot)
             const attemptedRefresh = !!cached?.refresh_token
             if (cached?.refresh_token) {
                 const refreshed = await this.refreshGrant(meta, reg, mcpBase, cached.refresh_token)
@@ -129,7 +172,9 @@ export class OAuthClient {
                 throw new Error(`OAuth authorization (PKCE) failed${suffix}: ${e?.message ?? String(e)}`)
             }
         } finally {
-            await new Promise<void>(res => server.close(() => res()))
+            if (server) {
+                await new Promise<void>(res => server!.close(() => res()))
+            }
         }
     }