fix: reuse pooled MCP client connections in pi extension

packages/pi/main.ts

@@ -1,6 +1,5 @@
 import type { ExtensionAPI } from '@earendil-works/pi-coding-agent'
-import { Client } from '@modelcontextprotocol/sdk/client/index.js'
-import { StreamableHTTPClientTransport } from '@modelcontextprotocol/sdk/client/streamableHttp.js'
+import { Client, type FetchLike, StreamableHTTPClientTransport } from '@modelcontextprotocol/client'
 import { Type } from 'typebox'
 import packageJson from './package.json' with { type: 'json' }
 
@@ -9,6 +8,7 @@ type PiToolDefinition = Parameters<ExtensionAPI['registerTool']>[0]
 type McpBridgeConfig = Omit<PiToolDefinition, 'execute' | 'parameters'> & {
   url?: string
   authenticated?: boolean
+  fetch?: FetchLike
 }
 
 type McpTool = {
@@ -20,6 +20,8 @@ type McpTool = {
 type McpServerConfig = {
   url: string
   authenticated?: boolean
+  /** Custom fetch for the Streamable HTTP transport — proxies, tests, non-standard runtimes. */
+  fetch?: FetchLike
   promptGuidelines: string[]
   /** Override the name the tool is registered under in Pi. The MCP callTool still uses the original server tool name. */
   registerAs?: (tool: McpTool) => string
@@ -55,23 +57,83 @@ const createHeaders = ({ authenticated = true }: { authenticated?: boolean } = {
   }
 }
 
-const discoveredToolsCache = new Map<string, Promise<McpTool[]>>()
+type McpServerTarget = Pick<McpServerConfig, 'authenticated' | 'fetch' | 'url'>
+
+type McpConnection = { client: Client; transport: StreamableHTTPClientTransport }
 
-const withMcpClient = async <T>(
-  { authenticated, url }: { authenticated?: boolean; url: string },
-  fn: (client: Client) => Promise<T>,
-): Promise<T> => {
+const connectMcpClient = async ({ authenticated, fetch, url }: McpServerTarget): Promise<McpConnection> => {
   const client = new Client(CLIENT_INFO)
   const transport = new StreamableHTTPClientTransport(new URL(url), {
     requestInit: { headers: createHeaders({ authenticated }) },
+    ...(fetch ? { fetch } : {}),
   })
 
   await client.connect(transport)
+  return { client, transport }
+}
+
+const closeMcpConnection = async ({ client, transport }: McpConnection) => {
+  // Over Streamable HTTP the clean disconnect terminates the server-side
+  // session before tearing down the transport; it is a no-op when the server
+  // never issued a session ID.
+  await transport.terminateSession().catch(() => {})
+  await client.close().catch(() => {})
+}
+
+const discoveredToolsCache = new Map<string, Promise<McpTool[]>>()
+
+const withMcpClient = async <T>(server: McpServerTarget, fn: (client: Client) => Promise<T>): Promise<T> => {
+  const connection = await connectMcpClient(server)
 
   try {
-    return await fn(client)
+    return await fn(connection.client)
   } finally {
-    await client.close()
+    await closeMcpConnection(connection)
+  }
+}
+
+// Lazily connected clients shared across callTool executions so the MCP
+// initialize handshake is paid once per server instead of once per call.
+// Closed from the session_shutdown handler registered in the extension entry point.
+const mcpClients = new Map<string, Promise<McpConnection>>()
+
+const mcpCacheKey = ({ authenticated, url }: McpServerTarget) => `${authenticated === false ? 'public' : 'auth'}:${url}`
+
+const getMcpClient = (server: McpServerTarget) => {
+  const cacheKey = mcpCacheKey(server)
+  let connection = mcpClients.get(cacheKey)
+  if (!connection) {
+    connection = connectMcpClient(server)
+    mcpClients.set(cacheKey, connection)
+    connection.catch(() => mcpClients.delete(cacheKey))
+  }
+  return connection
+}
+
+const resetMcpClient = async (server: McpServerTarget) => {
+  const cacheKey = mcpCacheKey(server)
+  const connection = mcpClients.get(cacheKey)
+  mcpClients.delete(cacheKey)
+  await connection?.then(closeMcpConnection).catch(() => {})
+}
+
+const closeMcpClients = async () => {
+  const connections = [...mcpClients.values()]
+  mcpClients.clear()
+  await Promise.all(connections.map((connection) => connection.then(closeMcpConnection).catch(() => {})))
+}
+
+const withPooledMcpClient = async <T>(server: McpServerTarget, fn: (client: Client) => Promise<T>): Promise<T> => {
+  try {
+    return await fn((await getMcpClient(server)).client)
+  } catch {
+    // A long-lived StreamableHTTP session can go stale (server restart, session
+    // expiry). Drop the cached client and retry once on a fresh connection.
+    // MINIMAL: any failure (not just transport errors) triggers one reconnect;
+    // permanent tool errors simply surface again after the retry. Upgrade path:
+    // only reset on SdkError/transport failures and rethrow ProtocolError.
+    await resetMcpClient(server)
+    return await fn((await getMcpClient(server)).client)
   }
 }
 
@@ -104,9 +166,10 @@ const registerMcpTool = (pi: ExtensionAPI, definition: McpBridgeConfig & { tool:
         throw new Error('params must be an object')
       }
 
-      const result = await withMcpClient(
+      const result = await withPooledMcpClient(
         {
           authenticated: definition.authenticated,
+          fetch: definition.fetch,
           url: definition.url ?? MCP_URL,
         },
         async (client) => await client.callTool({ name: definition.tool.name, arguments: params }),
@@ -123,6 +186,7 @@ const registerMcpServerTools = async (pi: ExtensionAPI, server: McpServerConfig)
     registerMcpTool(pi, {
       description: server.promptGuidelines[0] ?? `Use ${tool.name} for You.com MCP calls.`,
       authenticated: server.authenticated,
+      fetch: server.fetch,
       label: registeredName,
       name: registeredName,
       promptGuidelines: server.promptGuidelines,
@@ -158,8 +222,8 @@ const SERVER_CONFIGS: McpServerConfig[] = [
   },
 ]
 
-const registerMcpTools = async (pi: ExtensionAPI) => {
-  await Promise.all(SERVER_CONFIGS.map((config) => registerMcpServerTools(pi, config)))
+const registerMcpTools = async (pi: ExtensionAPI, servers: McpServerConfig[]) => {
+  await Promise.all(servers.map((config) => registerMcpServerTools(pi, config)))
 }
 
 const HOST_CONTEXT = [
@@ -184,14 +248,19 @@ const registerHostContext = (pi: ExtensionAPI) => {
  * Registers the minimal You.com MCP bridge and bundled Pi skill resources.
  *
  * @param pi - Pi extension API.
+ * @param servers - MCP servers to bridge; defaults to the You.com endpoints.
  *
  * @public
  */
-export default async function youPiPlugin(pi: ExtensionAPI) {
+export default async function youPiPlugin(pi: ExtensionAPI, servers: McpServerConfig[] = SERVER_CONFIGS) {
   pi.on('resources_discover', () => ({
     skillPaths: [SKILLS_PATH],
   }))
 
-  await registerMcpTools(pi)
+  // Reload, quit, and session switches all fire session_shutdown; close the
+  // pooled clients there (idempotent) and reconnect lazily on the next call.
+  pi.on('session_shutdown', closeMcpClients)
+
+  await registerMcpTools(pi, servers)
   registerHostContext(pi)
 }

packages/pi/package.json

@@ -49,10 +49,11 @@
     ]
   },
   "dependencies": {
-    "@modelcontextprotocol/sdk": "^1.29.0",
+    "@modelcontextprotocol/client": "^2.0.0",
     "typebox": "^1.3.6"
   },
   "devDependencies": {
-    "@earendil-works/pi-coding-agent": "^0.81.1"
+    "@earendil-works/pi-coding-agent": "^0.81.1",
+    "@modelcontextprotocol/server": "^2.0.0"
   }
 }

packages/pi/tests/mcp-connection.spec.ts

@@ -0,0 +1,162 @@
+import { describe, expect, test } from 'bun:test'
+import type { ExtensionAPI } from '@earendil-works/pi-coding-agent'
+import { createMcpHandler, fromJsonSchema, McpServer } from '@modelcontextprotocol/server'
+
+type RegisteredTool = {
+  name: string
+  execute: (_toolCallId: string, params: unknown) => Promise<unknown>
+}
+
+type RegisteredEvent = {
+  eventName: string
+  handler: (...args: never[]) => unknown
+}
+
+const loadExtension = async () => (await import(`../main.ts?test=${Date.now()}-${Math.random()}`)).default
+
+const createPiHarness = () => {
+  const events: RegisteredEvent[] = []
+  const tools: RegisteredTool[] = []
+
+  const pi = {
+    on: (eventName: string, handler: (...args: never[]) => unknown) => {
+      events.push({ eventName, handler })
+    },
+    registerTool: (tool: RegisteredTool) => {
+      tools.push(tool)
+    },
+  } as unknown as ExtensionAPI
+
+  return { events, pi, tools }
+}
+
+/** Real in-process MCP server: the extension's transport fetch is wired straight to handler.fetch. */
+const createTestServer = () => {
+  let initializeCount = 0
+  let failNextToolCall = false
+
+  const handler = createMcpHandler(() => {
+    const server = new McpServer({ name: 'fixture', version: '1.0.0' })
+    server.registerTool(
+      'echo',
+      {
+        description: 'Echo the query back',
+        inputSchema: fromJsonSchema<{ query: string }>({
+          type: 'object',
+          properties: { query: { type: 'string' } },
+          required: ['query'],
+        }),
+      },
+      async ({ query }) => ({ content: [{ type: 'text' as const, text: `echo:${query}` }] }),
+    )
+    return server
+  })
+
+  const fetchHandler = async (url: string | URL, init?: RequestInit) => {
+    const request = new Request(String(url), init)
+    if (request.method === 'POST') {
+      const body = (await request.clone().json()) as { method?: string }
+      if (body.method === 'initialize') initializeCount += 1
+      if (body.method === 'tools/call' && failNextToolCall) {
+        failNextToolCall = false
+        throw new TypeError('fetch failed')
+      }
+    }
+    return handler.fetch(request)
+  }
+
+  return {
+    close: () => handler.close(),
+    failNextToolCall: () => {
+      failNextToolCall = true
+    },
+    initializeCount: () => initializeCount,
+    serverConfig: {
+      url: 'http://fixture.local/mcp',
+      authenticated: false,
+      fetch: fetchHandler,
+      promptGuidelines: ['fixture server'],
+    },
+  }
+}
+
+const findTool = (tools: RegisteredTool[], name: string) => {
+  const tool = tools.find((registeredTool) => registeredTool.name === name)
+  expect(tool).toBeDefined()
+  if (!tool) throw new Error(`${name} tool was not registered`)
+  return tool
+}
+
+describe('MCP connection lifecycle', () => {
+  test('reuses one MCP connection across tool executions', async () => {
+    const server = createTestServer()
+    try {
+      const extension = await loadExtension()
+      const { pi, tools } = createPiHarness()
+
+      await extension(pi, [server.serverConfig])
+      // Tool discovery pays its own one-shot handshake; start counting from here.
+      const afterDiscovery = server.initializeCount()
+
+      const tool = findTool(tools, 'echo')
+      const first = (await tool.execute('call-1', { query: 'one' })) as { content: Array<{ text: string }> }
+      const second = (await tool.execute('call-2', { query: 'two' })) as { content: Array<{ text: string }> }
+
+      expect(first.content[0]?.text).toBe('echo:one')
+      expect(second.content[0]?.text).toBe('echo:two')
+      expect(server.initializeCount()).toBe(afterDiscovery + 1)
+    } finally {
+      await server.close()
+    }
+  })
+
+  test('reconnects and retries once when the connection drops', async () => {
+    const server = createTestServer()
+    try {
+      const extension = await loadExtension()
+      const { pi, tools } = createPiHarness()
+
+      await extension(pi, [server.serverConfig])
+      const afterDiscovery = server.initializeCount()
+
+      const tool = findTool(tools, 'echo')
+      await tool.execute('call-1', { query: 'one' })
+      expect(server.initializeCount()).toBe(afterDiscovery + 1)
+
+      server.failNextToolCall()
+      const result = (await tool.execute('call-2', { query: 'two' })) as { content: Array<{ text: string }> }
+
+      expect(result.content[0]?.text).toBe('echo:two')
+      // The dropped pooled client was replaced: one new handshake, then success.
+      expect(server.initializeCount()).toBe(afterDiscovery + 2)
+    } finally {
+      await server.close()
+    }
+  })
+
+  test('closes pooled clients on session_shutdown', async () => {
+    const server = createTestServer()
+    try {
+      const extension = await loadExtension()
+      const { events, pi, tools } = createPiHarness()
+
+      await extension(pi, [server.serverConfig])
+      const afterDiscovery = server.initializeCount()
+
+      const tool = findTool(tools, 'echo')
+      await tool.execute('call-1', { query: 'one' })
+      expect(server.initializeCount()).toBe(afterDiscovery + 1)
+
+      const shutdown = events.find((event) => event.eventName === 'session_shutdown')
+      expect(shutdown).toBeDefined()
+      if (!shutdown) throw new Error('session_shutdown handler was not registered')
+      await shutdown.handler()
+
+      await tool.execute('call-2', { query: 'two' })
+      // Shutdown dropped the pooled client: the next execution reconnects.
+      expect(server.initializeCount()).toBe(afterDiscovery + 2)
+    } finally {
+      await server.close()
+    }
+  })
+})