pi-code 1.0.13 → 1.0.15

package/extensions/mcp/mapping.ts

@@ -0,0 +1,173 @@
+/**
+ * MCP name and content mapping: the pi tool/prompt-command name formatting, prompt
+ * argument mapping, input-schema normalization, and the tool/prompt/resource content
+ * blocks mapped into pi's output budget.
+ */
+
+import { DEFAULT_MAX_BYTES } from '@earendil-works/pi-coding-agent'
+import { splitArgs } from '../internal/command-file.js'
+import { capForContext } from '../internal/output-guard.js'
+
+export function formatToolName(server: string, tool: string): string {
+  return `${server}_${tool}`.replaceAll('-', '_')
+}
+
+/** Claude exposes server prompts as /mcp__<server>__<prompt> slash commands. Both
+ * names normalize like formatToolName, extended to spaces: dashes and spaces each
+ * become an underscore. */
+export function formatPromptCommandName(server: string, prompt: string): string {
+  const normalize = (name: string): string => name.replace(/[\s-]/g, '_')
+  return `mcp__${normalize(server)}__${normalize(prompt)}`
+}
+
+export interface McpPromptArgumentInfo {
+  name: string
+  description?: string
+  required?: boolean
+}
+
+export interface McpPromptInfo {
+  name: string
+  description?: string
+  arguments?: McpPromptArgumentInfo[]
+}
+
+/** Claude passes prompt arguments space-separated after the command. Tokens map
+ * positionally onto the declared arguments, split the way slash-command args are
+ * (quoted runs stay together); the last declared argument absorbs any trailing
+ * tokens so free text at the end is not silently dropped. Declared arguments with
+ * no token are omitted, and the server enforces its own `required`. */
+export function mapPromptArguments(declared: ReadonlyArray<{ name: string }> | undefined, args: string): Record<string, string> {
+  const tokens = splitArgs(args)
+  const names = (declared ?? []).map((argument) => argument.name)
+  const mapped: Record<string, string> = {}
+  for (let index = 0; index < names.length && index < tokens.length; index++) {
+    mapped[names[index]] = index === names.length - 1 ? tokens.slice(index).join(' ') : tokens[index]
+  }
+  return mapped
+}
+
+/** The content blocks a getPrompt result injects. Each message carries one content
+ * block; the blocks ride the same mapContent budget as tool output, and image blocks
+ * are carried through rather than dropped, since sendUserMessage accepts them and a
+ * vision prompt is worthless flattened to text. An empty message list yields no
+ * blocks, and messages that carry only empty text yield none either, so the caller
+ * can skip the turn rather than drive it on an empty or sentinel message. */
+export function promptMessageContent(messages: ReadonlyArray<{ content: unknown }>): ToolContent[] {
+  if (messages.length === 0) return []
+  return mapContent(messages.map((message) => message.content as McpContentBlock)).filter((block) => block.type !== 'text' || block.text.trim() !== '')
+}
+
+/** Merge the `properties` (and, for allOf, the `required`) of a root-level combinator's
+ * branches into one flat object schema. Without this a tool whose input schema is a bare
+ * anyOf/oneOf/allOf (no top-level `type`) would present no properties at all, so the model
+ * would be forced to call it with no arguments. */
+function mergeCombinatorBranches(branches: unknown[]): { properties: Record<string, unknown>; required: string[] } {
+  const properties: Record<string, unknown> = {}
+  const required = new Set<string>()
+  for (const branch of branches) {
+    if (!branch || typeof branch !== 'object') continue
+    const b = branch as Record<string, unknown>
+    if (b.properties && typeof b.properties === 'object') Object.assign(properties, b.properties as Record<string, unknown>)
+    if (Array.isArray(b.required)) for (const name of b.required) if (typeof name === 'string') required.add(name)
+  }
+  return { properties, required: [...required] }
+}
+
+export function normalizeSchema(schema: unknown): object {
+  const base = (schema as Record<string, unknown>) ?? {}
+  const { $schema: _dropSchema, additionalProperties: _dropAdditional, ...rest } = base
+  if (rest.type) return rest
+  // A root-level combinator carries the real parameters in its branches; flatten them
+  // into one object schema rather than emptying it. allOf means every branch applies, so
+  // its required union is kept; anyOf/oneOf branches are alternatives, so required is left
+  // open (the server still enforces its own).
+  const allOf = Array.isArray(rest.allOf) ? rest.allOf : undefined
+  let branches = allOf
+  if (!branches && Array.isArray(rest.anyOf)) branches = rest.anyOf
+  if (!branches && Array.isArray(rest.oneOf)) branches = rest.oneOf
+  if (!branches) return { type: 'object', properties: {} }
+  const { properties, required } = mergeCombinatorBranches(branches)
+  const merged: Record<string, unknown> = { type: 'object', properties }
+  if (typeof rest.description === 'string') merged.description = rest.description
+  if (allOf && required.length > 0) merged.required = required
+  return merged
+}
+
+export interface McpContentBlock {
+  type: string
+  text?: string
+  data?: string
+  mimeType?: string
+  resource?: { uri?: string; text?: string }
+}
+
+export type ToolContent = { type: 'text'; text: string } | { type: 'image'; data: string; mimeType: string }
+
+export function mapContent(content: McpContentBlock[] | undefined, structured?: unknown): ToolContent[] {
+  // capForContext every text output, whatever its source: a server can blow the tool-output
+  // budget through a resource block, a JSON-stringified block, or the structured fallback,
+  // not only a text block. The per-block cap alone is not a budget, though: a server
+  // answering with one block per file multiplies it by the block count, so the blocks
+  // are capped again as a whole below.
+  const text = (value: string): ToolContent => ({ type: 'text', text: capForContext(value) })
+  if (!content || content.length === 0) {
+    return [text(structured !== undefined ? JSON.stringify(structured, null, 2) : '(empty result)')]
+  }
+  const mapped: ToolContent[] = content.map((block): ToolContent => {
+    if (block.type === 'text') {
+      return text(block.text ?? '')
+    }
+    if (block.type === 'image' && block.data) {
+      return { type: 'image', data: block.data, mimeType: block.mimeType ?? 'image/png' }
+    }
+    if (block.type === 'resource' && block.resource) {
+      return text(`[Resource: ${block.resource.uri ?? 'unknown'}]\n${block.resource.text ?? ''}`)
+    }
+    return text(JSON.stringify(block))
+  })
+  return capTotal(mapped)
+}
+
+/**
+ * Bound a result's text as a whole, not each block. The per-block cap multiplies by
+ * the block count, so a server answering with one block per file still injects
+ * megabytes.
+ *
+ * Blocks are kept whole. Each has already been capped on its own, so keeping the one
+ * that crosses the budget bounds the text at roughly a single cap rather than at the
+ * block count times it, and it preserves that block's own truncation notice, which
+ * states how much of it was dropped. Blocks after it are omitted rather than skipped
+ * over, so what reaches the model is a prefix of what the server sent, and the number
+ * omitted is stated so a truncated set is distinguishable from a complete one.
+ *
+ * Images pass through uncut and do not spend the budget: base64 cut short is a broken
+ * image rather than a smaller one, so nothing here can bound them, and charging the
+ * budget for one would only delete the caption that accompanies a screenshot.
+ */
+export function capTotal(blocks: ToolContent[]): ToolContent[] {
+  const kept: ToolContent[] = []
+  let spent = 0
+  let full = false
+  let dropped = 0
+  for (const block of blocks) {
+    if (block.type !== 'text') {
+      kept.push(block)
+      continue
+    }
+    const size = Buffer.byteLength(block.text, 'utf-8')
+    // The first text block always goes through: a lone oversized one is better read
+    // truncated, with its own notice, than replaced by a marker saying it existed.
+    if (full || (spent > 0 && spent + size > DEFAULT_MAX_BYTES)) {
+      full = true
+      dropped++
+      continue
+    }
+    kept.push(block)
+    spent += size
+  }
+  if (dropped > 0) {
+    kept.push({ type: 'text', text: `[${dropped} further content block${dropped === 1 ? '' : 's'} omitted: tool output budget spent]` })
+  }
+  return kept
+}

package/extensions/mcp/oauth-flow.ts

@@ -0,0 +1,80 @@
+/**
+ * MCP interactive OAuth: the serialization queue that stops two browser logins from
+ * stacking dialogs and tabs, and the interactive login itself (confirm, open the
+ * browser, catch the loopback redirect, exchange the code via the SDK's finishAuth).
+ * The FileOAuthProvider and callback server live in internal/mcp-oauth.
+ */
+
+import type { Client } from '@modelcontextprotocol/sdk/client/index.js'
+import { FileOAuthProvider, 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. */
+const OAUTH_FLOW_TIMEOUT_MS = 180_000
+
+/** Wrap a login-flow failure as OAuthRequiredError, passing an existing one through
+ * unchanged so its message is not doubled. */
+function asOAuthRequiredError(name: string, error: unknown): OAuthRequiredError {
+  if (error instanceof OAuthRequiredError) return error
+  const detail = error instanceof Error ? error.message : String(error)
+  return new OAuthRequiredError(`login for ${name} failed: ${detail}`)
+}
+
+/** Interactive OAuth logins block on a confirm dialog and open a browser tab, so two
+ * at once (a user-scope and a consented project-scope server both 401ing, connecting in
+ * parallel) would stack dialogs and browser tabs. This chains them so a second
+ * interactive login waits for the first to settle; the tail is reset to a resolved
+ * promise regardless of outcome, so a failed login never poisons the queue. Silent
+ * (stored-token) connects do not pass through here and stay fully parallel. */
+let oauthQueue: Promise<unknown> = Promise.resolve()
+
+export function serializeInteractiveOAuth<T>(run: () => Promise<T>): Promise<T> {
+  const result = oauthQueue.then(run, run)
+  oauthQueue = result.then(
+    () => {},
+    () => {},
+  )
+  return result
+}
+
+/**
+ * The interactive half of the OAuth login, reached only once a silent connect has
+ * failed with a 401 and a UI is present: confirm, open the browser, catch the loopback
+ * redirect, and exchange the code via the SDK's finishAuth. Past the confirm the server
+ * is known to need OAuth, so any failure here (a denied consent page, the 180s wait, a
+ * 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> {
+  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 { server, port } = await startCallbackServer(provider.savedRedirectPort())
+  provider.bindRedirectPort(port)
+  try {
+    const transport = makeTransport(provider)
+    // Verify the redirect echoes this login's state, so a stray or forged callback to the
+    // loopback port cannot inject a code or abort the login (see waitForAuthCode).
+    const pendingCode = waitForAuthCode(server, OAUTH_FLOW_TIMEOUT_MS, provider.state())
+    pendingCode.catch(() => {}) // consumed below; an abandoned login must not surface as unhandled
+    const client = newClient()
+    try {
+      await connectWithTimeout(client, transport, label)
+      return client // authorized between attempts; nothing left to exchange
+    } catch (retryError) {
+      if (!isUnauthorized(retryError)) throw retryError
+      const code = await pendingCode
+      await transport.finishAuth(code)
+      const authed = newClient()
+      await connectWithTimeout(authed, makeTransport(provider), label)
+      return authed
+    }
+  } catch (flowError) {
+    throw asOAuthRequiredError(name, flowError)
+  } finally {
+    server.close()
+  }
+}

package/extensions/mcp/policy.ts

@@ -0,0 +1,140 @@
+/**
+ * MCP server policy: the per-project .mcp.json approvals, the managed allow/deny
+ * lists, and the managed-mcp.json exclusive-control loader. The managed settings
+ * path (and its test-seam override) is owned by internal/managed-settings.
+ */
+
+import * as fs from 'node:fs'
+import * as path from 'node:path'
+import { claudeConfigDir } from '../internal/config-dir.js'
+import { managedSettingsFile } from '../internal/managed-settings.js'
+import { findNearestFile } from '../internal/project-root.js'
+import type { ServerConfig } from './config.js'
+
+export interface ProjectServerPolicy {
+  disabled: Set<string>
+  consented: Set<string>
+  consentAll: boolean
+}
+
+/** Claude's per-server approvals for project .mcp.json servers.
+ *
+ * Consent-granting keys (enabledMcpjsonServers, enableAllProjectMcpServers) count
+ * from the user's own settings always, and from the project's settings.local.json
+ * only once the project itself is approved. That file is gitignored by convention,
+ * not by enforcement: a repository can commit one, and honoring it unconditionally
+ * let a hostile repo self-approve a server whose `command` runs on connect, even
+ * after the user declined the trust prompt.
+ *
+ * disabledMcpjsonServers counts from every file, including the repo's own, and wins
+ * over consent: a repo may always restrict itself further, never less. */
+export function projectServerPolicy(cwd: string, home: string, projectApproved: boolean): ProjectServerPolicy {
+  const read = (file: string): Record<string, unknown> => {
+    try {
+      return JSON.parse(fs.readFileSync(file, 'utf-8'))
+    } catch {
+      return {}
+    }
+  }
+  const names = (value: unknown): string[] => (Array.isArray(value) ? value.filter((entry): entry is string => typeof entry === 'string') : [])
+  const userSettings = read(path.join(claudeConfigDir(home), 'settings.json'))
+  const projectSettings = read(findNearestFile(cwd, path.join('.claude', 'settings.json')) ?? path.join(cwd, '.claude', 'settings.json'))
+  const localSettings = read(findNearestFile(cwd, path.join('.claude', 'settings.local.json')) ?? path.join(cwd, '.claude', 'settings.local.json'))
+  const disabled = new Set([...names(userSettings.disabledMcpjsonServers), ...names(projectSettings.disabledMcpjsonServers), ...names(localSettings.disabledMcpjsonServers)])
+  const consentSources = projectApproved ? [userSettings, localSettings] : [userSettings]
+  const consented = new Set(consentSources.flatMap((settings) => names(settings.enabledMcpjsonServers)))
+  const consentAll = consentSources.some((settings) => settings.enableAllProjectMcpServers === true)
+  return { disabled, consented, consentAll }
+}
+
+/** Split project servers by the per-server policy: never-connect, connect without the
+ * whole-project confirm, and still gated behind it. */
+export function splitByPolicy(candidates: Record<string, ServerConfig>, policy: ProjectServerPolicy): { consented: Record<string, ServerConfig>; gated: Record<string, ServerConfig> } {
+  const consented: Record<string, ServerConfig> = {}
+  const gated: Record<string, ServerConfig> = {}
+  for (const [name, config] of Object.entries(candidates)) {
+    if (policy.disabled.has(name)) continue
+    if (policy.consentAll || policy.consented.has(name)) consented[name] = config
+    else gated[name] = config
+  }
+  return { consented, gated }
+}
+
+/** Claude's `allowedMcpServers`/`deniedMcpServers`, read from managed settings only
+ * (not user or project settings, so a repo can neither widen nor narrow the policy),
+ * applied globally to every server across scopes. Entries are `{ serverName }` objects
+ * (bare strings tolerated). `allowed` is null when unset (no restriction); an empty
+ * set is an explicit lockdown, as Claude documents (empty allow array = deny all). */
+export function mcpAllowDeny(managedFile: string = managedSettingsFile()): { allowed: Set<string> | null; denied: Set<string> } {
+  let settings: Record<string, unknown> = {}
+  try {
+    const parsed = JSON.parse(fs.readFileSync(managedFile, 'utf-8'))
+    if (parsed && typeof parsed === 'object') settings = parsed
+  } catch {
+    // No managed policy on this machine: no restriction.
+  }
+  const entryName = (entry: unknown): string | undefined => {
+    if (typeof entry === 'string') return entry
+    const serverName = (entry as { serverName?: unknown })?.serverName
+    return typeof serverName === 'string' ? serverName : undefined
+  }
+  const names = (value: unknown): string[] => (Array.isArray(value) ? value.map(entryName).filter((name): name is string => typeof name === 'string' && name.length > 0) : [])
+  return {
+    allowed: Array.isArray(settings.allowedMcpServers) ? new Set(names(settings.allowedMcpServers)) : null,
+    denied: new Set(names(settings.deniedMcpServers)),
+  }
+}
+
+/** The managed-mcp.json path: a sibling of managed-settings.json (same directory). Derived
+ * through the same test seam so a test can write both into one temp dir. */
+export function managedMcpPath(managedFile: string = managedSettingsFile()): string {
+  return path.join(path.dirname(managedFile), 'managed-mcp.json')
+}
+
+/** Claude's managed-mcp.json: when it exists beside managed-settings.json it takes
+ * exclusive control of MCP. Only its `mcpServers` load; user, project, and plugin servers
+ * are all suppressed (and the project-approval flow with them), and an empty map disables
+ * MCP entirely. Returns the managed server map (possibly empty) when the file exists and
+ * parses, or null only when the file is absent, in which case MCP loads from the usual
+ * scopes exactly as before. A file that parses but carries no `mcpServers` object is an
+ * empty managed set, so a deployed-but-bodyless policy locks down rather than silently
+ * reopening the other scopes. A file that is PRESENT but not valid JSON fails closed to
+ * the same empty set (deny-all) rather than reopening those scopes: the lockdown intent
+ * means a corrupt or truncated policy file must not become an allow-all. The allow/deny
+ * lists still filter the returned set. */
+export function loadManagedMcpServers(managedFile: string = managedSettingsFile()): Record<string, ServerConfig> | null {
+  const file = managedMcpPath(managedFile)
+  let raw: string
+  try {
+    raw = fs.readFileSync(file, 'utf-8')
+  } catch {
+    // Absent (or unreadable) managed-mcp.json: no managed MCP control, load normally.
+    return null
+  }
+  let parsed: unknown
+  try {
+    parsed = JSON.parse(raw)
+  } catch (error) {
+    // Present but corrupt: fail closed to an empty managed set, exactly like an empty map,
+    // rather than reopening the user/project/plugin scopes.
+    console.warn(`pi-code-mcp: managed-mcp.json is present but not valid JSON (${file}); failing closed to no MCP servers: ${error instanceof Error ? error.message : String(error)}`)
+    return {}
+  }
+  if (parsed === null || typeof parsed !== 'object') return {}
+  const servers = (parsed as { mcpServers?: unknown }).mcpServers
+  if (servers === null || typeof servers !== 'object' || Array.isArray(servers)) return {}
+  return servers as Record<string, ServerConfig>
+}
+
+/** Claude's managed allow/deny lists: `allowed` null means no allow list (keep all);
+ * a set (even empty) is exclusive, so only its members survive; a deny list removes
+ * servers on top, deny winning over allow. */
+export function applyServerPolicy(servers: Record<string, ServerConfig>, allowed: ReadonlySet<string> | null, denied: ReadonlySet<string>): Record<string, ServerConfig> {
+  const out: Record<string, ServerConfig> = {}
+  for (const [name, config] of Object.entries(servers)) {
+    if (denied.has(name)) continue
+    if (allowed !== null && !allowed.has(name)) continue
+    out[name] = config
+  }
+  return out
+}

package/extensions/mcp/transport.ts

@@ -0,0 +1,258 @@
+/**
+ * MCP transport: the connect/call timeouts, the SDK request options for pi's two-tier
+ * timeout, the stdio and HTTP-family (streamable with SSE fallback) connect paths, the
+ * headersHelper, and the auth primitives shared with the interactive OAuth flow.
+ */
+
+import { execFile } from 'node:child_process'
+// SSE is deprecated in favour of Streamable HTTP, but the SDK notes servers still on
+// the old spec exist, so this stays as a fallback for the migration period.
+import { type OAuthClientProvider, UnauthorizedError } from '@modelcontextprotocol/sdk/client/auth.js'
+import { Client } from '@modelcontextprotocol/sdk/client/index.js'
+import { SSEClientTransport } from '@modelcontextprotocol/sdk/client/sse.js' // NOSONAR
+import { getDefaultEnvironment, StdioClientTransport } from '@modelcontextprotocol/sdk/client/stdio.js'
+import { StreamableHTTPClientTransport } from '@modelcontextprotocol/sdk/client/streamableHttp.js'
+import { WebSocketClientTransport } from '@modelcontextprotocol/sdk/client/websocket.js'
+import { FileOAuthProvider } from '../internal/mcp-oauth.js'
+import { expandCwd, interpolateEnv, type ServerConfig, type StdioServerConfig } from './config.js'
+import { runInteractiveOAuth, serializeInteractiveOAuth } from './oauth-flow.js'
+
+const DEFAULT_CONNECT_TIMEOUT_MS = 10_000
+// Claude's MCP_TOOL_TIMEOUT default is effectively hours: the per-call wall-clock budget
+// is only a ceiling, and the idle timeout below is the real guard. 4h matches that model,
+// so a legitimately slow-but-progressing tool is not killed at the old 2 minutes.
+const DEFAULT_CALL_TIMEOUT_MS = 14_400_000
+// The idle timeout: the longest a call may go with no response or progress before it is
+// abandoned. Claude uses a separate idle guard (minutes) rather than the hours-long
+// wall-clock budget; the SDK resets this window on every progress notification.
+const DEFAULT_CALL_IDLE_TIMEOUT_MS = 300_000
+
+/** A positive-integer env override, or the default when unset or unparseable. */
+function envTimeout(name: string, fallback: number): number {
+  const raw = process.env[name]
+  if (raw === undefined) return fallback
+  const value = Number.parseInt(raw, 10)
+  return Number.isInteger(value) && value > 0 ? value : fallback
+}
+
+// Claude honors MCP_TIMEOUT (connect) and MCP_TOOL_TIMEOUT (per-call), both in ms.
+export const connectTimeoutMs = (): number => envTimeout('MCP_TIMEOUT', DEFAULT_CONNECT_TIMEOUT_MS)
+export const callTimeoutMs = (): number => envTimeout('MCP_TOOL_TIMEOUT', DEFAULT_CALL_TIMEOUT_MS)
+
+/** The idle timeout in ms: the longest a call may go with no response or progress before
+ * it is abandoned, overridable by CLAUDE_CODE_MCP_TOOL_IDLE_TIMEOUT, with 0 disabling it
+ * (leaving only the wall-clock budget). Unlike envTimeout, an explicit 0 is honored as
+ * "disabled" rather than falling back to the default. */
+function idleTimeoutMs(): number {
+  const raw = process.env.CLAUDE_CODE_MCP_TOOL_IDLE_TIMEOUT
+  if (raw === undefined) return DEFAULT_CALL_IDLE_TIMEOUT_MS
+  const value = Number.parseInt(raw, 10)
+  if (value === 0) return 0
+  return Number.isInteger(value) && value > 0 ? value : DEFAULT_CALL_IDLE_TIMEOUT_MS
+}
+
+/** The SDK RequestOptions for a call under pi's two-tier timeout: a wall-clock ceiling and,
+ * under it, an idle timeout the SDK resets on every progress notification. When the idle
+ * window is enabled and tighter than the wall budget, `timeout` is that per-quiet-period
+ * deadline (resetTimeoutOnProgress), maxTotalTimeout caps the wall clock, and an onprogress
+ * handler is required: it makes the server address progress to this request and lets the
+ * SDK reset the timer on it. When the idle timeout is disabled, or already looser than the
+ * wall budget, only the wall budget applies. The outer withTimeout race is a wall-clock
+ * backstop and must be raced against `wall`, never the idle window, so a legitimately
+ * progressing call is not cut off. */
+export function callRequestOptions(wall: number): { timeout: number; resetTimeoutOnProgress?: boolean; maxTotalTimeout?: number; onprogress?: () => void } {
+  const idle = idleTimeoutMs()
+  if (idle === 0 || idle >= wall) return { timeout: wall }
+  return { timeout: idle, resetTimeoutOnProgress: true, maxTotalTimeout: wall, onprogress: () => {} }
+}
+
+/** Claude reports a config entry that has a url but no type as an error; pi-code
+ * still connects (streamable HTTP with SSE fallback) but says the entry is wrong. */
+/** An inline bearerToken (interpolated) wins over bearerTokenEnv, which names an
+ * environment variable read as-is. */
+export function resolveBearerToken(config: { bearerToken?: string; bearerTokenEnv?: string }): string | undefined {
+  if (config.bearerToken) return interpolateEnv(config.bearerToken)
+  if (config.bearerTokenEnv) return process.env[config.bearerTokenEnv]
+  return undefined
+}
+
+/** Claude's `headersHelper` output: a flat JSON object of header name -> string,
+ * merged into the connect headers. Non-string values and non-object output are
+ * ignored so a broken helper cannot poison the request. */
+export function parseHelperHeaders(stdout: string): Record<string, string> {
+  let parsed: unknown
+  try {
+    parsed = JSON.parse(stdout)
+  } catch {
+    return {}
+  }
+  if (parsed === null || typeof parsed !== 'object' || Array.isArray(parsed)) return {}
+  const out: Record<string, string> = {}
+  for (const [key, value] of Object.entries(parsed)) if (typeof value === 'string') out[key] = value
+  return out
+}
+
+function isStdio(config: ServerConfig): config is StdioServerConfig {
+  // An explicit type wins; without one, a command field means stdio.
+  return 'command' in config && (config.type === undefined || config.type === 'stdio')
+}
+
+export async function withTimeout<T>(promise: Promise<T>, ms: number, label: string): Promise<T> {
+  let timer: ReturnType<typeof setTimeout> | undefined
+  const timeout = new Promise<never>((_, reject) => {
+    timer = setTimeout(() => reject(new Error(`${label} timed out after ${ms}ms`)), ms)
+  })
+  // If the timeout wins, `promise` stays pending; swallow any late rejection so it can never
+  // surface as an unhandled rejection that crashes the host.
+  promise.catch(() => {})
+  try {
+    return await Promise.race([promise, timeout])
+  } finally {
+    clearTimeout(timer)
+  }
+}
+
+export async function connect(name: string, config: ServerConfig, authUi?: AuthUi): Promise<Client> {
+  const client = new Client({ name: 'pi-code-mcp', version: '0.1.0' })
+  // Names referenced by ${VAR} with no value and no default, gathered across this
+  // server's interpolated fields so the connect can warn once rather than fail with a
+  // mystery 401 or a command that lost an argument.
+  const missing = new Set<string>()
+  const fill = (value: string): string => interpolateEnv(value, process.env, (varName) => missing.add(varName))
+  const warnMissing = (): void => {
+    if (missing.size > 0) console.warn(`pi-code-mcp: server ${name} references undefined variable(s) ${[...missing].join(', ')}; leaving them unexpanded`)
+  }
+  if (isStdio(config)) {
+    // Start from the SDK's allowlist (PATH, HOME, SHELL, ...) rather than the whole
+    // process env: a server should not receive ANTHROPIC_API_KEY or GITHUB_TOKEN just
+    // for being launched. A server that needs a variable names it in its own env block.
+    const env: Record<string, string> = { ...getDefaultEnvironment() }
+    for (const [key, value] of Object.entries(config.env ?? {})) env[key] = fill(value)
+    const transport = new StdioClientTransport({
+      command: fill(config.command),
+      args: (config.args ?? []).map((arg) => fill(arg)),
+      env,
+      cwd: expandCwd(config.cwd),
+      stderr: 'ignore',
+    })
+    warnMissing()
+    await connectWithTimeout(client, transport, `connect ${name}`)
+    return client
+  }
+  const url = new URL(fill(config.url))
+  if (config.type === 'ws' || config.type === 'websocket') {
+    // The SDK's WebSocket transport takes only a url: it carries no headers, bearer
+    // token, or headersHelper output. Warn rather than silently dropping configured
+    // auth, and skip the helper entirely (running it would block the connect for up to
+    // 10s while contributing nothing). A ws server must be reachable without auth.
+    if (config.headers || config.bearerToken || config.bearerTokenEnv || config.headersHelper) {
+      console.warn(`pi-code-mcp: server ${name} is a WebSocket server; the SDK ws transport is url-only, so its headers/bearerToken/headersHelper are ignored`)
+    }
+    const transport = new WebSocketClientTransport(url)
+    warnMissing()
+    await connectWithTimeout(client, transport, `connect ${name} (ws)`)
+    return client
+  }
+  const headers: Record<string, string> = {}
+  for (const [key, value] of Object.entries(config.headers ?? {})) headers[key] = fill(value)
+  const token = resolveBearerToken(config)
+  if (token) headers.Authorization = `Bearer ${token}`
+  // A headersHelper generates connect-time headers for non-OAuth auth schemes; its
+  // JSON stdout merges over the static headers.
+  if (config.headersHelper) Object.assign(headers, await runHeadersHelper(fill(config.headersHelper)))
+  warnMissing()
+  const sseTransport = (authProvider?: OAuthClientProvider) => new SSEClientTransport(url, { requestInit: { headers }, authProvider }) // NOSONAR: explicitly declared or deliberate legacy transport
+  if (config.type === 'sse') {
+    return await connectHttpFamily(name, config, sseTransport, `connect ${name} (sse)`, token, authUi)
+  }
+  try {
+    return await connectHttpFamily(name, config, (authProvider) => new StreamableHTTPClientTransport(url, { requestInit: { headers }, authProvider }), `connect ${name}`, token, authUi)
+  } catch (error) {
+    // An explicitly declared streamable transport must not silently degrade to SSE.
+    if (config.type !== undefined || isUnauthorized(error)) throw error
+    return await connectHttpFamily(name, config, sseTransport, `connect ${name} (sse)`, token, authUi)
+  }
+}
+
+/** Run a headersHelper command and parse its JSON stdout into headers. A failure or
+ * a 10s timeout yields no extra headers rather than blocking the connection. */
+function runHeadersHelper(command: string): Promise<Record<string, string>> {
+  return new Promise((resolve) => {
+    execFile('/bin/sh', ['-c', command], { timeout: 10_000 }, (error, stdout) => {
+      resolve(error ? {} : parseHelperHeaders(stdout))
+    })
+  })
+}
+
+/** UI seams the OAuth flow needs; absent in headless runs, which fail with advice. */
+export interface AuthUi {
+  confirm: (title: string, body: string) => Promise<boolean>
+  notify: (message: string, level: 'info' | 'warning' | 'error') => void
+}
+
+/** A server needs OAuth pi could not complete (headless, declined, or the flow
+ * failed). A typed marker so the SSE-fallback caller can tell an auth failure
+ * from a transport mismatch without matching on message text. */
+export class OAuthRequiredError extends Error {}
+
+/** Whether a connect failure is an authentication problem: the SDK's own
+ * UnauthorizedError, a transport error carrying HTTP 401 (which is what a 401
+ * throws when no authProvider was attached, so a first-time login is detected),
+ * or our own marker. */
+export function isUnauthorized(error: unknown): boolean {
+  if (error instanceof UnauthorizedError || error instanceof OAuthRequiredError) return true
+  return typeof error === 'object' && error !== null && (error as { code?: unknown }).code === 401
+}
+
+// SSEClientTransport is deprecated in favour of Streamable HTTP, but both concrete
+// transports expose finishAuth (the base Transport interface does not), so the union
+// stays as the http-family fallback type through the migration period.
+type HttpFamilyTransport = SSEClientTransport | StreamableHTTPClientTransport // NOSONAR typescript:S1874 - SSE fallback still required by the MCP SDK
+
+export type MakeTransport = (authProvider?: OAuthClientProvider) => HttpFamilyTransport
+
+/**
+ * Connect an http-family server, running Claude's OAuth login when the server
+ * demands one. Stored tokens ride the first attempt so the SDK refreshes
+ * silently; a 401 without tokens asks the user, opens the browser, catches the
+ * loopback redirect, and exchanges the code via the SDK's finishAuth.
+ * Bearer-token servers never enter the OAuth path: an explicit token is the
+ * user saying how auth works.
+ */
+async function connectHttpFamily(name: string, config: { url: string }, makeTransport: MakeTransport, label: string, bearerToken: string | undefined, authUi: AuthUi | undefined): Promise<Client> {
+  const newClient = () => new Client({ name: 'pi-code-mcp', version: '0.1.0' })
+  // 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 = bearerToken ? undefined : new FileOAuthProvider(name, () => {})
+  try {
+    const client = newClient()
+    await connectWithTimeout(client, makeTransport(silent?.hasTokens() ? silent : undefined), label)
+    return client
+  } catch (error) {
+    if (bearerToken || !isUnauthorized(error)) throw error
+    if (!authUi) throw new OAuthRequiredError(`${name} requires a login; run pi interactively to authenticate`)
+    return await serializeInteractiveOAuth(() => runInteractiveOAuth(name, config, makeTransport, label, authUi, newClient))
+  }
+}
+
+/**
+ * Connect with a deadline, closing the client if the deadline (not a connect error) wins.
+ * Without this, a slow-but-successful server finishes connecting after the race is lost and
+ * lingers unreferenced: process/socket alive, never in `clients`, invisible to shutdown.
+ */
+export async function connectWithTimeout(client: Client, transport: Parameters<Client['connect']>[0], label: string): Promise<void> {
+  const connecting = client.connect(transport)
+  try {
+    await withTimeout(connecting, connectTimeoutMs(), label)
+  } catch (error) {
+    // Only a timeout can orphan a still-opening transport; a connect rejection means the
+    // SDK already tore it down, so closing again would be redundant.
+    if (String(error).includes('timed out after')) {
+      connecting.catch(() => {}) // a late rejection must not surface as unhandled
+      void client.close().catch(() => {})
+    }
+    throw error
+  }
+}