Adding support for Streamable MCP server

.changeset/silver-trains-learn.md

@@ -0,0 +1,5 @@
+---
+"claude-dev": patch
+---
+
+adding support for streamable mcp server

package.json

@@ -318,7 +318,7 @@
 		"@typescript-eslint/eslint-plugin": "^7.14.1",
 		"@typescript-eslint/parser": "^7.11.0",
 		"@vscode/test-cli": "^0.0.10",
-		"@vscode/test-electron": "^2.4.1",
+		"@vscode/test-electron": "^2.5.2",
 		"chai": "^4.3.10",
 		"chalk": "^5.3.0",
 		"esbuild": "^0.25.0",
@@ -341,15 +341,16 @@
 		"@anthropic-ai/bedrock-sdk": "^0.12.4",
 		"@anthropic-ai/sdk": "^0.37.0",
 		"@anthropic-ai/vertex-sdk": "^0.6.4",
-		"@aws-sdk/client-bedrock-runtime": "^3.758.0",
+		"@aws-sdk/client-bedrock-runtime": "^3.779.0",
 		"@bufbuild/protobuf": "^2.2.5",
 		"@cerebras/cerebras_cloud_sdk": "^1.35.0",
 		"@google-cloud/vertexai": "^1.9.3",
 		"@google/genai": "^0.13.0",
 		"@grpc/grpc-js": "^1.9.15",
 		"@grpc/reflection": "^1.0.4",
-		"@mistralai/mistralai": "^1.5.0",
-		"@modelcontextprotocol/sdk": "^1.11.1",
+		"@mistralai/mistralai": "^1.3.6",
+		"@modelcontextprotocol/sdk": "^1.11.4",
+		"reconnecting-eventsource": "^1.6.4",
 		"@opentelemetry/api": "^1.4.1",
 		"@opentelemetry/exporter-trace-otlp-http": "^0.39.1",
 		"@opentelemetry/resources": "^1.30.1",
@@ -359,6 +360,7 @@
 		"@sentry/browser": "^9.12.0",
 		"@streamparser/json": "^0.0.22",
 		"@vscode/codicons": "^0.0.36",
+		"@vscode/vsce": "^3.4.1",
 		"archiver": "^7.0.1",
 		"axios": "^1.8.2",
 		"cheerio": "^1.0.0",

src/services/mcp/McpHub.ts

@@ -1,5 +1,8 @@
 import { Client } from "@modelcontextprotocol/sdk/client/index.js"
 import { StdioClientTransport } from "@modelcontextprotocol/sdk/client/stdio.js"
+import { SSEClientTransport } from "@modelcontextprotocol/sdk/client/sse.js"
+import { StreamableHTTPClientTransport } from "@modelcontextprotocol/sdk/client/streamableHttp.js"
+import ReconnectingEventSource from "reconnecting-eventsource"
 import {
 	CallToolResultSchema,
 	ListResourcesResultSchema,
@@ -29,17 +32,18 @@ import { fileExistsAtPath } from "@utils/fs"
 import { arePathsEqual } from "@utils/path"
 import { secondsToMs } from "@utils/time"
 import { GlobalFileNames } from "@core/storage/disk"
-import { SSEClientTransport } from "@modelcontextprotocol/sdk/client/sse.js"
-import { StreamableHTTPClientTransport } from "@modelcontextprotocol/sdk/client/streamableHttp.js"
+// import { injectEnv } from "../../utils/config" // Commented out as injectEnv is not found
 import { ExtensionMessage } from "@shared/ExtensionMessage"
 
 // Default timeout for internal MCP data requests in milliseconds; is not the same as the user facing timeout stored as DEFAULT_MCP_TIMEOUT_SECONDS
 const DEFAULT_REQUEST_TIMEOUT_MS = 5000
 
+type Transport = StdioClientTransport | SSEClientTransport | StreamableHTTPClientTransport
+
 export type McpConnection = {
 	server: McpServer
 	client: Client
-	transport: StdioClientTransport | SSEClientTransport | StreamableHTTPClientTransport
+	transport: Transport
 }
 
 export type McpTransportType = "stdio" | "sse" | "http"
@@ -48,37 +52,91 @@ export type McpServerConfig = z.infer<typeof ServerConfigSchema>
 
 const AutoApproveSchema = z.array(z.string()).default([])
 
+// Base configuration schema for common settings
 const BaseConfigSchema = z.object({
 	autoApprove: AutoApproveSchema.optional(),
 	disabled: z.boolean().optional(),
 	timeout: z.number().min(MIN_MCP_TIMEOUT_SECONDS).optional().default(DEFAULT_MCP_TIMEOUT_SECONDS),
 })
 
-const SseConfigSchema = BaseConfigSchema.extend({
-	url: z.string().url(),
-}).transform((config) => ({
-	...config,
-	transportType: "sse" as const,
-}))
+// Custom error messages for better user feedback
+const typeErrorMessage = "Server type must be one of: 'stdio', 'sse', or 'streamableHttp'"
+const stdioFieldsErrorMessage =
+	"For 'stdio' type servers, you must provide a 'command' field and can optionally include 'args' and 'env'"
+const urlFieldsErrorMessage = "For url based type servers, you must provide a 'url' field and can optionally include 'headers'"
+const mixedFieldsErrorMessage =
+	"Cannot mix 'stdio' and 'sse' fields. For 'stdio' use 'command', 'args', and 'env'. For 'sse' use 'url' and 'headers'"
+const missingFieldsErrorMessage = "Server configuration must include either 'command' (for stdio) or 'url' (for sse)"
 
-const StdioConfigSchema = BaseConfigSchema.extend({
-	command: z.string(),
-	args: z.array(z.string()).optional(),
-	env: z.record(z.string()).optional(),
-}).transform((config) => ({
-	...config,
-	transportType: "stdio" as const,
-}))
+function inferUrlBasedType(config: any): "sse" | "streamableHttp" | null {
+	if (!config.headers || typeof config.headers !== "object") {
+		return "streamableHttp"
+	}
 
-const StreamableHTTPConfigSchema = BaseConfigSchema.extend({
-	transportType: z.literal("http"),
-	url: z.string().url(),
-}).transform((config) => ({
-	...config,
-	transportType: "http" as const,
-}))
+	const headers = Object.fromEntries(Object.entries(config.headers).map(([k, v]) => [k.toLowerCase(), v]))
 
-const ServerConfigSchema = z.union([StdioConfigSchema, SseConfigSchema, StreamableHTTPConfigSchema])
+	if (typeof headers["accept"] === "string" && headers["accept"].includes("text/event-stream")) {
+		return "sse"
+	}
+
+	return "streamableHttp"
+}
+
+// Helper function to create a refined schema with better error messages
+const createServerTypeSchema = () => {
+	return z.union([
+		// Stdio config (has command field)
+		BaseConfigSchema.extend({
+			type: z.literal("stdio").optional(),
+			command: z.string(),
+			args: z.array(z.string()).optional(),
+			cwd: z.string().optional(),
+			env: z.record(z.string()).optional(),
+			// Explicitly disallow other types' fields
+			url: z.undefined().optional(),
+			headers: z.undefined().optional(),
+		})
+			.transform((data) => ({
+				...data,
+				type: "stdio" as const,
+			}))
+			.refine((data) => data.type === undefined || data.type === "stdio", { message: typeErrorMessage }),
+		// SSE config (has url field)
+		BaseConfigSchema.extend({
+			type: z.literal("sse").optional(),
+			url: z.string().url("URL must be a valid URL format"),
+			headers: z.record(z.string()).optional(),
+			// Explicitly disallow other types' fields
+			command: z.undefined().optional(),
+			args: z.undefined().optional(),
+			env: z.undefined().optional(),
+		})
+			.transform((data) => ({
+				...data,
+				type: "sse" as const,
+			}))
+			.refine((data) => data.type === undefined || data.type === "sse", { message: typeErrorMessage }),
+		// Streamable HTTP config (has url field)
+		BaseConfigSchema.extend({
+			type: z.literal("streamableHttp").optional(),
+			url: z.string().url("URL must be a valid URL format"),
+			headers: z.record(z.string()).optional(),
+			// Explicitly disallow other types' fields
+			command: z.undefined().optional(),
+			args: z.undefined().optional(),
+			env: z.undefined().optional(),
+		})
+			.transform((data) => ({
+				...data,
+				type: "streamableHttp" as const,
+			}))
+			.refine((data) => data.type === undefined || data.type === "streamableHttp", {
+				message: typeErrorMessage,
+			}),
+	])
+}
+
+const ServerConfigSchema = createServerTypeSchema()
 
 const McpSettingsSchema = z.object({
 	mcpServers: z.record(ServerConfigSchema),
@@ -190,125 +248,14 @@ export class McpHub {
 		}
 	}
 
-	private async connectToServerRPC(
-		name: string,
-		config: z.infer<typeof StdioConfigSchema> | z.infer<typeof SseConfigSchema> | z.infer<typeof StreamableHTTPConfigSchema>,
-	): Promise<void> {
-		// Remove existing connection if it exists (should never happen, the connection should be deleted beforehand)
-		this.connections = this.connections.filter((conn) => conn.server.name !== name)
-
-		try {
-			// Each MCP server requires its own transport connection and has unique capabilities, configurations, and error handling. Having separate clients also allows proper scoping of resources/tools and independent server management like reconnection.
-			const client = new Client(
-				{
-					name: "Cline",
-					version: this.clientVersion,
-				},
-				{
-					capabilities: {},
-				},
-			)
-
-			let transport: StdioClientTransport | SSEClientTransport | StreamableHTTPClientTransport
-
-			if (config.transportType === "sse") {
-				transport = new SSEClientTransport(new URL(config.url), {})
-			} else if (config.transportType === "http") {
-				transport = new StreamableHTTPClientTransport(new URL(config.url), {})
-			} else {
-				transport = new StdioClientTransport({
-					command: config.command,
-					args: config.args,
-					env: {
-						...config.env,
-						...(process.env.PATH ? { PATH: process.env.PATH } : {}),
-						// ...(process.env.NODE_PATH ? { NODE_PATH: process.env.NODE_PATH } : {}),
-					},
-					stderr: "pipe", // necessary for stderr to be available
-				})
-			}
-
-			transport.onerror = async (error) => {
-				console.error(`Transport error for "${name}":`, error)
-				const connection = this.connections.find((conn) => conn.server.name === name)
-				if (connection) {
-					connection.server.status = "disconnected"
-					this.appendErrorMessage(connection, error.message)
-				}
-			}
-
-			transport.onclose = async () => {
-				const connection = this.connections.find((conn) => conn.server.name === name)
-				if (connection) {
-					connection.server.status = "disconnected"
-				}
-			}
-
-			const connection: McpConnection = {
-				server: {
-					name,
-					config: JSON.stringify(config),
-					status: "connecting",
-					disabled: config.disabled,
-				},
-				client,
-				transport,
-			}
-			this.connections.push(connection)
-
-			if (config.transportType === "stdio") {
-				// transport.stderr is only available after the process has been started. However we can't start it separately from the .connect() call because it also starts the transport. And we can't place this after the connect call since we need to capture the stderr stream before the connection is established, in order to capture errors during the connection process.
-				// As a workaround, we start the transport ourselves, and then monkey-patch the start method to no-op so that .connect() doesn't try to start it again.
-				await transport.start()
-				const stderrStream = (transport as StdioClientTransport).stderr
-				if (stderrStream) {
-					stderrStream.on("data", async (data: Buffer) => {
-						const output = data.toString()
-						// Check if output contains INFO level log
-						const isInfoLog = !/\berror\b/i.test(output)
-
-						if (isInfoLog) {
-							// Log normal informational messages
-							console.info(`Server "${name}" info:`, output)
-						} else {
-							// Treat as error log
-							console.error(`Server "${name}" stderr:`, output)
-							const connection = this.connections.find((conn) => conn.server.name === name)
-							if (connection) {
-								this.appendErrorMessage(connection, output)
-							}
-						}
-					})
-				} else {
-					console.error(`No stderr stream for ${name}`)
-				}
-				transport.start = async () => {} // No-op now, .connect() won't fail
-			}
-
-			// Connect
-			await client.connect(transport)
-
-			connection.server.status = "connected"
-			connection.server.error = ""
-
-			// Initial fetch of tools and resources
-			connection.server.tools = await this.fetchToolsList(name)
-			connection.server.resources = await this.fetchResourcesList(name)
-			connection.server.resourceTemplates = await this.fetchResourceTemplatesList(name)
-		} catch (error) {
-			// Update status with error
-			const connection = this.connections.find((conn) => conn.server.name === name)
-			if (connection) {
-				connection.server.status = "disconnected"
-				this.appendErrorMessage(connection, error instanceof Error ? error.message : String(error))
-			}
-			throw error
-		}
+	private findConnection(name: string, source: "rpc" | "internal"): McpConnection | undefined {
+		return this.connections.find((conn) => conn.server.name === name)
 	}
 
 	private async connectToServer(
 		name: string,
-		config: z.infer<typeof StdioConfigSchema> | z.infer<typeof SseConfigSchema> | z.infer<typeof StreamableHTTPConfigSchema>,
+		config: z.infer<typeof ServerConfigSchema>,
+		source: "rpc" | "internal",
 	): Promise<void> {
 		// Remove existing connection if it exists (should never happen, the connection should be deleted beforehand)
 		this.connections = this.connections.filter((conn) => conn.server.name !== name)
@@ -327,39 +274,110 @@ export class McpHub {
 
 			let transport: StdioClientTransport | SSEClientTransport | StreamableHTTPClientTransport
 
-			if (config.transportType === "sse") {
-				transport = new SSEClientTransport(new URL(config.url), {})
-			} else if (config.transportType === "http") {
-				transport = new StreamableHTTPClientTransport(new URL(config.url), {})
-			} else {
-				transport = new StdioClientTransport({
-					command: config.command,
-					args: config.args,
-					env: {
-						...config.env,
-						...(process.env.PATH ? { PATH: process.env.PATH } : {}),
-						// ...(process.env.NODE_PATH ? { NODE_PATH: process.env.NODE_PATH } : {}),
-					},
-					stderr: "pipe", // necessary for stderr to be available
-				})
-			}
+			switch (config.type) {
+				case "stdio": {
+					transport = new StdioClientTransport({
+						command: config.command,
+						args: config.args,
+						cwd: config.cwd,
+						env: {
+							// ...(config.env ? await injectEnv(config.env) : {}), // Commented out as injectEnv is not found
+							...(config.env || {}), // Use config.env directly or an empty object
+							...(process.env.PATH ? { PATH: process.env.PATH } : {}),
+						},
+						stderr: "pipe",
+					})
 
-			transport.onerror = async (error) => {
-				console.error(`Transport error for "${name}":`, error)
-				const connection = this.connections.find((conn) => conn.server.name === name)
-				if (connection) {
-					connection.server.status = "disconnected"
-					this.appendErrorMessage(connection, error.message)
-				}
-				await this.notifyWebviewOfServerChanges()
-			}
+					transport.onerror = async (error) => {
+						console.error(`Transport error for "${name}":`, error)
+						const connection = this.findConnection(name, source)
+						if (connection) {
+							connection.server.status = "disconnected"
+							this.appendErrorMessage(connection, error instanceof Error ? error.message : `${error}`)
+						}
+						await this.notifyWebviewOfServerChanges()
+					}
 
-			transport.onclose = async () => {
-				const connection = this.connections.find((conn) => conn.server.name === name)
-				if (connection) {
-					connection.server.status = "disconnected"
+					transport.onclose = async () => {
+						const connection = this.findConnection(name, source)
+						if (connection) {
+							connection.server.status = "disconnected"
+						}
+						await this.notifyWebviewOfServerChanges()
+					}
+
+					await transport.start()
+					const stderrStream = transport.stderr
+					if (stderrStream) {
+						stderrStream.on("data", async (data: Buffer) => {
+							const output = data.toString()
+							const isInfoLog = /INFO/i.test(output)
+
+							if (isInfoLog) {
+								console.log(`Server "${name}" info:`, output)
+							} else {
+								console.error(`Server "${name}" stderr:`, output)
+								const connection = this.findConnection(name, source)
+								if (connection) {
+									this.appendErrorMessage(connection, output)
+									if (connection.server.status === "disconnected") {
+										await this.notifyWebviewOfServerChanges()
+									}
+								}
+							}
+						})
+					} else {
+						console.error(`No stderr stream for ${name}`)
+					}
+					transport.start = async () => {}
+					break
 				}
-				await this.notifyWebviewOfServerChanges()
+				case "sse": {
+					const sseOptions = {
+						requestInit: {
+							headers: config.headers,
+						},
+					}
+					const reconnectingEventSourceOptions = {
+						max_retry_time: 5000,
+						withCredentials: config.headers?.["Authorization"] ? true : false,
+					}
+					global.EventSource = ReconnectingEventSource
+					transport = new SSEClientTransport(new URL(config.url), {
+						...sseOptions,
+						eventSourceInit: reconnectingEventSourceOptions,
+					})
+
+					transport.onerror = async (error) => {
+						console.error(`Transport error for "${name}":`, error)
+						const connection = this.findConnection(name, source)
+						if (connection) {
+							connection.server.status = "disconnected"
+							this.appendErrorMessage(connection, error instanceof Error ? error.message : `${error}`)
+						}
+						await this.notifyWebviewOfServerChanges()
+					}
+					break
+				}
+				case "streamableHttp": {
+					transport = new StreamableHTTPClientTransport(new URL(config.url), {
+						requestInit: {
+							headers: config.headers,
+						},
+					})
+					transport.onerror = async (error) => {
+						console.error(`Transport error for "${name}":`, error)
+						const connection = this.findConnection(name, source)
+						if (connection) {
+							connection.server.status = "disconnected"
+							this.appendErrorMessage(connection, error instanceof Error ? error.message : `${error}`)
+						}
+						await this.notifyWebviewOfServerChanges()
+					}
+					break
+				}
+				default:
+					throw new Error(`Unknown transport type: ${(config as any).type}`)
 			}
 
 			const connection: McpConnection = {
@@ -374,39 +392,6 @@ export class McpHub {
 			}
 			this.connections.push(connection)
 
-			if (config.transportType === "stdio") {
-				// transport.stderr is only available after the process has been started. However we can't start it separately from the .connect() call because it also starts the transport. And we can't place this after the connect call since we need to capture the stderr stream before the connection is established, in order to capture errors during the connection process.
-				// As a workaround, we start the transport ourselves, and then monkey-patch the start method to no-op so that .connect() doesn't try to start it again.
-				await transport.start()
-				const stderrStream = (transport as StdioClientTransport).stderr
-				if (stderrStream) {
-					stderrStream.on("data", async (data: Buffer) => {
-						const output = data.toString()
-						// Check if output contains INFO level log
-						const isInfoLog = !/\berror\b/i.test(output)
-
-						if (isInfoLog) {
-							// Log normal informational messages
-							console.info(`Server "${name}" info:`, output)
-						} else {
-							// Treat as error log
-							console.error(`Server "${name}" stderr:`, output)
-							const connection = this.connections.find((conn) => conn.server.name === name)
-							if (connection) {
-								this.appendErrorMessage(connection, output)
-								// Only notify webview if server is already disconnected
-								if (connection.server.status === "disconnected") {
-									await this.notifyWebviewOfServerChanges()
-								}
-							}
-						}
-					})
-				} else {
-					console.error(`No stderr stream for ${name}`)
-				}
-				transport.start = async () => {} // No-op now, .connect() won't fail
-			}
-
 			// Connect
 			await client.connect(transport)
 
@@ -419,7 +404,7 @@ export class McpHub {
 			connection.server.resourceTemplates = await this.fetchResourceTemplatesList(name)
 		} catch (error) {
 			// Update status with error
-			const connection = this.connections.find((conn) => conn.server.name === name)
+			const connection = this.findConnection(name, source)
 			if (connection) {
 				connection.server.status = "disconnected"
 				this.appendErrorMessage(connection, error instanceof Error ? error.message : String(error))
@@ -525,21 +510,21 @@ export class McpHub {
 			if (!currentConnection) {
 				// New server
 				try {
-					if (config.transportType === "stdio") {
+					if (config.type === "stdio") {
 						this.setupFileWatcher(name, config)
 					}
-					await this.connectToServer(name, config)
+					await this.connectToServer(name, config, "rpc")
 				} catch (error) {
 					console.error(`Failed to connect to new MCP server ${name}:`, error)
 				}
 			} else if (!deepEqual(JSON.parse(currentConnection.server.config), config)) {
 				// Existing server with changed config
 				try {
-					if (config.transportType === "stdio") {
+					if (config.type === "stdio") {
 						this.setupFileWatcher(name, config)
 					}
 					await this.deleteConnection(name)
-					await this.connectToServer(name, config)
+					await this.connectToServer(name, config, "rpc")
 					console.log(`Reconnected MCP server with updated config: ${name}`)
 				} catch (error) {
 					console.error(`Failed to reconnect MCP server ${name}:`, error)
@@ -572,21 +557,21 @@ export class McpHub {
 			if (!currentConnection) {
 				// New server
 				try {
-					if (config.transportType === "stdio") {
+					if (config.type === "stdio") {
 						this.setupFileWatcher(name, config)
 					}
-					await this.connectToServer(name, config)
+					await this.connectToServer(name, config, "internal")
 				} catch (error) {
 					console.error(`Failed to connect to new MCP server ${name}:`, error)
 				}
 			} else if (!deepEqual(JSON.parse(currentConnection.server.config), config)) {
 				// Existing server with changed config
 				try {
-					if (config.transportType === "stdio") {
+					if (config.type === "stdio") {
 						this.setupFileWatcher(name, config)
 					}
 					await this.deleteConnection(name)
-					await this.connectToServer(name, config)
+					await this.connectToServer(name, config, "internal")
 					console.log(`Reconnected MCP server with updated config: ${name}`)
 				} catch (error) {
 					console.error(`Failed to reconnect MCP server ${name}:`, error)
@@ -598,7 +583,7 @@ export class McpHub {
 		this.isConnecting = false
 	}
 
-	private setupFileWatcher(name: string, config: Extract<McpServerConfig, { transportType: "stdio" }>) {
+	private setupFileWatcher(name: string, config: Extract<McpServerConfig, { type: "stdio" }>) {
 		const filePath = config.args?.find((arg: string) => arg.includes("build/index.js"))
 		if (filePath) {
 			// we use chokidar instead of onDidSaveTextDocument because it doesn't require the file to be open in the editor. The settings config is better suited for onDidSave since that will be manually updated by the user or Cline (and we want to detect save events, not every file change)
@@ -635,7 +620,7 @@ export class McpHub {
 			try {
 				await this.deleteConnection(serverName)
 				// Try to connect again using existing config
-				await this.connectToServerRPC(serverName, JSON.parse(inMemoryConfig))
+				await this.connectToServer(serverName, JSON.parse(inMemoryConfig), "rpc")
 			} catch (error) {
 				console.error(`Failed to restart connection for ${serverName}:`, error)
 			}
@@ -667,7 +652,7 @@ export class McpHub {
 			try {
 				await this.deleteConnection(serverName)
 				// Try to connect again using existing config
-				await this.connectToServer(serverName, JSON.parse(config))
+				await this.connectToServer(serverName, JSON.parse(config), "internal")
 				vscode.window.showInformationMessage(`${serverName} MCP server connected`)
 			} catch (error) {
 				console.error(`Failed to restart connection for ${serverName}:`, error)