feat: support streamable http server connection (pathintegral-institute#150)

* feat: support streamable http server connection

* test: fix test

* Update src/mcpm/router/client_connection.py

Author: JoJoJoJoJoJoJo

pyproject.toml

@@ -24,7 +24,7 @@ dependencies = [
     "rich>=12.0.0",
     "requests>=2.28.0",
     "pydantic>=2.5.1",
-    "mcp==1.6.0",
+    "mcp>=1.8.0",
     "ruamel-yaml>=0.18.10",
     "watchfiles>=1.0.4",
     "duckdb>=1.2.2",

src/mcpm/router/client_connection.py

@@ -2,8 +2,10 @@
 import logging
 from typing import Optional, TextIO, cast
 
+import requests
 from mcp import ClientSession, InitializeResult, StdioServerParameters, stdio_client
 from mcp.client.sse import sse_client
+from mcp.client.streamable_http import streamablehttp_client
 
 from mcpm.core.schema import ServerConfig, SSEServerConfig, STDIOServerConfig
 
@@ -21,6 +23,11 @@ def _sse_transport_context(server_config: ServerConfig):
     return sse_client(server_config.url, headers=server_config.headers)
 
 
+def _streamable_http_transport_context(server_config: ServerConfig):
+    server_config = cast(SSEServerConfig, server_config)
+    return streamablehttp_client(server_config.url, headers=server_config.headers)
+
+
 class ServerConnection:
     def __init__(self, server_config: ServerConfig, errlog: TextIO) -> None:
         self.session: Optional[ClientSession] = None
@@ -31,14 +38,21 @@ def __init__(self, server_config: ServerConfig, errlog: TextIO) -> None:
         self._shutdown_event = asyncio.Event()
         self._errlog = errlog
 
-        self._transport_context_factory = (
-            lambda config: _stdio_transport_context(config, errlog=self._errlog)
-            if isinstance(config, STDIOServerConfig)
-            else _sse_transport_context(config)
-        )
-
         self._server_task = asyncio.create_task(self._server_lifespan_cycle())
 
+    def _transport_context_factory(self, server_config: ServerConfig):
+        if isinstance(server_config, STDIOServerConfig):
+            return _stdio_transport_context(server_config, self._errlog)
+        elif isinstance(server_config, SSEServerConfig):
+            r = requests.head(server_config.url)
+            if r.status_code != 200:
+                return _streamable_http_transport_context(server_config)
+            if r.headers.get("connection") == "keep-alive" and r.headers.get("content-type", "").startswith(
+                "text/event-stream"
+            ):
+                return _sse_transport_context(server_config)
+            return _streamable_http_transport_context(server_config)
+
     def healthy(self) -> bool:
         return self.session is not None and self._initialized
 
@@ -56,7 +70,7 @@ async def wait_for_shutdown_request(self):
 
     async def _server_lifespan_cycle(self):
         try:
-            async with self._transport_context_factory(self.server_config) as (read, write):
+            async with self._transport_context_factory(self.server_config) as (read, write, *_):
                 async with ClientSession(read, write) as session:
                     self.session_initialized_response = await session.initialize()
 

src/mcpm/router/transport.py

@@ -10,6 +10,7 @@
 from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
 from mcp import types
 from mcp.server.sse import SseServerTransport
+from mcp.shared.message import SessionMessage
 from pydantic import ValidationError
 from sse_starlette import EventSourceResponse
 from starlette.background import BackgroundTask
@@ -86,11 +87,11 @@ async def connect_sse(self, scope: Scope, receive: Receive, send: Send):
             return
 
         logger.debug("Setting up SSE connection")
-        read_stream: MemoryObjectReceiveStream[types.JSONRPCMessage | Exception]
-        read_stream_writer: MemoryObjectSendStream[types.JSONRPCMessage | Exception]
+        read_stream: MemoryObjectReceiveStream[SessionMessage | Exception]
+        read_stream_writer: MemoryObjectSendStream[SessionMessage | Exception]
 
-        write_stream: MemoryObjectSendStream[types.JSONRPCMessage]
-        write_stream_reader: MemoryObjectReceiveStream[types.JSONRPCMessage]
+        write_stream: MemoryObjectSendStream[SessionMessage]
+        write_stream_reader: MemoryObjectReceiveStream[SessionMessage]
 
         read_stream_writer, read_stream = anyio.create_memory_object_stream(0)
         write_stream, write_stream_reader = anyio.create_memory_object_stream(0)
@@ -118,12 +119,12 @@ async def sse_writer():
                 await sse_stream_writer.send({"event": "endpoint", "data": session_uri})
                 logger.debug(f"Sent endpoint event: {session_uri}")
 
-                async for message in write_stream_reader:
-                    logger.debug(f"Sending message via SSE: {message}")
+                async for session_message in write_stream_reader:
+                    logger.debug(f"Sending message via SSE: {session_message}")
                     await sse_stream_writer.send(
                         {
                             "event": "message",
-                            "data": message.model_dump_json(by_alias=True, exclude_none=True),
+                            "data": session_message.message.model_dump_json(by_alias=True, exclude_none=True),
                         }
                     )
 
@@ -228,7 +229,7 @@ async def handle_post_message(self, scope: Scope, receive: Receive, send: Send):
 
         # add error handling, catch possible pipe errors
         try:
-            await writer.send(message)
+            await writer.send(SessionMessage(message=message))
         except (BrokenPipeError, ConnectionError, OSError) as e:
             # if it's EPIPE error or other connection error, log it but don't throw an exception
             if isinstance(e, OSError) and e.errno == 32:  # EPIPE

src/mcpm/utils/errlog_manager.py

@@ -6,6 +6,7 @@
 
 DEFAULT_ROOT_STDERR_LOG_DIR = get_log_directory("mcpm") / "errlogs"
 
+
 class ServerErrorLogManager:
     """
     A manager for server error logs.
@@ -30,5 +31,5 @@ def close_errlog_file(self, server_id: str) -> None:
             del self._log_files[server_id]
 
     def close_all(self) -> None:
-        for server_id in self._log_files:
+        for server_id in list(self._log_files.keys()):
             self.close_errlog_file(server_id)

tests/test_add.py

@@ -96,7 +96,7 @@ def test_add_server_with_missing_arg(windsurf_manager, monkeypatch):
         patch("rich.progress.Progress.add_task"),
     ):
         # Use CliRunner which provides its own isolated environment
-        runner = CliRunner(mix_stderr=False)
+        runner = CliRunner()
         result = runner.invoke(add, ["server-test", "--force", "--alias", "test-missing-arg"])
 
         if result.exit_code != 0:
@@ -164,7 +164,7 @@ def test_add_server_with_empty_args(windsurf_manager, monkeypatch):
         patch("rich.progress.Progress.stop"),
         patch("rich.progress.Progress.add_task"),
     ):
-        runner = CliRunner(mix_stderr=False)
+        runner = CliRunner()
         result = runner.invoke(add, ["server-test", "--force", "--alias", "test-empty-args"])
 
         assert result.exit_code == 0