From e178483971a7c98c6c1b782df973a5d0d25adf9a Mon Sep 17 00:00:00 2001
From: suntp <605682931@qq.com>
Date: Tue, 29 Jul 2025 14:04:04 +0800
Subject: [PATCH] fix: Add streamable MCP server support (#468)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* fix: Add streamable MCP server support(#349)

* “Revert-timeout”

* fix lint and test check

* modify streamable error notify

---------

Co-authored-by: Willem Jiang <willem.jiang@gmail.com>
---
 pyproject.toml                                      |  2 +-
 src/server/mcp_utils.py                             | 13 ++++++++++++-
 tests/unit/server/test_mcp_utils.py                 |  7 ++++++-
 .../app/settings/dialogs/add-mcp-server-dialog.tsx  |  2 +-
 web/src/core/mcp/schema.ts                          |  5 +++++
 web/src/core/mcp/types.ts                           |  4 ++--
 6 files changed, 27 insertions(+), 6 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 7d83ae6..52e67f3 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -30,7 +30,7 @@ dependencies = [
     "duckduckgo-search>=8.0.0",
     "inquirerpy>=0.3.4",
     "arxiv>=2.2.0",
-    "mcp>=1.6.0",
+    "mcp>=1.11.0",
     "langchain-mcp-adapters>=0.0.9",
     "langchain-deepseek>=0.1.3",
     "wikipedia>=1.4.0",
diff --git a/src/server/mcp_utils.py b/src/server/mcp_utils.py
index 3b69c3a..d485688 100644
--- a/src/server/mcp_utils.py
+++ b/src/server/mcp_utils.py
@@ -9,6 +9,7 @@ from fastapi import HTTPException
 from mcp import ClientSession, StdioServerParameters
 from mcp.client.stdio import stdio_client
 from mcp.client.sse import sse_client
+from mcp.client.streamable_http import streamablehttp_client
 
 logger = logging.getLogger(__name__)
 
@@ -29,7 +30,7 @@ async def _get_tools_from_client_session(
     Raises:
         Exception: If there's an error during the process
     """
-    async with client_context_manager as (read, write):
+    async with client_context_manager as (read, write, _):
         async with ClientSession(
             read, write, read_timeout_seconds=timedelta(seconds=timeout_seconds)
         ) as session:
@@ -92,6 +93,16 @@ async def load_mcp_tools(
                 sse_client(url=url), timeout_seconds
             )
 
+        elif server_type == "streamable_http":
+            if not url:
+                raise HTTPException(
+                    status_code=400, detail="URL is required for streamable_http type"
+                )
+
+            return await _get_tools_from_client_session(
+                streamablehttp_client(url=url), timeout_seconds
+            )
+
         else:
             raise HTTPException(
                 status_code=400, detail=f"Unsupported server type: {server_type}"
diff --git a/tests/unit/server/test_mcp_utils.py b/tests/unit/server/test_mcp_utils.py
index 3975b70..cb130b1 100644
--- a/tests/unit/server/test_mcp_utils.py
+++ b/tests/unit/server/test_mcp_utils.py
@@ -13,8 +13,13 @@ import src.server.mcp_utils as mcp_utils
 async def test__get_tools_from_client_session_success(mock_ClientSession):
     mock_read = AsyncMock()
     mock_write = AsyncMock()
+    mock_callback = AsyncMock()
     mock_context_manager = AsyncMock()
-    mock_context_manager.__aenter__.return_value = (mock_read, mock_write)
+    mock_context_manager.__aenter__.return_value = (
+        mock_read,
+        mock_write,
+        mock_callback,
+    )
     mock_context_manager.__aexit__.return_value = None
 
     mock_session = AsyncMock()
diff --git a/web/src/app/settings/dialogs/add-mcp-server-dialog.tsx b/web/src/app/settings/dialogs/add-mcp-server-dialog.tsx
index 95295cd..92f1a10 100644
--- a/web/src/app/settings/dialogs/add-mcp-server-dialog.tsx
+++ b/web/src/app/settings/dialogs/add-mcp-server-dialog.tsx
@@ -96,7 +96,7 @@ export function AddMCPServerDialog({
         addingServers.push(metadata);
       } else if ("url" in server) {
         const metadata: SimpleSSEMCPServerMetadata = {
-          transport: "sse",
+          transport: server.transport,
           name: key,
           url: server.url,
         };
diff --git a/web/src/core/mcp/schema.ts b/web/src/core/mcp/schema.ts
index 0e4f879..fc472cb 100644
--- a/web/src/core/mcp/schema.ts
+++ b/web/src/core/mcp/schema.ts
@@ -47,6 +47,11 @@ export const MCPConfigSchema = z.object({
               message: "`env` must be an object of key-value pairs",
             })
             .optional(),
+          transport: z
+            .enum(["sse", "streamable_http"], {
+              message: "transport must be either sse or streamable_http"
+            })
+            .default("sse"),
         }),
       ],
       {
diff --git a/web/src/core/mcp/types.ts b/web/src/core/mcp/types.ts
index 854e21b..0c1d73b 100644
--- a/web/src/core/mcp/types.ts
+++ b/web/src/core/mcp/types.ts
@@ -28,8 +28,8 @@ export type SimpleStdioMCPServerMetadata = Omit<
   "enabled" | "tools" | "createdAt" | "updatedAt"
 >;
 
-export interface SSEMCPServerMetadata extends GenericMCPServerMetadata<"sse"> {
-  transport: "sse";
+export interface SSEMCPServerMetadata extends GenericMCPServerMetadata<string> {
+  transport: string;
   url: string;
 }
 export type SimpleSSEMCPServerMetadata = Omit<