terminate on close instead of callback to terminate

Author: ihrpr

src/mcp/client/streamable_http.py

@@ -35,7 +35,6 @@
 SessionMessageOrError = SessionMessage | Exception
 StreamWriter = MemoryObjectSendStream[SessionMessageOrError]
 StreamReader = MemoryObjectReceiveStream[SessionMessage]
-TerminateCallback = Callable[[], Awaitable[None]]
 GetSessionIdCallback = Callable[[], str | None]
 
 MCP_SESSION_ID = "mcp-session-id"
@@ -413,11 +412,11 @@ async def streamablehttp_client(
     headers: dict[str, Any] | None = None,
     timeout: timedelta = timedelta(seconds=30),
     sse_read_timeout: timedelta = timedelta(seconds=60 * 5),
+    terminate_on_close: bool = True,
 ) -> AsyncGenerator[
     tuple[
         MemoryObjectReceiveStream[SessionMessage | Exception],
         MemoryObjectSendStream[SessionMessage],
-        TerminateCallback,
         GetSessionIdCallback,
     ],
     None,
@@ -432,7 +431,6 @@ async def streamablehttp_client(
         Tuple containing:
             - read_stream: Stream for reading messages from the server
             - write_stream: Stream for sending messages to the server
-            - terminate_callback: Async function to terminate the session - send DELETE
             - get_session_id_callback: Function to retrieve the current session ID
     """
     transport = StreamableHTTPTransport(url, headers, timeout, sse_read_timeout)
@@ -461,9 +459,6 @@ def start_get_stream() -> None:
                         transport.handle_get_stream, client, read_stream_writer
                     )
 
-                async def terminate_session() -> None:
-                    await transport.terminate_session(client)
-
                 tg.start_soon(
                     transport.post_writer,
                     client,
@@ -477,10 +472,11 @@ async def terminate_session() -> None:
                     yield (
                         read_stream,
                         write_stream,
-                        terminate_session,
                         transport.get_session_id,
                     )
                 finally:
+                    if transport.session_id and terminate_on_close:
+                        await transport.terminate_session(client)
                     tg.cancel_scope.cancel()
         finally:
             await read_stream_writer.aclose()

tests/shared/test_streamable_http.py

@@ -779,7 +779,6 @@ async def initialized_client_session(basic_server, basic_server_url):
         read_stream,
         write_stream,
         _,
-        _,
     ):
         async with ClientSession(
             read_stream,
@@ -796,7 +795,6 @@ async def test_streamablehttp_client_basic_connection(basic_server, basic_server
         read_stream,
         write_stream,
         _,
-        _,
     ):
         async with ClientSession(
             read_stream,
@@ -854,7 +852,6 @@ async def test_streamablehttp_client_session_persistence(
         read_stream,
         write_stream,
         _,
-        _,
     ):
         async with ClientSession(
             read_stream,
@@ -885,7 +882,6 @@ async def test_streamablehttp_client_json_response(
         read_stream,
         write_stream,
         _,
-        _,
     ):
         async with ClientSession(
             read_stream,
@@ -928,7 +924,6 @@ async def message_handler(
         read_stream,
         write_stream,
         _,
-        _,
     ):
         async with ClientSession(
             read_stream, write_stream, message_handler=message_handler
@@ -961,24 +956,36 @@ async def test_streamablehttp_client_session_termination(
 ):
     """Test client session termination functionality."""
 
+    captured_session_id = None
+
     # Create the streamablehttp_client with a custom httpx client to capture headers
     async with streamablehttp_client(f"{basic_server_url}/mcp") as (
         read_stream,
         write_stream,
-        terminate_session,
-        _,
+        get_session_id,
     ):
         async with ClientSession(read_stream, write_stream) as session:
             # Initialize the session
             result = await session.initialize()
             assert isinstance(result, InitializeResult)
+            captured_session_id = get_session_id()
+            assert captured_session_id is not None
 
             # Make a request to confirm session is working
             tools = await session.list_tools()
             assert len(tools.tools) == 3
 
-            # After exiting ClientSession context, explicitly terminate the session
-            await terminate_session()
+    headers = {}
+    if captured_session_id:
+        headers[MCP_SESSION_ID_HEADER] = captured_session_id
+
+    async with streamablehttp_client(f"{basic_server_url}/mcp", headers=headers) as (
+        read_stream,
+        write_stream,
+        _,
+    ):
+        async with ClientSession(read_stream, write_stream) as session:
+            # Attempt to make a request after termination
             with pytest.raises(
                 McpError,
                 match="Session terminated",
@@ -1015,10 +1022,9 @@ async def on_resumption_token_update(token: str) -> None:
         captured_resumption_token = token
 
     # First, start the client session and begin the long-running tool
-    async with streamablehttp_client(f"{server_url}/mcp") as (
+    async with streamablehttp_client(f"{server_url}/mcp", terminate_on_close=False) as (
         read_stream,
         write_stream,
-        _,
         get_session_id,
     ):
         async with ClientSession(
@@ -1072,7 +1078,6 @@ async def run_tool():
         read_stream,
         write_stream,
         _,
-        _,
     ):
         async with ClientSession(
             read_stream, write_stream, message_handler=message_handler