From d3efb8473c5e41a91e3e02e56c81703162f5e3a0 Mon Sep 17 00:00:00 2001
From: Nate Barbettini <nate@arcade-ai.com>
Date: Mon, 12 May 2025 11:25:07 -0700
Subject: [PATCH] fix: 204 is an acceptable response to DELETEing the session

---
 src/mcp/client/streamable_http.py    |  2 +-
 tests/shared/test_streamable_http.py | 66 ++++++++++++++++++++++++++++
 2 files changed, 67 insertions(+), 1 deletion(-)

diff --git a/src/mcp/client/streamable_http.py b/src/mcp/client/streamable_http.py
index 893aeb84a..3324dab5a 100644
--- a/src/mcp/client/streamable_http.py
+++ b/src/mcp/client/streamable_http.py
@@ -410,7 +410,7 @@ async def terminate_session(self, client: httpx.AsyncClient) -> None:
 
             if response.status_code == 405:
                 logger.debug("Server does not allow session termination")
-            elif response.status_code != 200:
+            elif response.status_code not in (200, 204):
                 logger.warning(f"Session termination failed: {response.status_code}")
         except Exception as exc:
             logger.warning(f"Session termination failed: {exc}")
diff --git a/tests/shared/test_streamable_http.py b/tests/shared/test_streamable_http.py
index 9b32254a9..f1c7ef809 100644
--- a/tests/shared/test_streamable_http.py
+++ b/tests/shared/test_streamable_http.py
@@ -960,6 +960,72 @@ async def test_streamablehttp_client_session_termination(
                 await session.list_tools()
 
 
+@pytest.mark.anyio
+async def test_streamablehttp_client_session_termination_204(
+    basic_server, basic_server_url, monkeypatch
+):
+    """Test client session termination functionality with a 204 response.
+
+    This test patches the httpx client to return a 204 response for DELETEs.
+    """
+
+    # Save the original delete method to restore later
+    original_delete = httpx.AsyncClient.delete
+
+    # Mock the client's delete method to return a 204
+    async def mock_delete(self, *args, **kwargs):
+        # Call the original method to get the real response
+        response = await original_delete(self, *args, **kwargs)
+
+        # Create a new response with 204 status code but same headers
+        mocked_response = httpx.Response(
+            204,
+            headers=response.headers,
+            content=response.content,
+            request=response.request,
+        )
+        return mocked_response
+
+    # Apply the patch to the httpx client
+    monkeypatch.setattr(httpx.AsyncClient, "delete", mock_delete)
+
+    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,
+        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) == 4
+
+    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",
+            ):
+                await session.list_tools()
+
+
 @pytest.mark.anyio
 async def test_streamablehttp_client_resumption(event_server):
     """Test client session to resume a long running tool."""