fix: re-initialize the session when the server expires it

A 404 to a request carrying a session id means the server dropped the
session. The spec requires the client to start a new one with a fresh
InitializeRequest; mcpProxy instead logged the error and left the
request unanswered, so every later call failed too.

Capture the client's initialize, and on a 404 clear the session id,
replay it under a sentinel id, send notifications/initialized and retry
the original message. Sends that fail for other reasons now return a
JSON-RPC error rather than hanging.

Author: hfreire

src/lib/utils.test.ts

@@ -1,6 +1,7 @@
 import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest'
 import { parseCommandLineArgs, shouldIncludeTool, mcpProxy, setupOAuthCallbackServerWithLongPoll, getServerUrlHash } from './utils'
 import { Transport } from '@modelcontextprotocol/sdk/shared/transport.js'
+import { StreamableHTTPError } from '@modelcontextprotocol/sdk/client/streamableHttp.js'
 import { EventEmitter } from 'events'
 import express from 'express'
 
@@ -979,6 +980,122 @@ describe('Feature: MCP Proxy', () => {
       }),
     )
   })
+
+  it('Scenario: Re-establish the session when the server has expired it', async () => {
+    // Given a client transport
+    const mockTransportToClient = {
+      send: vi.fn().mockResolvedValue(undefined),
+      close: vi.fn().mockResolvedValue(undefined),
+      start: vi.fn().mockResolvedValue(undefined),
+      onmessage: vi.fn(),
+      onclose: vi.fn(),
+      onerror: vi.fn(),
+    } as unknown as Transport
+
+    // And a server transport that 404s the first tools/call, as a server does
+    // once it has dropped the session, then answers the fresh handshake
+    const sent: any[] = []
+    let expireNextCall = true
+    const mockTransportToServer = {
+      send: vi.fn(async (message: any) => {
+        sent.push(message)
+        if (typeof message.id === 'string' && message.id.startsWith('mcp-remote-reinit-')) {
+          setTimeout(() => (mockTransportToServer as any).onmessage?.({ jsonrpc: '2.0', id: message.id, result: {} }), 0)
+          return
+        }
+        if (expireNextCall && message.method === 'tools/call') {
+          expireNextCall = false
+          throw new StreamableHTTPError(404, 'Error POSTing to endpoint: Session terminated')
+        }
+      }),
+      close: vi.fn().mockResolvedValue(undefined),
+      start: vi.fn().mockResolvedValue(undefined),
+      onmessage: vi.fn(),
+      onclose: vi.fn(),
+      onerror: vi.fn(),
+      sessionId: 'session-2',
+    } as unknown as Transport
+
+    mcpProxy({
+      transportToClient: mockTransportToClient,
+      transportToServer: mockTransportToServer,
+      ignoredTools: [],
+    })
+
+    // When the client initializes and then calls a tool
+    mockTransportToClient.onmessage?.({
+      jsonrpc: '2.0' as const,
+      method: 'initialize',
+      id: '1',
+      params: { clientInfo: { name: 'Test Client', version: '1.0.0' } },
+    } as any)
+    mockTransportToClient.onmessage?.({
+      jsonrpc: '2.0' as const,
+      method: 'tools/call',
+      id: '2',
+      params: { name: 'ping', arguments: {} },
+    } as any)
+
+    await vi.waitFor(() => expect(sent.map((m) => m.method)).toContain('notifications/initialized'))
+
+    // Then a fresh initialize was sent, carrying the client's own parameters
+    const reinitialize = sent.find((m) => typeof m.id === 'string' && m.id.startsWith('mcp-remote-reinit-'))
+    expect(reinitialize).toBeDefined()
+    expect(reinitialize.method).toBe('initialize')
+    expect(reinitialize.params.clientInfo.name).toContain('Test Client')
+
+    // And the call that triggered it was retried on the new session
+    expect(sent.filter((m) => m.method === 'tools/call' && m.id === '2')).toHaveLength(2)
+
+    // And the handshake response was consumed by the proxy, never shown to the client
+    expect(mockTransportToClient.send).not.toHaveBeenCalledWith(expect.objectContaining({ id: reinitialize.id }))
+  })
+
+  it('Scenario: Answer the client when a request cannot be delivered', async () => {
+    // Given a server transport that fails for a reason a new session cannot fix
+    const mockTransportToClient = {
+      send: vi.fn().mockResolvedValue(undefined),
+      close: vi.fn().mockResolvedValue(undefined),
+      start: vi.fn().mockResolvedValue(undefined),
+      onmessage: vi.fn(),
+      onclose: vi.fn(),
+      onerror: vi.fn(),
+    } as unknown as Transport
+
+    const mockTransportToServer = {
+      send: vi.fn().mockRejectedValue(new Error('connection reset')),
+      close: vi.fn().mockResolvedValue(undefined),
+      start: vi.fn().mockResolvedValue(undefined),
+      onmessage: vi.fn(),
+      onclose: vi.fn(),
+      onerror: vi.fn(),
+    } as unknown as Transport
+
+    mcpProxy({
+      transportToClient: mockTransportToClient,
+      transportToServer: mockTransportToServer,
+      ignoredTools: [],
+    })
+
+    // When the client sends a request
+    mockTransportToClient.onmessage?.({
+      jsonrpc: '2.0' as const,
+      method: 'tools/call',
+      id: '7',
+      params: { name: 'ping', arguments: {} },
+    } as any)
+
+    // Then it gets an error rather than waiting forever for a reply
+    await vi.waitFor(() =>
+      expect(mockTransportToClient.send).toHaveBeenCalledWith(
+        expect.objectContaining({
+          jsonrpc: '2.0',
+          id: '7',
+          error: expect.objectContaining({ code: -32001 }),
+        }),
+      ),
+    )
+  })
 })
 
 describe('setupOAuthCallbackServerWithLongPoll', () => {

src/lib/utils.ts

@@ -139,6 +139,9 @@ export function mcpProxy({
 }) {
   let transportToClientClosed = false
   let transportToServerClosed = false
+  let lastInitialize: Message | null = null
+  let reinitSeq = 0
+  const pendingReinit = new Map<string, (message: Message) => void>()
 
   const messageTransformer = createMessageTransformer({
     transformRequestFunction: (request: Message) => {
@@ -199,12 +202,23 @@ export function mcpProxy({
       log(JSON.stringify(message, null, 2))
 
       debugLog('Initialize message with modified client info', { clientInfo })
+
+      lastInitialize = message
     }
 
-    transportToServer.send(message).catch(onServerError)
+    sendToServer(message)
   }
 
   transportToServer.onmessage = (_message) => {
+    // Responses to our own re-initialize handshake are ours to consume, not the client's
+    const reinitId = (_message as any).id
+    if (typeof reinitId === 'string' && pendingReinit.has(reinitId)) {
+      const settle = pendingReinit.get(reinitId)!
+      pendingReinit.delete(reinitId)
+      settle(_message as any)
+      return
+    }
+
     // TODO: fix types
     const message = messageTransformer.interceptResponse(_message as any)
     log('[Remote→Local]', message.method || message.id)
@@ -250,6 +264,91 @@ export function mcpProxy({
     log('Error from remote server:', error)
     debugLog('Error from remote server', { stack: error.stack })
   }
+
+  /**
+   * A 404 to a request that carried a session id means the server dropped the
+   * session (idle expiry, restart, eviction). The spec says the client must then
+   * start a new session with a fresh InitializeRequest.
+   */
+  function isSessionExpired(error: Error) {
+    return error instanceof StreamableHTTPError && error.code === 404
+  }
+
+  async function reinitializeSession() {
+    if (!lastInitialize) {
+      throw new Error('no initialize request was seen, cannot re-establish the session')
+    }
+
+    // Must be cleared before we send, or the transport re-attaches the dead id
+    // and the server 404s the handshake too. The SDK exposes sessionId read-only.
+    ;(transportToServer as unknown as { _sessionId?: string })._sessionId = undefined
+
+    const id = `mcp-remote-reinit-${++reinitSeq}`
+    const response = await new Promise<Message>((resolve, reject) => {
+      const timer = setTimeout(() => {
+        pendingReinit.delete(id)
+        reject(new Error('timed out waiting for the re-initialize response'))
+      }, 30000)
+      pendingReinit.set(id, (message) => {
+        clearTimeout(timer)
+        resolve(message)
+      })
+      transportToServer.send({ ...lastInitialize, id }).catch((error) => {
+        clearTimeout(timer)
+        pendingReinit.delete(id)
+        reject(error)
+      })
+    })
+
+    if (response.error) {
+      throw new Error(`server rejected re-initialize: ${JSON.stringify(response.error)}`)
+    }
+
+    await transportToServer.send({ jsonrpc: '2.0', method: 'notifications/initialized' })
+    log(`Re-established session ${transportToServer.sessionId ?? '(none)'} after server expiry`)
+  }
+
+  async function sendToServer(message: Message) {
+    try {
+      await transportToServer.send(message)
+      return
+    } catch (error) {
+      // Re-initializing in response to a failed initialize would loop
+      if (!isSessionExpired(error as Error) || message.method === 'initialize') {
+        onServerError(error as Error)
+        replyWithError(message, error as Error)
+        return
+      }
+
+      log('Remote session expired, re-initializing')
+      debugLog('Remote session expired', { id: message.id, method: message.method })
+
+      try {
+        await reinitializeSession()
+        await transportToServer.send(message)
+      } catch (retryError) {
+        onServerError(retryError as Error)
+        replyWithError(message, retryError as Error)
+      }
+    }
+  }
+
+  /**
+   * Without this a failed send leaves the client waiting forever on a request
+   * that will never be answered.
+   */
+  function replyWithError(message: Message, error: Error) {
+    if (message.id === undefined || message.id === null) {
+      return
+    }
+    transportToClient
+      .send({
+        jsonrpc: '2.0',
+        id: message.id,
+        error: { code: -32001, message: `mcp-remote: ${error.message ?? String(error)}` },
+      })
+      .catch(onClientError)
+  }
 }
 
 /**