Skip to content

Commit 501340a

Browse files
committed
refactor(mcp): migrate from MCP SDK v1 to v2.0.0b1
Breaking changes applied: - streamablehttp_client → streamable_http_client (renamed import) - McpError → MCPError (exception renamed) - mcp.types → mcp_types (split into separate package) - FastMCP → MCPServer (base class renamed) - SSLCapableFastMCP → SSLCapableMCPServer - MCPServer.__init__ no longer accepts host/port/transport_security; these are passed to run methods and streamable_http_app() instead - streamable_http_client returns 2-tuple (read, write) instead of 3-tuple; session ID tracking and manual termination removed (terminate_on_close=True) - Remove _http_client_factory; use pre-built httpx.AsyncClient directly - Replace self.settings.host/port with self.server_config.host/port Deprecations removed: - __terminate_http_session(), _session_id, _get_session_id, _http_client_factory Tests: 190 passing
1 parent 09b234b commit 501340a

10 files changed

Lines changed: 686 additions & 237 deletions

File tree

‎MIGRATION_MCP_V2.md‎

Lines changed: 524 additions & 0 deletions
Large diffs are not rendered by default.

‎cpex/framework/external/mcp/client.py‎

Lines changed: 9 additions & 41 deletions
Original file line numberDiff line numberDiff line change
@@ -21,10 +21,10 @@
2121
# Third-Party
2222
import httpx
2323
import orjson
24-
from mcp import ClientSession, McpError, StdioServerParameters
24+
from mcp import ClientSession, MCPError, StdioServerParameters
2525
from mcp.client.stdio import stdio_client
26-
from mcp.client.streamable_http import streamablehttp_client
27-
from mcp.types import TextContent
26+
from mcp.client.streamable_http import streamable_http_client
27+
from mcp_types import TextContent
2828

2929
# First-Party
3030
from cpex.framework.base import HookRef, Plugin, PluginRef
@@ -82,9 +82,6 @@ def __init__(self, config: PluginConfig) -> None:
8282
self._stdio_ready: Optional[asyncio.Event] = None
8383
self._stdio_stop: Optional[asyncio.Event] = None
8484
self._stdio_error: Optional[BaseException] = None
85-
self._get_session_id: Optional[Callable[[], str | None]] = None
86-
self._session_id: Optional[str] = None
87-
self._http_client_factory: Optional[Callable[..., httpx.AsyncClient]] = None
8885
self._reconnect_attempts: int = 3
8986
self._reconnect_delay: float = 0.1
9087
self._reconnect_lock: asyncio.Lock = asyncio.Lock()
@@ -373,23 +370,20 @@ def _tls_httpx_client_factory(
373370

374371
return httpx.AsyncClient(**kwargs)
375372

376-
self._http_client_factory = _tls_httpx_client_factory
377373
max_retries = 3
378374
base_delay = 1.0
379375

380376
for attempt in range(max_retries):
381377
try:
382-
client_factory = _tls_httpx_client_factory
383-
streamable_client = streamablehttp_client(
384-
uri, httpx_client_factory=client_factory, terminate_on_close=False
378+
http_client_instance = _tls_httpx_client_factory()
379+
streamable_client = streamable_http_client(
380+
uri, http_client=http_client_instance, terminate_on_close=True
385381
)
386382
http_transport = await self._exit_stack.enter_async_context(streamable_client)
387-
self._http, self._write, get_session_id = http_transport
388-
self._get_session_id = get_session_id
383+
self._http, self._write = http_transport
389384
self._session = await self._exit_stack.enter_async_context(ClientSession(self._http, self._write))
390385

391386
await self._session.initialize()
392-
self._session_id = self._get_session_id() if self._get_session_id else None
393387
response = await self._session.list_tools()
394388
tools = response.tools
395389
logger.info(
@@ -446,8 +440,6 @@ async def _cleanup_session(self) -> None:
446440
self._http = None
447441
self._write = None
448442
self._stdio = None
449-
self._get_session_id = None
450-
self._session_id = None
451443

452444
async def _reconnect_session(self) -> None:
453445
"""Tear down old session and reconnect to MCP server with linear backoff.
@@ -570,8 +562,8 @@ async def _execute_call() -> PluginResult:
570562
) from reconn_err
571563
logger.exception(pe)
572564
raise
573-
except McpError as e:
574-
logger.warning("McpError for plugin %s: %s", self.name, e)
565+
except MCPError as e:
566+
logger.warning("MCPError for plugin %s: %s", self.name, e)
575567
try:
576568
async with self._reconnect_lock:
577569
await self._reconnect_session()
@@ -637,30 +629,6 @@ async def shutdown(self) -> None:
637629

638630
if self._exit_stack:
639631
await self._exit_stack.aclose()
640-
if self._config and self._config.mcp and self._config.mcp.proto == TransportType.STREAMABLEHTTP:
641-
await self.__terminate_http_session()
642-
self._get_session_id = None
643-
self._session_id = None
644-
self._http_client_factory = None
645-
646-
async def __terminate_http_session(self) -> None:
647-
"""Terminate streamable HTTP session explicitly to avoid lingering server state."""
648-
if not self._session_id or not self._config or not self._config.mcp or not self._config.mcp.url:
649-
return
650-
# Third-Party
651-
from mcp.server.streamable_http import MCP_SESSION_ID_HEADER # pylint: disable=import-outside-toplevel
652-
653-
client_factory = self._http_client_factory
654-
try:
655-
if client_factory:
656-
client = client_factory()
657-
else:
658-
client = httpx.AsyncClient(follow_redirects=True)
659-
async with client:
660-
headers = {MCP_SESSION_ID_HEADER: self._session_id}
661-
await client.delete(self._config.mcp.url, headers=headers)
662-
except Exception as exc:
663-
logger.debug("Failed to terminate streamable HTTP session: %s", exc)
664632

665633

666634
class ExternalHookRef(HookRef):

‎cpex/framework/external/mcp/server/runtime.py‎

Lines changed: 43 additions & 41 deletions
Original file line numberDiff line numberDiff line change
@@ -5,10 +5,10 @@
55
SPDX-License-Identifier: Apache-2.0
66
Authors: Fred Araujo, Teryl Taylor
77
8-
MCP Plugin Runtime using FastMCP with SSL/TLS support.
8+
MCP Plugin Runtime using MCPServer with SSL/TLS support.
99
1010
This runtime does the following:
11-
- Uses FastMCP from the MCP Python SDK
11+
- Uses MCPServer from the MCP Python SDK
1212
- Supports both mTLS and non-mTLS configurations
1313
- Reads configuration from PLUGINS_SERVER_* environment variables or uses configurations
1414
the plugin config.yaml
@@ -19,7 +19,7 @@
1919
2020
>>> from cpex.framework.models import MCPServerConfig
2121
>>> config = MCPServerConfig(host="localhost", port=8000)
22-
>>> server = SSLCapableFastMCP(server_config=config, name="TestServer")
22+
>>> server = SSLCapableMCPServer(server_config=config, name="TestServer")
2323
>>> server.settings.host
2424
'localhost'
2525
>>> server.settings.port
@@ -29,7 +29,7 @@
2929
3030
>>> from cpex.framework.models import MCPServerConfig
3131
>>> config = MCPServerConfig(host="127.0.0.1", port=8000, tls=None)
32-
>>> server = SSLCapableFastMCP(server_config=config, name="NoTLSServer")
32+
>>> server = SSLCapableMCPServer(server_config=config, name="NoTLSServer")
3333
>>> ssl_config = server._get_ssl_config()
3434
>>> ssl_config
3535
{}
@@ -38,17 +38,17 @@
3838
3939
>>> from cpex.framework.models import MCPServerConfig
4040
>>> config = MCPServerConfig(host="localhost", port=9000)
41-
>>> server = SSLCapableFastMCP(server_config=config, name="ConfigTest")
41+
>>> server = SSLCapableMCPServer(server_config=config, name="ConfigTest")
4242
>>> server.server_config.host
4343
'localhost'
4444
>>> server.server_config.port
4545
9000
4646
47-
Settings are properly passed to FastMCP:
47+
Settings are properly passed to MCPServer:
4848
4949
>>> from cpex.framework.models import MCPServerConfig
5050
>>> config = MCPServerConfig(host="0.0.0.0", port=8080)
51-
>>> server = SSLCapableFastMCP(server_config=config, name="SettingsTest")
51+
>>> server = SSLCapableMCPServer(server_config=config, name="SettingsTest")
5252
>>> server.settings.host
5353
'0.0.0.0'
5454
>>> server.settings.port
@@ -66,7 +66,7 @@
6666

6767
# Third-Party
6868
from fastapi import Response, status
69-
from mcp.server.fastmcp import FastMCP
69+
from mcp.server.mcpserver import MCPServer
7070
from mcp.server.transport_security import TransportSecuritySettings
7171
from prometheus_client import REGISTRY, Gauge, generate_latest
7272

@@ -185,15 +185,15 @@ async def invoke_hook(hook_type: str, plugin_name: str, payload: Dict[str, Any],
185185
return await SERVER.invoke_hook(hook_type, plugin_name, payload, context)
186186

187187

188-
class SSLCapableFastMCP(FastMCP):
189-
"""FastMCP server with SSL/TLS support using MCPServerConfig.
188+
class SSLCapableMCPServer(MCPServer):
189+
"""MCPServer with SSL/TLS support using MCPServerConfig.
190190
191191
Examples:
192-
Create an SSL-capable FastMCP server:
192+
Create an SSL-capable MCPServer:
193193
194194
>>> from cpex.framework.models import MCPServerConfig
195195
>>> config = MCPServerConfig(host="127.0.0.1", port=8000)
196-
>>> server = SSLCapableFastMCP(server_config=config, name="TestServer")
196+
>>> server = SSLCapableMCPServer(server_config=config, name="TestServer")
197197
>>> server.settings.host
198198
'127.0.0.1'
199199
>>> server.settings.port
@@ -205,13 +205,13 @@ def __init__(self, server_config: MCPServerConfig, *args, **kwargs):
205205
206206
Args:
207207
server_config: the MCP server configuration including mTLS information.
208-
*args: Additional positional arguments passed to FastMCP.
209-
**kwargs: Additional keyword arguments passed to FastMCP.
208+
*args: Additional positional arguments passed to MCPServer.
209+
**kwargs: Additional keyword arguments passed to MCPServer.
210210
211211
Examples:
212212
>>> from cpex.framework.models import MCPServerConfig
213213
>>> config = MCPServerConfig(host="0.0.0.0", port=9000)
214-
>>> server = SSLCapableFastMCP(server_config=config, name="PluginServer")
214+
>>> server = SSLCapableMCPServer(server_config=config, name="PluginServer")
215215
>>> server.server_config.host
216216
'0.0.0.0'
217217
>>> server.server_config.port
@@ -220,13 +220,14 @@ def __init__(self, server_config: MCPServerConfig, *args, **kwargs):
220220
# Load server config from environment
221221

222222
self.server_config = server_config
223-
# Override FastMCP settings with our server config
224-
if "host" not in kwargs:
225-
kwargs["host"] = self.server_config.host
226-
if "port" not in kwargs:
227-
kwargs["port"] = self.server_config.port
228-
if self.server_config.uds and kwargs.get("transport_security") is None:
229-
kwargs["transport_security"] = TransportSecuritySettings(
223+
# MCPServer v2 does not accept host/port/transport_security in __init__;
224+
# transport_security is passed to streamable_http_app(), host/port to run methods.
225+
kwargs.pop("host", None)
226+
kwargs.pop("port", None)
227+
228+
transport_security = kwargs.pop("transport_security", None)
229+
if self.server_config.uds and transport_security is None:
230+
transport_security = TransportSecuritySettings(
230231
enable_dns_rebinding_protection=True,
231232
allowed_hosts=[
232233
"127.0.0.1",
@@ -245,6 +246,7 @@ def __init__(self, server_config: MCPServerConfig, *args, **kwargs):
245246
"http://[::1]:*",
246247
],
247248
)
249+
self._transport_security = transport_security
248250

249251
super().__init__(*args, **kwargs)
250252

@@ -257,7 +259,7 @@ def _get_ssl_config(self) -> dict:
257259
Examples:
258260
>>> from cpex.framework.models import MCPServerConfig
259261
>>> config = MCPServerConfig(host="127.0.0.1", port=8000, tls=None)
260-
>>> server = SSLCapableFastMCP(server_config=config, name="TestServer")
262+
>>> server = SSLCapableMCPServer(server_config=config, name="TestServer")
261263
>>> ssl_config = server._get_ssl_config()
262264
>>> ssl_config
263265
{}
@@ -361,10 +363,10 @@ async def metrics_disabled():
361363
# Create a minimal Starlette app with only the health endpoint
362364
health_app = Starlette(routes=routes)
363365

364-
logger.info(f"Starting HTTP health check server on {self.settings.host}:{health_port}")
366+
logger.info(f"Starting HTTP health check server on { self.server_config.host}:{health_port}")
365367
config = uvicorn.Config(
366368
app=health_app,
367-
host=self.settings.host,
369+
host= self.server_config.host,
368370
port=health_port,
369371
log_level="warning", # Reduce noise from health checks
370372
)
@@ -379,13 +381,13 @@ async def run_streamable_http_async(self) -> None:
379381
380382
>>> from cpex.framework.models import MCPServerConfig
381383
>>> config = MCPServerConfig(host="0.0.0.0", port=9000)
382-
>>> server = SSLCapableFastMCP(server_config=config, name="HTTPServer")
384+
>>> server = SSLCapableMCPServer(server_config=config, name="HTTPServer")
383385
>>> server.settings.host
384386
'0.0.0.0'
385387
>>> server.settings.port
386388
9000
387389
"""
388-
starlette_app = self.streamable_http_app()
390+
starlette_app = self.streamable_http_app(transport_security=getattr(self, '_transport_security', None))
389391

390392
# Add health check endpoint to main app
391393
# Third-Party
@@ -438,8 +440,8 @@ async def metrics_disabled():
438440
ssl_config = self._get_ssl_config()
439441
config_kwargs = {
440442
"app": starlette_app,
441-
"host": self.settings.host,
442-
"port": self.settings.port,
443+
"host": self.server_config.host,
444+
"port": self.server_config.port,
443445
"log_level": self.settings.log_level.lower(),
444446
}
445447
config_kwargs.update(ssl_config)
@@ -450,13 +452,13 @@ async def metrics_disabled():
450452
config_kwargs["uds"] = self.server_config.uds
451453
logger.info(f"Starting plugin server on unix socket {self.server_config.uds}")
452454
else:
453-
logger.info(f"Starting plugin server on {self.settings.host}:{self.settings.port}")
455+
logger.info(f"Starting plugin server on { self.server_config.host}:{ self.server_config.port}")
454456
config = uvicorn.Config(**config_kwargs) # type: ignore[arg-type]
455457
server = uvicorn.Server(config)
456458

457459
# If SSL is enabled, start a separate HTTP health check server
458460
if ssl_config and not self.server_config.uds:
459-
health_port = self.settings.port + 1000 # Use port+1000 for health checks
461+
health_port = self.server_config.port + 1000 # Use port+1000 for health checks
460462
logger.info(f"SSL enabled - starting separate HTTP health check on port {health_port}")
461463
# Run both servers concurrently
462464
await asyncio.gather(server.serve(), self._start_health_check_server(health_port))
@@ -466,7 +468,7 @@ async def metrics_disabled():
466468

467469

468470
async def run() -> None:
469-
"""Run the external plugin server with FastMCP.
471+
"""Run the external plugin server with MCPServer.
470472
471473
Supports both stdio and HTTP transports. Auto-detects transport based on stdin
472474
(if stdin is not a TTY, uses stdio mode), or you can explicitly set PLUGINS_TRANSPORT.
@@ -491,7 +493,7 @@ async def run() -> None:
491493
>>> SERVER is None
492494
True
493495
494-
FastMCP server names are defined as constants:
496+
MCPServer names are defined as constants:
495497
496498
>>> from cpex.framework.constants import MCP_SERVER_NAME
497499
>>> isinstance(MCP_SERVER_NAME, str)
@@ -524,33 +526,33 @@ async def run() -> None:
524526

525527
try:
526528
if transport == "stdio":
527-
# Create basic FastMCP server for stdio (no SSL support needed for stdio)
528-
mcp = FastMCP(
529+
# Create basic MCPServer for stdio (no SSL support needed for stdio)
530+
mcp = MCPServer(
529531
name=MCP_SERVER_NAME,
530532
instructions=MCP_SERVER_INSTRUCTIONS,
531533
)
532534

533-
# Register module-level tool functions with FastMCP
535+
# Register module-level tool functions with MCPServer
534536
mcp.tool(name=GET_PLUGIN_CONFIGS)(get_plugin_configs)
535537
mcp.tool(name=GET_PLUGIN_CONFIG)(get_plugin_config)
536538
mcp.tool(name=INVOKE_HOOK)(invoke_hook)
537539
# set the plugin_info gauge on startup
538540
PLUGIN_INFO.labels(server_name=MCP_SERVER_NAME, transport="stdio", ssl_enabled="false").set(1)
539541

540542
# Run with stdio transport
541-
logger.info("Starting MCP plugin server with FastMCP (stdio transport)")
543+
logger.info("Starting MCP plugin server with MCPServer (stdio transport)")
542544
await mcp.run_stdio_async()
543545

544546
else: # http or streamablehttp
545547
server_config: MCPServerConfig = SERVER.get_server_config()
546-
# Create FastMCP server with SSL support
547-
mcp = SSLCapableFastMCP(
548+
# Create MCPServer with SSL support
549+
mcp = SSLCapableMCPServer(
548550
server_config,
549551
name=MCP_SERVER_NAME,
550552
instructions=MCP_SERVER_INSTRUCTIONS,
551553
)
552554

553-
# Register module-level tool functions with FastMCP
555+
# Register module-level tool functions with MCPServer
554556
mcp.tool(name=GET_PLUGIN_CONFIGS)(get_plugin_configs)
555557
mcp.tool(name=GET_PLUGIN_CONFIG)(get_plugin_config)
556558
mcp.tool(name=INVOKE_HOOK)(invoke_hook)
@@ -564,7 +566,7 @@ async def run() -> None:
564566
f"Prometheus metrics available at http://{server_config.host}:{server_config.port}/metrics/prometheus"
565567
)
566568
# Run with streamable-http transport
567-
logger.info("Starting MCP plugin server with FastMCP (HTTP transport)")
569+
logger.info("Starting MCP plugin server with MCPServer (HTTP transport)")
568570
await mcp.run_streamable_http_async()
569571

570572
except Exception:

‎pyproject.toml‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,8 @@ dependencies = [
2626
"httpx>=0.28.1",
2727
"httpx[http2]>=0.28.1",
2828
"jinja2>=3.1.6",
29-
"mcp>=1.26.0",
29+
"mcp==2.0.0b1",
30+
"mcp-types==2.0.0b1",
3031
"orjson>=3.11.7",
3132
"prometheus-fastapi-instrumentator>=7.1.0",
3233
"prometheus_client>=0.24.1",

0 commit comments

Comments
 (0)