Add on_session_created callback option (#1710)

src/mcp/client/sse.py

@@ -1,7 +1,8 @@
 import logging
+from collections.abc import Callable
 from contextlib import asynccontextmanager
 from typing import Any
-from urllib.parse import urljoin, urlparse
+from urllib.parse import parse_qs, urljoin, urlparse
 
 import anyio
 import httpx
@@ -21,6 +22,11 @@ def remove_request_params(url: str) -> str:
     return urljoin(url, urlparse(url).path)
 
 
+def _extract_session_id_from_endpoint(endpoint_url: str) -> str | None:
+    query_params = parse_qs(urlparse(endpoint_url).query)
+    return query_params.get("sessionId", [None])[0] or query_params.get("session_id", [None])[0]
+
+
 @asynccontextmanager
 async def sse_client(
     url: str,
@@ -29,6 +35,7 @@ async def sse_client(
     sse_read_timeout: float = 60 * 5,
     httpx_client_factory: McpHttpClientFactory = create_mcp_http_client,
     auth: httpx.Auth | None = None,
+    on_session_created: Callable[[str], None] | None = None,
 ):
     """
     Client transport for SSE.
@@ -42,6 +49,7 @@ async def sse_client(
         timeout: HTTP timeout for regular operations.
         sse_read_timeout: Timeout for SSE read operations.
         auth: Optional HTTPX authentication handler.
+        on_session_created: Optional callback invoked with the session ID when received.
     """
     read_stream: MemoryObjectReceiveStream[SessionMessage | Exception]
     read_stream_writer: MemoryObjectSendStream[SessionMessage | Exception]
@@ -89,6 +97,11 @@ async def sse_client(
                                             logger.error(error_msg)  # pragma: no cover
                                             raise ValueError(error_msg)  # pragma: no cover
 
+                                        if on_session_created:
+                                            session_id = _extract_session_id_from_endpoint(endpoint_url)
+                                            if session_id:
+                                                on_session_created(session_id)
+
                                         task_status.started(endpoint_url)
 
                                     case "message":