Skip to content

Commit a79e7ba

Browse files
committed
feat: mcp proxy keep alive (ping) mechanism
1 parent 5c71b26 commit a79e7ba

File tree

3 files changed

+85
-4
lines changed

3 files changed

+85
-4
lines changed

src/lib/types.ts

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,3 +33,11 @@ export interface OAuthCallbackServerOptions {
3333
/** Event emitter to signal when auth code is received */
3434
events: EventEmitter
3535
}
36+
37+
/*
38+
* Configuration for the ping mechanism
39+
*/
40+
export interface PingConfig {
41+
enabled: boolean
42+
interval: number
43+
}

src/lib/utils.ts

Lines changed: 65 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,10 +3,12 @@ import { Client } from '@modelcontextprotocol/sdk/client/index.js'
33
import { SSEClientTransport } from '@modelcontextprotocol/sdk/client/sse.js'
44
import { StreamableHTTPClientTransport } from '@modelcontextprotocol/sdk/client/streamableHttp.js'
55
import { Transport } from '@modelcontextprotocol/sdk/shared/transport.js'
6+
import { PingConfig } from './types'
67

78
// Connection constants
89
export const REASON_AUTH_NEEDED = 'authentication-needed'
910
export const REASON_TRANSPORT_FALLBACK = 'falling-back-to-alternate-transport'
11+
export const PING_INTERVAL_DEFAULT = 30000
1012

1113
// Transport strategy types
1214
export type TransportStrategy = 'sse-only' | 'http-only' | 'sse-first' | 'http-first'
@@ -81,6 +83,44 @@ export type AuthInitializer = () => Promise<{
8183
skipBrowserAuth: boolean
8284
}>
8385

86+
/**
87+
* Sets up periodic ping to keep the connection alive
88+
* @param transport The transport to ping
89+
* @param config Ping configuration
90+
* @returns A cleanup function to stop pinging
91+
*/
92+
export function setupPing(transport: Transport, config: PingConfig): () => void {
93+
if (!config.enabled) {
94+
return () => {}
95+
}
96+
97+
let pingTimeout: NodeJS.Timeout | null = null
98+
let lastPingId = 0
99+
100+
const interval = config.interval * 1000 // convert ms to s
101+
const pingInterval = setInterval(async () => {
102+
const pingId = ++lastPingId
103+
try {
104+
// Docs: https://modelcontextprotocol.io/specification/2025-03-26/basic/utilities/ping
105+
await transport.send({
106+
jsonrpc: '2.0',
107+
id: `ping-${pingId}`,
108+
method: 'ping',
109+
})
110+
log(`Ping ${pingId} successful`)
111+
} catch (error) {
112+
log(`Ping ${pingId} failed:`, error)
113+
}
114+
}, interval)
115+
116+
return () => {
117+
if (pingTimeout) {
118+
clearTimeout(pingTimeout)
119+
}
120+
clearInterval(pingInterval)
121+
}
122+
}
123+
84124
/**
85125
* Creates and connects to a remote server with OAuth authentication
86126
* @param client The client to connect with
@@ -402,6 +442,21 @@ export async function parseCommandLineArgs(args: string[], defaultPort: number,
402442
i++
403443
}
404444

445+
// Parse ping configuration
446+
const keepAlive = args.includes('--keep-alive')
447+
const pingIntervalIndex = args.indexOf('--ping-interval')
448+
let pingInterval = PING_INTERVAL_DEFAULT
449+
if (pingIntervalIndex !== -1 && pingIntervalIndex < args.length - 1) {
450+
const intervalStr = args[pingIntervalIndex + 1]
451+
const interval = parseInt(intervalStr)
452+
if (!isNaN(interval) && interval > 0) {
453+
pingInterval = interval
454+
log(`Using ping interval: ${pingInterval} seconds`)
455+
} else {
456+
log(`Warning: Invalid ping interval "${args[pingIntervalIndex + 1]}". Using default: ${PING_INTERVAL_DEFAULT} seconds`)
457+
}
458+
}
459+
405460
const serverUrl = args[0]
406461
const specifiedPort = args[1] ? parseInt(args[1]) : undefined
407462
const allowHttp = args.includes('--allow-http')
@@ -461,7 +516,16 @@ export async function parseCommandLineArgs(args: string[], defaultPort: number,
461516
})
462517
}
463518

464-
return { serverUrl, callbackPort, headers, transportStrategy }
519+
return {
520+
serverUrl,
521+
callbackPort,
522+
headers,
523+
transportStrategy,
524+
pingConfig: {
525+
enabled: keepAlive,
526+
interval: pingInterval,
527+
},
528+
}
465529
}
466530

467531
/**

src/proxy.ts

Lines changed: 12 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -18,11 +18,12 @@ import {
1818
parseCommandLineArgs,
1919
setupSignalHandlers,
2020
getServerUrlHash,
21-
MCP_REMOTE_VERSION,
2221
TransportStrategy,
22+
setupPing,
2323
} from './lib/utils'
2424
import { NodeOAuthClientProvider } from './lib/node-oauth-client-provider'
2525
import { createLazyAuthCoordinator } from './lib/coordination'
26+
import { PingConfig } from './lib/types'
2627

2728
/**
2829
* Main function to run the proxy
@@ -32,6 +33,7 @@ async function runProxy(
3233
callbackPort: number,
3334
headers: Record<string, string>,
3435
transportStrategy: TransportStrategy = 'http-first',
36+
pingConfig: PingConfig,
3537
) {
3638
// Set up event emitter for auth flow
3739
const events = new EventEmitter()
@@ -80,6 +82,9 @@ async function runProxy(
8082
// Connect to remote server with lazy authentication
8183
const remoteTransport = await connectToRemoteServer(null, serverUrl, authProvider, headers, authInitializer, transportStrategy)
8284

85+
// Set up ping mechanism for remote transport
86+
const stopPing = setupPing(remoteTransport, pingConfig)
87+
8388
// Set up bidirectional proxy between local and remote transports
8489
mcpProxy({
8590
transportToClient: localTransport,
@@ -89,11 +94,15 @@ async function runProxy(
8994
// Start the local STDIO server
9095
await localTransport.start()
9196
log('Local STDIO server running')
97+
if (pingConfig.enabled) {
98+
log(`Automatic ping enabled with ${pingConfig.interval} second interval`)
99+
}
92100
log(`Proxy established successfully between local STDIO and remote ${remoteTransport.constructor.name}`)
93101
log('Press Ctrl+C to exit')
94102

95103
// Setup cleanup handler
96104
const cleanup = async () => {
105+
stopPing()
97106
await remoteTransport.close()
98107
await localTransport.close()
99108
// Only close the server if it was initialized
@@ -136,8 +145,8 @@ to the CA certificate file. If using claude_desktop_config.json, this might look
136145

137146
// Parse command-line arguments and run the proxy
138147
parseCommandLineArgs(process.argv.slice(2), 3334, 'Usage: npx tsx proxy.ts <https://server-url> [callback-port]')
139-
.then(({ serverUrl, callbackPort, headers, transportStrategy }) => {
140-
return runProxy(serverUrl, callbackPort, headers, transportStrategy)
148+
.then(({ serverUrl, callbackPort, headers, transportStrategy, pingConfig }) => {
149+
return runProxy(serverUrl, callbackPort, headers, transportStrategy, pingConfig)
141150
})
142151
.catch((error) => {
143152
log('Fatal error:', error)

0 commit comments

Comments
 (0)