pi-code 1.0.27 → 1.0.28

package/extensions/internal/mcp-oauth.ts

@@ -28,6 +28,17 @@ interface StoredAuth {
   redirectPort?: number
 }
 
+/** Claude's per-server `oauth` config object: a pre-registered client (secret via
+ * MCP_CLIENT_SECRET), a fixed callback port for pre-registered redirect URIs, and
+ * pinned scopes. authServerMetadataUrl is accepted but unsupported (the MCP SDK
+ * offers no discovery override); the connect path warns when it is set. */
+export interface OAuthServerConfig {
+  clientId?: string
+  callbackPort?: number
+  scopes?: string
+  authServerMetadataUrl?: string
+}
+
 /** A server name is config-controlled text; the digest keeps hostile names inside
  * the store directory and distinct names from colliding after sanitization. */
 function storeFileFor(serverName: string): string {
@@ -55,9 +66,10 @@ export class FileOAuthProvider implements OAuthClientProvider {
   // page cannot inject an authorization code into this login (RFC 8252 8.9).
   private readonly loginState = crypto.randomBytes(16).toString('hex')
 
-  constructor(serverName: string, onRedirect: (authorizationUrl: URL) => void) {
+  constructor(serverName: string, onRedirect: (authorizationUrl: URL) => void, oauth?: OAuthServerConfig) {
     this.storePath = storeFileFor(serverName)
     this.onRedirect = onRedirect
+    this.oauth = oauth
     try {
       this.data = JSON.parse(fs.readFileSync(this.storePath, 'utf-8'))
     } catch {
@@ -65,14 +77,22 @@ export class FileOAuthProvider implements OAuthClientProvider {
     }
   }
 
+  private readonly oauth?: OAuthServerConfig
+
+  /** The client secret for a pre-configured client, from Claude's MCP_CLIENT_SECRET. */
+  private clientSecret(): string | undefined {
+    return this.oauth?.clientId ? process.env.MCP_CLIENT_SECRET : undefined
+  }
+
   private persist(): void {
     fs.mkdirSync(path.dirname(this.storePath), { recursive: true })
     fs.writeFileSync(this.storePath, JSON.stringify(this.data), { mode: 0o600 })
   }
 
-  /** The port a prior login registered, so a re-login can bind the same one. */
+  /** The configured callbackPort (Claude: for pre-registered redirect URIs), else
+   * the port a prior login registered, so a re-login can bind the same one. */
   savedRedirectPort(): number | undefined {
-    return this.data.redirectPort
+    return this.oauth?.callbackPort ?? this.data.redirectPort
   }
 
   /** Record the loopback port the callback server actually bound; the redirect
@@ -96,11 +116,19 @@ export class FileOAuthProvider implements OAuthClientProvider {
       grant_types: ['authorization_code', 'refresh_token'],
       response_types: ['code'],
       // A local CLI is a public client; PKCE carries the proof instead of a secret.
-      token_endpoint_auth_method: 'none',
+      // A pre-configured client with an MCP_CLIENT_SECRET authenticates with it.
+      token_endpoint_auth_method: this.clientSecret() ? 'client_secret_post' : 'none',
+      // Claude: oauth.scopes pins the scopes requested during authorization.
+      ...(this.oauth?.scopes ? { scope: this.oauth.scopes } : {}),
     }
   }
 
   clientInformation(): OAuthClientInformationMixed | undefined {
+    // A pre-configured clientId replaces dynamic registration entirely.
+    if (this.oauth?.clientId) {
+      const secret = this.clientSecret()
+      return { client_id: this.oauth.clientId, ...(secret ? { client_secret: secret } : {}) }
+    }
     return this.data.client
   }
 

package/extensions/mcp/config.ts

@@ -7,6 +7,7 @@ import * as fs from 'node:fs'
 import * as os from 'node:os'
 import * as path from 'node:path'
 import { claudeConfigDir } from '../internal/config-dir.js'
+import type { OAuthServerConfig } from '../internal/mcp-oauth.js'
 import type { InstalledPlugin } from '../internal/plugins.js'
 import { findNearestFile } from '../internal/project-root.js'
 
@@ -32,6 +33,8 @@ export interface HttpServerConfig {
   headers?: Record<string, string>
   bearerToken?: string
   bearerTokenEnv?: string
+  /** Claude's oauth object: pre-registered client, fixed callback port, pinned scopes. */
+  oauth?: OAuthServerConfig
   /** A command whose JSON stdout is merged into the connect headers, for auth
    * schemes other than OAuth/static tokens (Claude's headersHelper). */
   headersHelper?: string

package/extensions/mcp/index.ts

@@ -48,7 +48,7 @@ import { disabledServerNames, loadConfigFrom, loadPluginServers, loadUserScope,
 import { collectServerResourceEntries, listAllPrompts, listAllTools, type McpToolInfo, resourceServerFilter } from './listing.js'
 import { formatPromptCommandName, formatToolName, type McpContentBlock, type McpPromptInfo, mapContent, mapPromptArguments, normalizeSchema, promptMessageContent } from './mapping.js'
 import { applyServerPolicy, loadManagedMcpServers, type McpPolicy, mcpAllowDeny, projectServerPolicy, splitByPolicy } from './policy.js'
-import { type AuthUi, callRequestOptions, callTimeoutMs, connect, connectTimeoutMs, type ServerCallTuning, type SessionDirs, serverCallTuning, withTimeout } from './transport.js'
+import { type AuthUi, callRequestOptions, callTimeoutMs, connect, connectTimeoutMs, connectWithRetries, isUnauthorized, type ServerCallTuning, type SessionDirs, serverCallTuning, withTimeout } from './transport.js'
 
 export { managedSettingsPath, setManagedSettingsPath } from '../internal/managed-settings.js'
 // Re-exports for consumers: the module split keeps the extension's public surface
@@ -91,16 +91,25 @@ export default async function mcpExtension(pi: ExtensionAPI) {
   // The session's directories, set at session_start before any connect: the launch
   // directory answers roots/list, the project root feeds CLAUDE_PROJECT_DIR.
   let sessionDirs: SessionDirs | undefined
+  // Shutdown closes clients while they are still in the map; the onclose handlers
+  // must not schedule reconnects for that deliberate teardown.
+  let shuttingDown = false
   const callTuning = (name: string): ServerCallTuning => {
     const config = serverConfigs.get(name)
     return config ? serverCallTuning(config) : {}
   }
   // Let other extensions (hooks' mcp_tool type) call a connected server's tool.
+  // The same 401/403 reconnect-and-retry-once as registered tools, when the
+  // server's config is known; a name with no stored config calls through once.
   setMcpToolCaller(async (server, tool, input) => {
-    const client = clients.get(server)
-    if (!client) throw new Error(`MCP server "${server}" is not connected`)
-    const tuning = callTuning(server)
-    const result = await client.callTool({ name: tool, arguments: input }, undefined, callRequestOptions(tuning.serverTimeoutMs ?? callTimeoutMs(), tuning))
+    const config = serverConfigs.get(server)
+    const result = config
+      ? await callToolWithAuthRetry(server, config, { name: tool, arguments: input }, `${server}: ${tool}`)
+      : await (async () => {
+          const client = clients.get(server)
+          if (!client) throw new Error(`MCP server "${server}" is not connected`)
+          return await client.callTool({ name: tool, arguments: input }, undefined, callRequestOptions(callTimeoutMs()))
+        })()
     const text = mapContent(result.content as McpContentBlock[], result.structuredContent)
       .filter((part): part is { type: 'text'; text: string } => part.type === 'text')
       .map((part) => part.text)
@@ -142,17 +151,13 @@ export default async function mcpExtension(pi: ExtensionAPI) {
           // present at registration: pi has no tool unregister, so after a server drops
           // and a later session_start reconnects it, registerTools skips re-registration
           // and this closure would otherwise keep calling the old, closed client.
-          const current = clients.get(name)
-          if (!current) throw new Error(`MCP server "${name}" is not connected`)
           // The per-server timeout (Claude's, 1s floor) or MCP_TOOL_TIMEOUT is the
           // wall-clock ceiling; callRequestOptions layers the idle timeout under it, which
-          // the SDK enforces (resetting on progress). Pass the options to the SDK too: its
-          // own default request timeout is 60s and would otherwise reject first. The outer
-          // race uses the wall budget, never the idle window, so a progressing call is not
-          // cut off at the idle timeout.
-          const tuning = serverCallTuning(config)
-          const wall = tuning.serverTimeoutMs ?? callTimeoutMs()
-          const result = await withTimeout(current.callTool({ name: tool.name, arguments: params as Record<string, unknown> }, undefined, callRequestOptions(wall, tuning)), wall, toolName)
+          // the SDK enforces (resetting on progress). The outer race uses the wall
+          // budget, never the idle window, so a progressing call is not cut off at the
+          // idle timeout. A 401/403 reconnects once (fresh helper headers or refreshed
+          // OAuth tokens) and retries once.
+          const result = await callToolWithAuthRetry(name, config, { name: tool.name, arguments: params as Record<string, unknown> }, toolName)
           const content = mapContent(result.content as McpContentBlock[], result.structuredContent)
           const details: { error?: string } = {}
           if (result.isError) {
@@ -167,6 +172,53 @@ export default async function mcpExtension(pi: ExtensionAPI) {
     return count
   }
 
+  /** Claude's mid-session reconnect for a dropped remote server: five attempts with
+   * a delay doubling from one second. connectServers redoes the full bring-up
+   * (tools, prompts, subscriptions, a fresh onclose) and its duplicate guard skips
+   * out if another path already reconnected the name. Runs without authUi: a server
+   * that now needs a login ends failed, and after the fifth failure the last
+   * attempt's failed status stands, with a session restart as the manual retry. */
+  async function reconnectWithBackoff(name: string, config: ServerConfig): Promise<void> {
+    for (let attempt = 0; attempt < 5; attempt++) {
+      await new Promise((resolve) => setTimeout(resolve, 1000 * 2 ** attempt))
+      if (shuttingDown || clients.has(name)) return
+      await connectServers({ [name]: config }, undefined, true)
+      if (clients.has(name)) return
+    }
+  }
+
+  /** Claude's 401/403 tool-call recovery: drop the client and reconnect once, so the
+   * headersHelper re-runs (fresh credential) or the OAuth tokens refresh, then the
+   * caller retries the call once. The map delete precedes the close so the onclose
+   * guard does not also schedule a backoff reconnect. */
+  async function reconnectForAuth(name: string, config: ServerConfig): Promise<void> {
+    const old = clients.get(name)
+    if (old) {
+      clients.delete(name)
+      await withTimeout(old.close(), 3000, 'close').catch(() => {})
+    }
+    await connectServers({ [name]: config }, undefined, true)
+  }
+
+  /** A tool call with the auth retry: on a 401/403 rejection, reconnect once and
+   * retry once; a second auth failure surfaces to the caller. */
+  async function callToolWithAuthRetry(name: string, config: ServerConfig, args: { name: string; arguments: Record<string, unknown> }, label: string): Promise<Awaited<ReturnType<Client['callTool']>>> {
+    const tuning = serverCallTuning(config)
+    const wall = tuning.serverTimeoutMs ?? callTimeoutMs()
+    const callOnce = async () => {
+      const current = clients.get(name)
+      if (!current) throw new Error(`MCP server "${name}" is not connected`)
+      return await withTimeout(current.callTool(args, undefined, callRequestOptions(wall, tuning)), wall, label)
+    }
+    try {
+      return await callOnce()
+    } catch (error) {
+      if (!isUnauthorized(error)) throw error
+      await reconnectForAuth(name, config)
+      return await callOnce()
+    }
+  }
+
   // Prompt command name -> the server and prompt that own it, so a refresh re-listing
   // the same prompt is told apart both from a cross-server collision and from a second
   // prompt on the same server whose name normalizes to the one already taken (e.g.
@@ -340,7 +392,7 @@ export default async function mcpExtension(pi: ExtensionAPI) {
     }
   }
 
-  async function connectServers(servers: Record<string, ServerConfig>, authUi?: AuthUi): Promise<void> {
+  async function connectServers(servers: Record<string, ServerConfig>, authUi?: AuthUi, noRetry = false): Promise<void> {
     const pending: [string, ServerConfig][] = []
     for (const [name, config] of Object.entries(servers)) {
       // A later scope must not take the name of a server that already connected: it
@@ -360,7 +412,9 @@ export default async function mcpExtension(pi: ExtensionAPI) {
       pending.map(async ([name, config]) => {
         warnOnTypelessUrl(name, config)
         try {
-          const client = await connect(name, config, authUi, sessionDirs)
+          // First connections retry transient failures (Claude: up to three times for
+          // HTTP/SSE); the backoff reconnect below carries its own schedule instead.
+          const client = noRetry ? await connect(name, config, authUi, sessionDirs) : await connectWithRetries(name, config, authUi, sessionDirs)
           clients.set(name, client)
           const tools = await withTimeout(listAllTools(client), connectTimeoutMs(), `list tools ${name}`)
           registerTools(name, config, tools)
@@ -382,6 +436,11 @@ export default async function mcpExtension(pi: ExtensionAPI) {
             if (clients.get(name) !== client) return
             clients.delete(name)
             status.set(name, { state: 'disconnected', tools: 0 })
+            // Claude reconnects a dropped remote server with exponential backoff;
+            // stdio servers are local processes and are not reconnected. Shutdown
+            // closes clients while they are still in the map, so the flag guards
+            // against scheduling a reconnect for a deliberate teardown.
+            if (!shuttingDown && !serverCallTuning(config).stdio) void reconnectWithBackoff(name, config)
           }
         } catch (error) {
           status.set(name, { state: `failed: ${error instanceof Error ? error.message : String(error)}`, tools: 0 })
@@ -487,6 +546,9 @@ export default async function mcpExtension(pi: ExtensionAPI) {
     // withdrawn tool keeps its registration and surfaces the server's own error), which
     // is why serverToolCount reads from `registered` to recover the true count here.
     status.clear()
+    // A same-process session switch (/new, /resume) shut the last session down;
+    // this one may reconnect again.
+    shuttingDown = false
     // Claude answers roots/list with the session's launch directory and exports the
     // project root as CLAUDE_PROJECT_DIR to stdio servers; both derive from ctx.cwd.
     sessionDirs = { projectDir: repoRoot(ctx.cwd) ?? ctx.cwd, launchDir: ctx.cwd }
@@ -519,6 +581,9 @@ export default async function mcpExtension(pi: ExtensionAPI) {
   })
 
   pi.on('session_shutdown', async () => {
+    // Closing fires each client's onclose while it is still in the map; the flag
+    // stops those handlers (and any in-flight backoff loop) from reconnecting.
+    shuttingDown = true
     // Close in parallel with a per-client timeout so one hung server can't stall pi's exit.
     await Promise.all([...clients.values()].map((client) => withTimeout(client.close(), 3000, 'close').catch(() => {})))
     // Drop the closed clients and their status now rather than waiting on each client's

package/extensions/mcp/oauth-flow.ts

@@ -6,7 +6,7 @@
  */
 
 import type { Client } from '@modelcontextprotocol/sdk/client/index.js'
-import { FileOAuthProvider, openBrowser, startCallbackServer, waitForAuthCode } from '../internal/mcp-oauth.js'
+import { FileOAuthProvider, type OAuthServerConfig, openBrowser, startCallbackServer, waitForAuthCode } from '../internal/mcp-oauth.js'
 import { type AuthUi, connectWithTimeout, isUnauthorized, type MakeTransport, OAuthRequiredError } from './transport.js'
 
 /** Browser logins are human-paced; a connect-sized timeout would cut them off. */
@@ -45,13 +45,17 @@ export function serializeInteractiveOAuth<T>(run: () => Promise<T>): Promise<T>
  * token exchange error) is wrapped as an auth failure, not a transport mismatch: that
  * keeps the typeless-url caller from retrying over SSE and prompting for a second login.
  */
-export async function runInteractiveOAuth(name: string, config: { url: string }, makeTransport: MakeTransport, label: string, authUi: AuthUi, newClient: () => Client): Promise<Client> {
+export async function runInteractiveOAuth(name: string, config: { url: string; oauth?: OAuthServerConfig }, makeTransport: MakeTransport, label: string, authUi: AuthUi, newClient: () => Client): Promise<Client> {
   const approved = await authUi.confirm(`MCP server "${name}" requires login`, `Open your browser to authorize ${config.url}?`)
   if (!approved) throw new OAuthRequiredError(`login declined for ${name}`)
-  const provider = new FileOAuthProvider(name, (authorizationUrl) => {
-    openBrowser(String(authorizationUrl))
-    authUi.notify(`Authorize "${name}" in the browser. If it did not open: ${authorizationUrl}`, 'info')
-  })
+  const provider = new FileOAuthProvider(
+    name,
+    (authorizationUrl) => {
+      openBrowser(String(authorizationUrl))
+      authUi.notify(`Authorize "${name}" in the browser. If it did not open: ${authorizationUrl}`, 'info')
+    },
+    config.oauth,
+  )
   const { server, port } = await startCallbackServer(provider.savedRedirectPort())
   provider.bindRedirectPort(port)
   try {

package/extensions/mcp/transport.ts

@@ -15,7 +15,7 @@ import { getDefaultEnvironment, StdioClientTransport } from '@modelcontextprotoc
 import { StreamableHTTPClientTransport } from '@modelcontextprotocol/sdk/client/streamableHttp.js'
 import { WebSocketClientTransport } from '@modelcontextprotocol/sdk/client/websocket.js'
 import { ListRootsRequestSchema } from '@modelcontextprotocol/sdk/types.js'
-import { FileOAuthProvider } from '../internal/mcp-oauth.js'
+import { FileOAuthProvider, type OAuthServerConfig } from '../internal/mcp-oauth.js'
 import { expandCwd, type HttpServerConfig, interpolateEnv, type ServerConfig, type StdioServerConfig } from './config.js'
 import { runInteractiveOAuth, serializeInteractiveOAuth } from './oauth-flow.js'
 
@@ -247,6 +247,11 @@ export async function connect(name: string, config: ServerConfig, authUi?: AuthU
     await connectWithTimeout(client, transport, `connect ${name} (ws)`)
     return client
   }
+  // The SDK's OAuth discovery has no override seam, so a configured metadata URL
+  // cannot be honored; say so instead of silently using standard discovery.
+  if (config.oauth?.authServerMetadataUrl) {
+    console.warn(`pi-code-mcp: server ${name} sets oauth.authServerMetadataUrl, which the MCP SDK cannot override; using standard discovery`)
+  }
   const headers: Record<string, string> = {}
   for (const [key, value] of Object.entries(config.headers ?? {})) headers[key] = fill(value)
   const token = resolveBearerToken(config)
@@ -271,6 +276,33 @@ export async function connect(name: string, config: ServerConfig, authUi?: AuthU
   }
 }
 
+/** Whether a connect failure is worth retrying: a 5xx response, a refused or reset
+ * connection, or a timeout. Auth and not-found errors need a configuration change. */
+function isTransientConnectError(error: unknown): boolean {
+  if (isUnauthorized(error)) return false
+  const code = typeof error === 'object' && error !== null ? (error as { code?: unknown }).code : undefined
+  if (typeof code === 'number') return code >= 500
+  return /ECONNREFUSED|ECONNRESET|ETIMEDOUT|timed out after/.test(String(error))
+}
+
+const delay = (ms: number): Promise<void> => new Promise((resolve) => setTimeout(resolve, ms))
+
+/** Connect with Claude's first-connection retry: an HTTP or SSE server's transient
+ * failure (5xx, connection refused, timeout) is retried up to three times with a
+ * 1s-doubling delay. Stdio and WebSocket connects are attempted once, as are auth
+ * and not-found failures. */
+export async function connectWithRetries(name: string, config: ServerConfig, authUi?: AuthUi, session?: SessionDirs): Promise<Client> {
+  const retriable = !isStdio(config) && config.type !== 'ws' && config.type !== 'websocket'
+  for (let attempt = 0; ; attempt++) {
+    try {
+      return await connect(name, config, authUi, session)
+    } catch (error) {
+      if (!retriable || attempt >= 3 || !isTransientConnectError(error)) throw error
+      await delay(1000 * 2 ** attempt)
+    }
+  }
+}
+
 /** Run a headersHelper command and parse its JSON stdout into headers, under the
  * environment helperEnv built. A failure or a 10s timeout yields no extra headers
  * rather than blocking the connection. */
@@ -318,13 +350,15 @@ export type MakeTransport = (authProvider?: OAuthClientProvider) => HttpFamilyTr
  * header, or one from a headersHelper) never enters the OAuth path: the credential
  * to fix is the configured one, so an auth failure reports as a failed connection.
  */
-async function connectHttpFamily(name: string, config: { url: string }, makeTransport: MakeTransport, label: string, hasConfiguredAuth: boolean, authUi: AuthUi | undefined, session?: SessionDirs): Promise<Client> {
+async function connectHttpFamily(name: string, config: { url: string; oauth?: OAuthServerConfig }, makeTransport: MakeTransport, label: string, hasConfiguredAuth: boolean, authUi: AuthUi | undefined, session?: SessionDirs): Promise<Client> {
   const newClient = () => makeClient(session)
   // Stored tokens ride the first attempt so the SDK refreshes them; with none, no
   // provider is attached, so a 401 surfaces as a transport error carrying code 401
   // (isUnauthorized detects it) and only the interactive provider below ever runs
-  // dynamic registration, keeping it bound to the real callback port.
-  const silent = hasConfiguredAuth ? undefined : new FileOAuthProvider(name, () => {})
+  // dynamic registration, keeping it bound to the real callback port. A
+  // pre-configured client (oauth.clientId) rides the silent provider too, so its
+  // stored tokens refresh with the configured credentials.
+  const silent = hasConfiguredAuth ? undefined : new FileOAuthProvider(name, () => {}, config.oauth)
   try {
     const client = newClient()
     await connectWithTimeout(client, makeTransport(silent?.hasTokens() ? silent : undefined), label)

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "pi-code",
-  "version": "1.0.27",
+  "version": "1.0.28",
   "description": "Claude Code experience for the pi coding agent: reads your .claude config (rules, commands, skills, hooks, output styles, MCP servers, agents) and adds todo, checkpoints, memory, web, and subagents",
   "keywords": [
     "pi",