feat(mcp): enhance MCP transport support and add server management UI
- Extend McpTransport to support "sse" in schemas. - Refactor McpToolClient to handle both "streamable_http" and "sse" transports. - Introduce McpServerDialog for managing MCP server configurations, including transport settings and tool synchronization. - Replace McpServersSection with the new dialog component for improved server management. - Add tests for MCP transport handling and server dialog functionality.
This commit is contained in:
@@ -22,7 +22,7 @@ KnowledgeRetrievalMode = Literal["automatic", "on_demand"]
|
||||
ToolType = Literal["end_call", "http", "mcp"]
|
||||
ToolStatus = Literal["active", "archived", "draft"]
|
||||
McpServerStatus = Literal["active", "archived", "draft"]
|
||||
McpTransport = Literal["streamable_http"]
|
||||
McpTransport = Literal["streamable_http", "sse"]
|
||||
ToolParameterType = Literal["string", "number", "integer", "boolean", "object", "array"]
|
||||
ToolParameterLocation = Literal["path", "query", "body", "header"]
|
||||
DynamicVariableType = Literal["string", "number", "boolean"]
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
"""One-shot Streamable HTTP MCP discovery and execution.
|
||||
"""One-shot HTTP MCP discovery and execution.
|
||||
|
||||
Connections intentionally live for one operation in the MVP. This keeps the
|
||||
async context ownership simple and correct; a session pool can be introduced
|
||||
@@ -15,6 +15,7 @@ from typing import Any, AsyncIterator
|
||||
|
||||
import httpx
|
||||
from mcp import ClientSession
|
||||
from mcp.client.sse import sse_client
|
||||
from mcp.client.streamable_http import streamable_http_client
|
||||
|
||||
from models import RuntimeMcpServer
|
||||
@@ -41,16 +42,14 @@ class McpToolClient:
|
||||
"""Minimal public MCP SDK wrapper shared by discovery and runtime calls."""
|
||||
|
||||
@asynccontextmanager
|
||||
async def _session(
|
||||
async def _transport_streams(
|
||||
self,
|
||||
server: RuntimeMcpServer,
|
||||
) -> AsyncIterator[ClientSession]:
|
||||
if server.transport != "streamable_http":
|
||||
raise McpClientError(f"暂不支持 MCP transport: {server.transport}")
|
||||
|
||||
headers = {**server.headers, **server.secret_headers}
|
||||
timeout = max(1, min(int(server.timeout_seconds), 120))
|
||||
try:
|
||||
headers: dict[str, str],
|
||||
timeout: int,
|
||||
) -> AsyncIterator[tuple[Any, Any]]:
|
||||
"""Open the configured MCP transport and expose its read/write streams."""
|
||||
if server.transport == "streamable_http":
|
||||
async with httpx.AsyncClient(
|
||||
headers=headers,
|
||||
timeout=httpx.Timeout(timeout),
|
||||
@@ -61,13 +60,38 @@ class McpToolClient:
|
||||
http_client=http_client,
|
||||
) as streams:
|
||||
read_stream, write_stream, _ = streams
|
||||
async with ClientSession(
|
||||
read_stream,
|
||||
write_stream,
|
||||
read_timeout_seconds=timedelta(seconds=timeout),
|
||||
) as session:
|
||||
await session.initialize()
|
||||
yield session
|
||||
yield read_stream, write_stream
|
||||
return
|
||||
|
||||
if server.transport == "sse":
|
||||
async with sse_client(
|
||||
server.url,
|
||||
headers=headers,
|
||||
timeout=timeout,
|
||||
sse_read_timeout=timeout,
|
||||
) as streams:
|
||||
yield streams
|
||||
return
|
||||
|
||||
raise McpClientError(f"暂不支持 MCP transport: {server.transport}")
|
||||
|
||||
@asynccontextmanager
|
||||
async def _session(
|
||||
self,
|
||||
server: RuntimeMcpServer,
|
||||
) -> AsyncIterator[ClientSession]:
|
||||
headers = {**server.headers, **server.secret_headers}
|
||||
timeout = max(1, min(int(server.timeout_seconds), 120))
|
||||
try:
|
||||
async with self._transport_streams(server, headers, timeout) as streams:
|
||||
read_stream, write_stream = streams
|
||||
async with ClientSession(
|
||||
read_stream,
|
||||
write_stream,
|
||||
read_timeout_seconds=timedelta(seconds=timeout),
|
||||
) as session:
|
||||
await session.initialize()
|
||||
yield session
|
||||
except McpClientError:
|
||||
raise
|
||||
except httpx.TimeoutException as exc:
|
||||
|
||||
51
backend/tests/test_mcp_client.py
Normal file
51
backend/tests/test_mcp_client.py
Normal file
@@ -0,0 +1,51 @@
|
||||
import unittest
|
||||
from contextlib import asynccontextmanager
|
||||
from unittest.mock import patch
|
||||
|
||||
from models import RuntimeMcpServer
|
||||
from schemas import McpServerUpsert
|
||||
from services.tools.mcp_client import McpToolClient
|
||||
|
||||
|
||||
class McpTransportTests(unittest.IsolatedAsyncioTestCase):
|
||||
def test_server_schema_accepts_sse(self):
|
||||
server = McpServerUpsert(
|
||||
name="旧版 MCP 服务",
|
||||
transport="sse",
|
||||
url="https://mcp.example.com/sse",
|
||||
)
|
||||
|
||||
self.assertEqual(server.transport, "sse")
|
||||
|
||||
async def test_sse_transport_uses_sdk_sse_client(self):
|
||||
captured: dict = {}
|
||||
|
||||
@asynccontextmanager
|
||||
async def fake_sse_client(url, **kwargs):
|
||||
captured["url"] = url
|
||||
captured.update(kwargs)
|
||||
yield "read-stream", "write-stream"
|
||||
|
||||
server = RuntimeMcpServer(
|
||||
id="mcp_sse",
|
||||
transport="sse",
|
||||
url="https://mcp.example.com/sse",
|
||||
)
|
||||
client = McpToolClient()
|
||||
|
||||
with patch("services.tools.mcp_client.sse_client", fake_sse_client):
|
||||
async with client._transport_streams(
|
||||
server,
|
||||
{"Authorization": "Bearer test"},
|
||||
20,
|
||||
) as streams:
|
||||
self.assertEqual(streams, ("read-stream", "write-stream"))
|
||||
|
||||
self.assertEqual(captured["url"], server.url)
|
||||
self.assertEqual(captured["headers"], {"Authorization": "Bearer test"})
|
||||
self.assertEqual(captured["timeout"], 20)
|
||||
self.assertEqual(captured["sse_read_timeout"], 20)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user