1
+ """Session management for different MCP transport types.
2
+
3
+ This module provides connection configurations and session management for various
4
+ MCP transport types including stdio, SSE, WebSocket, and streamable HTTP.
5
+ """
6
+
1
7
from __future__ import annotations
2
8
3
9
import os
30
36
31
37
32
38
class McpHttpClientFactory (Protocol ):
39
+ """Protocol for creating httpx.AsyncClient instances for MCP connections."""
40
+
33
41
def __call__ (
34
42
self ,
35
43
headers : dict [str , str ] | None = None ,
36
44
timeout : httpx .Timeout | None = None ,
37
45
auth : httpx .Auth | None = None ,
38
- ) -> httpx .AsyncClient : ...
46
+ ) -> httpx .AsyncClient :
47
+ """Create an httpx.AsyncClient instance.
48
+
49
+ Args:
50
+ headers: HTTP headers to include in requests.
51
+ timeout: Request timeout configuration.
52
+ auth: Authentication configuration.
53
+
54
+ Returns:
55
+ Configured httpx.AsyncClient instance.
56
+ """
57
+ ...
39
58
40
59
41
60
class StdioConnection (TypedDict ):
61
+ """Configuration for stdio transport connections to MCP servers."""
62
+
42
63
transport : Literal ["stdio" ]
43
64
44
65
command : str
@@ -74,6 +95,8 @@ class StdioConnection(TypedDict):
74
95
75
96
76
97
class SSEConnection (TypedDict ):
98
+ """Configuration for Server-Sent Events (SSE) transport connections to MCP servers."""
99
+
77
100
transport : Literal ["sse" ]
78
101
79
102
url : str
@@ -107,9 +130,10 @@ class SSEConnection(TypedDict):
107
130
108
131
109
132
class StreamableHttpConnection (TypedDict ):
110
- transport : Literal ["streamable_http" ]
111
133
"""Connection configuration for Streamable HTTP transport."""
112
134
135
+ transport : Literal ["streamable_http" ]
136
+
113
137
url : str
114
138
"""The URL of the endpoint to connect to."""
115
139
@@ -137,6 +161,8 @@ class StreamableHttpConnection(TypedDict):
137
161
138
162
139
163
class WebsocketConnection (TypedDict ):
164
+ """Configuration for WebSocket transport connections to MCP servers."""
165
+
140
166
transport : Literal ["websocket" ]
141
167
142
168
url : str
@@ -163,14 +189,16 @@ async def _create_stdio_session( # noqa: PLR0913
163
189
"""Create a new session to an MCP server using stdio.
164
190
165
191
Args:
166
- command: Command to execute
167
- args: Arguments for the command
168
- env: Environment variables for the command
169
- cwd: Working directory for the command
170
- encoding: Character encoding
171
- encoding_error_handler: How to handle encoding errors
172
- session_kwargs: Additional keyword arguments to pass to the ClientSession
192
+ command: Command to execute.
193
+ args: Arguments for the command.
194
+ env: Environment variables for the command.
195
+ cwd: Working directory for the command.
196
+ encoding: Character encoding.
197
+ encoding_error_handler: How to handle encoding errors.
198
+ session_kwargs: Additional keyword arguments to pass to the ClientSession.
173
199
200
+ Yields:
201
+ An initialized ClientSession.
174
202
"""
175
203
# NOTE: execution commands (e.g., `uvx` / `npx`) require PATH envvar to be set.
176
204
# To address this, we automatically inject existing PATH envvar into the `env` value,
@@ -210,14 +238,16 @@ async def _create_sse_session( # noqa: PLR0913
210
238
"""Create a new session to an MCP server using SSE.
211
239
212
240
Args:
213
- url: URL of the SSE server
214
- headers: HTTP headers to send to the SSE endpoint
215
- timeout: HTTP timeout
216
- sse_read_timeout: SSE read timeout
217
- session_kwargs: Additional keyword arguments to pass to the ClientSession
218
- httpx_client_factory: Custom factory for httpx.AsyncClient (optional)
219
- auth: httpx.Auth | None = None
241
+ url: URL of the SSE server.
242
+ headers: HTTP headers to send to the SSE endpoint.
243
+ timeout: HTTP timeout.
244
+ sse_read_timeout: SSE read timeout.
245
+ session_kwargs: Additional keyword arguments to pass to the ClientSession.
246
+ httpx_client_factory: Custom factory for httpx.AsyncClient (optional).
247
+ auth: Authentication for the HTTP client.
220
248
249
+ Yields:
250
+ An initialized ClientSession.
221
251
"""
222
252
# Create and store the connection
223
253
kwargs = {}
@@ -249,15 +279,17 @@ async def _create_streamable_http_session( # noqa: PLR0913
249
279
"""Create a new session to an MCP server using Streamable HTTP.
250
280
251
281
Args:
252
- url: URL of the endpoint to connect to
253
- headers: HTTP headers to send to the endpoint
254
- timeout: HTTP timeout
255
- sse_read_timeout: How long (in seconds) the client will wait for a new event before disconnecting
256
- terminate_on_close: Whether to terminate the session on close
257
- session_kwargs: Additional keyword arguments to pass to the ClientSession
258
- httpx_client_factory: Custom factory for httpx.AsyncClient (optional)
259
- auth: httpx.Auth | None = None
282
+ url: URL of the endpoint to connect to.
283
+ headers: HTTP headers to send to the endpoint.
284
+ timeout: HTTP timeout.
285
+ sse_read_timeout: How long the client will wait for a new event before disconnecting.
286
+ terminate_on_close: Whether to terminate the session on close.
287
+ session_kwargs: Additional keyword arguments to pass to the ClientSession.
288
+ httpx_client_factory: Custom factory for httpx.AsyncClient (optional).
289
+ auth: Authentication for the HTTP client.
260
290
291
+ Yields:
292
+ An initialized ClientSession.
261
293
"""
262
294
# Create and store the connection
263
295
kwargs = {}
@@ -288,12 +320,14 @@ async def _create_websocket_session(
288
320
"""Create a new session to an MCP server using Websockets.
289
321
290
322
Args:
291
- url: URL of the Websocket endpoint
292
- session_kwargs: Additional keyword arguments to pass to the ClientSession
323
+ url: URL of the Websocket endpoint.
324
+ session_kwargs: Additional keyword arguments to pass to the ClientSession.
293
325
294
- Raises :
295
- ImportError: If websockets package is not installed
326
+ Yields :
327
+ An initialized ClientSession.
296
328
329
+ Raises:
330
+ ImportError: If websockets package is not installed.
297
331
"""
298
332
try :
299
333
from mcp .client .websocket import websocket_client
0 commit comments