add on_session_created callback option

Author: crondinini-ant

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
@@ -29,6 +30,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 +44,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 +92,13 @@ async def sse_reader(
                                             logger.error(error_msg)  # pragma: no cover
                                             raise ValueError(error_msg)  # pragma: no cover
 
+                                        if on_session_created:
+                                            session_id = parse_qs(endpoint_parsed.query).get(
+                                                "sessionId", [None]
+                                            )[0]
+                                            if session_id:
+                                                on_session_created(session_id)
+
                                         task_status.started(endpoint_url)
 
                                     case "message":

tests/shared/test_sse.py

@@ -184,6 +184,25 @@ async def test_sse_client_basic_connection(server: None, server_url: str) -> Non
             assert isinstance(ping_result, EmptyResult)
 
 
+@pytest.mark.anyio
+async def test_sse_client_on_session_created(server: None, server_url: str) -> None:
+    captured_session_id: str | None = None
+
+    def on_session_created(session_id: str) -> None:
+        nonlocal captured_session_id
+        captured_session_id = session_id
+
+    async with sse_client(
+        server_url + "/sse", on_session_created=on_session_created
+    ) as streams:
+        async with ClientSession(*streams) as session:
+            result = await session.initialize()
+            assert isinstance(result, InitializeResult)
+
+    assert captured_session_id is not None
+    assert len(captured_session_id) > 0
+
+
 @pytest.fixture
 async def initialized_sse_client_session(server: None, server_url: str) -> AsyncGenerator[ClientSession, None]:
     async with sse_client(server_url + "/sse", sse_read_timeout=0.5) as streams: