From 270d8c3712aa7933fb0b41d88bb7dbf5994a344e Mon Sep 17 00:00:00 2001 From: Anoyer-lzh <1316895941@qq.com> Date: Wed, 20 Aug 2025 17:23:57 +0800 Subject: [PATCH] fix: env parameters exception when configuring SSE or HTTP MCP server (#513) * fix: _create_streamable_http_session() got an unexpected keyword argument 'env' fix unit error * update md --------- Co-authored-by: Willem Jiang --- docs/mcp_integrations.md | 13 ++++++++++++- src/graph/nodes.py | 2 +- src/server/app.py | 2 ++ src/server/mcp_request.py | 24 ++++++++++++++++++++---- src/server/mcp_utils.py | 12 +++++++----- tests/unit/server/test_mcp_utils.py | 7 ++++++- web/src/core/mcp/types.ts | 6 ++++-- web/src/core/store/settings-store.ts | 4 ++-- 8 files changed, 54 insertions(+), 16 deletions(-) diff --git a/docs/mcp_integrations.md b/docs/mcp_integrations.md index 9ccf6fe..f5348b4 100644 --- a/docs/mcp_integrations.md +++ b/docs/mcp_integrations.md @@ -46,7 +46,18 @@ For `sse` type: { "transport": "sse", "url": "http://localhost:3000/sse", - "env": { + "headers": { + "API_KEY": "value" + } +} +``` + +For `streamable_http` type: +```json +{ + "transport": "streamable_http", + "url": "http://localhost:3000/mcp", + "headers": { "API_KEY": "value" } } diff --git a/src/graph/nodes.py b/src/graph/nodes.py index 2062eac..e006ca2 100644 --- a/src/graph/nodes.py +++ b/src/graph/nodes.py @@ -450,7 +450,7 @@ async def _setup_and_execute_agent_step( mcp_servers[server_name] = { k: v for k, v in server_config.items() - if k in ("transport", "command", "args", "url", "env") + if k in ("transport", "command", "args", "url", "env", "headers") } for tool_name in server_config["enabled_tools"]: enabled_tools[tool_name] = server_name diff --git a/src/server/app.py b/src/server/app.py index efea5b8..b0e03a2 100644 --- a/src/server/app.py +++ b/src/server/app.py @@ -552,6 +552,7 @@ async def mcp_server_metadata(request: MCPServerMetadataRequest): args=request.args, url=request.url, env=request.env, + headers=request.headers, timeout_seconds=timeout, ) @@ -562,6 +563,7 @@ async def mcp_server_metadata(request: MCPServerMetadataRequest): args=request.args, url=request.url, env=request.env, + headers=request.headers, tools=tools, ) diff --git a/src/server/mcp_request.py b/src/server/mcp_request.py index 85a490b..617bc8a 100644 --- a/src/server/mcp_request.py +++ b/src/server/mcp_request.py @@ -10,7 +10,10 @@ class MCPServerMetadataRequest(BaseModel): """Request model for MCP server metadata.""" transport: str = Field( - ..., description="The type of MCP server connection (stdio or sse)" + ..., + description=( + "The type of MCP server connection (stdio or sse or streamable_http)" + ), ) command: Optional[str] = Field( None, description="The command to execute (for stdio type)" @@ -21,7 +24,12 @@ class MCPServerMetadataRequest(BaseModel): url: Optional[str] = Field( None, description="The URL of the SSE server (for sse type)" ) - env: Optional[Dict[str, str]] = Field(None, description="Environment variables") + env: Optional[Dict[str, str]] = Field( + None, description="Environment variables (for stdio type)" + ) + headers: Optional[Dict[str, str]] = Field( + None, description="HTTP headers (for sse/streamable_http type)" + ) timeout_seconds: Optional[int] = Field( None, description="Optional custom timeout in seconds for the operation" ) @@ -31,7 +39,10 @@ class MCPServerMetadataResponse(BaseModel): """Response model for MCP server metadata.""" transport: str = Field( - ..., description="The type of MCP server connection (stdio or sse)" + ..., + description=( + "The type of MCP server connection (stdio or sse or streamable_http)" + ), ) command: Optional[str] = Field( None, description="The command to execute (for stdio type)" @@ -42,7 +53,12 @@ class MCPServerMetadataResponse(BaseModel): url: Optional[str] = Field( None, description="The URL of the SSE server (for sse type)" ) - env: Optional[Dict[str, str]] = Field(None, description="Environment variables") + env: Optional[Dict[str, str]] = Field( + None, description="Environment variables (for stdio type)" + ) + headers: Optional[Dict[str, str]] = Field( + None, description="HTTP headers (for sse/streamable_http type)" + ) tools: List = Field( default_factory=list, description="Available tools from the MCP server" ) diff --git a/src/server/mcp_utils.py b/src/server/mcp_utils.py index c18c6b5..dcb9340 100644 --- a/src/server/mcp_utils.py +++ b/src/server/mcp_utils.py @@ -47,17 +47,19 @@ async def load_mcp_tools( args: Optional[List[str]] = None, url: Optional[str] = None, env: Optional[Dict[str, str]] = None, + headers: Optional[Dict[str, str]] = None, timeout_seconds: int = 60, # Longer default timeout for first-time executions ) -> List: """ Load tools from an MCP server. Args: - server_type: The type of MCP server connection (stdio or sse) + server_type: The type of MCP server connection (stdio, sse, or streamable_http) command: The command to execute (for stdio type) args: Command arguments (for stdio type) - url: The URL of the SSE server (for sse type) - env: Environment variables + url: The URL of the SSE/HTTP server (for sse/streamable_http type) + env: Environment variables (for stdio type) + headers: HTTP headers (for sse/streamable_http type) timeout_seconds: Timeout in seconds (default: 60 for first-time executions) Returns: @@ -90,7 +92,7 @@ async def load_mcp_tools( ) return await _get_tools_from_client_session( - sse_client(url=url), timeout_seconds + sse_client(url=url, headers=headers, timeout=timeout_seconds), timeout_seconds ) elif server_type == "streamable_http": @@ -100,7 +102,7 @@ async def load_mcp_tools( ) return await _get_tools_from_client_session( - streamablehttp_client(url=url), timeout_seconds + streamablehttp_client(url=url, headers=headers, timeout=timeout_seconds), timeout_seconds, ) else: diff --git a/tests/unit/server/test_mcp_utils.py b/tests/unit/server/test_mcp_utils.py index c37ed8d..d9a5a2b 100644 --- a/tests/unit/server/test_mcp_utils.py +++ b/tests/unit/server/test_mcp_utils.py @@ -87,10 +87,15 @@ async def test_load_mcp_tools_sse_success(mock_sse_client, mock_get_tools): result = await mcp_utils.load_mcp_tools( server_type="sse", url="http://localhost:1234", + headers={"Authorization": "Bearer 1234567890"}, timeout_seconds=7, ) assert result == ["toolB"] - mock_sse_client.assert_called_once_with(url="http://localhost:1234") + mock_sse_client.assert_called_once_with( + url="http://localhost:1234", + headers={"Authorization": "Bearer 1234567890"}, + timeout=7, + ) mock_get_tools.assert_awaited_once_with(mock_client, 7) diff --git a/web/src/core/mcp/types.ts b/web/src/core/mcp/types.ts index 0cd8806..352c61e 100644 --- a/web/src/core/mcp/types.ts +++ b/web/src/core/mcp/types.ts @@ -12,6 +12,7 @@ export interface GenericMCPServerMetadata { transport: T; enabled: boolean; env?: Record; + headers?: Record; tools: MCPToolMetadata[]; createdAt: number; updatedAt: number; @@ -28,8 +29,9 @@ export type SimpleStdioMCPServerMetadata = Omit< "enabled" | "tools" | "createdAt" | "updatedAt" >; -export interface SSEMCPServerMetadata extends GenericMCPServerMetadata<"sse"|"streamable_http"> { - transport: "sse"|"streamable_http" +export interface SSEMCPServerMetadata + extends GenericMCPServerMetadata<"sse" | "streamable_http"> { + transport: "sse" | "streamable_http"; url: string; } diff --git a/web/src/core/store/settings-store.ts b/web/src/core/store/settings-store.ts index d3d2fe8..717bcde 100644 --- a/web/src/core/store/settings-store.ts +++ b/web/src/core/store/settings-store.ts @@ -94,7 +94,7 @@ export const getChatStreamSettings = () => { if (mcpServers.length > 0) { mcpSettings = { servers: mcpServers.reduce((acc, cur) => { - const { transport, env } = cur; + const { transport, env, headers } = cur; let server: SimpleMCPServerMetadata; if (transport === "stdio") { server = { @@ -108,7 +108,7 @@ export const getChatStreamSettings = () => { server = { name: cur.name, transport, - env, + headers, url: cur.url, }; }