Clean up code

Author: aravind-segu

src/mcp/client/streamable_http.py

@@ -42,6 +42,7 @@
 MCP_SESSION_ID = "mcp-session-id"
 LAST_EVENT_ID = "last-event-id"
 CONTENT_TYPE = "content-type"
+HEADER_CAPTURE = "[TESTING_HEADER_CAPTURE]"
 ACCEPT = "Accept"
 
 
@@ -275,7 +276,7 @@ async def _handle_resumption_request(self, ctx: RequestContext) -> None:
     async def _is_testing_header_capture(self, response: httpx.Response) -> str | None:
         try:
             content = await response.aread()
-            if content.decode().startswith("[TESTING_HEADER_CAPTURE]"):
+            if content.decode().startswith(HEADER_CAPTURE):
                 return content.decode()
         except Exception as _:
             return None
@@ -306,6 +307,10 @@ async def _handle_post_request(self, ctx: RequestContext) -> None:
                     )
                 return
 
+            # To test if headers are being forwarded correctly, in unit tests
+            # we have a mock server that returns a 418 status code with the
+            # HEADER_CAPTURE prefix. If the response has this status code
+            # with the prefix, return the response content as part of the error message.
             if response.status_code == 418:
                 test_error_message = await self._is_testing_header_capture(response)
                 # If this is coming from the test case return the response content

tests/shared/test_streamable_http.py

@@ -25,7 +25,7 @@
 
 import mcp.types as types
 from mcp.client.session import ClientSession
-from mcp.client.streamable_http import streamablehttp_client
+from mcp.client.streamable_http import HEADER_CAPTURE, streamablehttp_client
 from mcp.server import Server
 from mcp.server.streamable_http import (
     MCP_SESSION_ID_HEADER,
@@ -262,7 +262,7 @@ async def header_capture_wrapper(scope, receive, send):
 
         # Return error response with headers in body
         response = Response(
-            "[TESTING_HEADER_CAPTURE]:" + json.dumps({"headers": headers}),
+            HEADER_CAPTURE + json.dumps({"headers": headers}),
             status_code=418,
         )
         await response(scope, receive, send)
@@ -279,7 +279,7 @@ async def header_capture_wrapper(scope, receive, send):
 
 
 def _get_captured_headrs(str) -> dict[str, str]:
-    return json.loads(str.split("[TESTING_HEADER_CAPTURE]:")[1])["headers"]
+    return json.loads(str.split(HEADER_CAPTURE)[1])["headers"]
 
 
 def run_server(
@@ -356,21 +356,23 @@ def _start_basic_server(
 
     # Wait for server to be running
     max_attempts = 20
-    for attempt in range(max_attempts):
+    attempt = 0
+    while attempt < max_attempts:
         try:
             with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
                 s.connect(("127.0.0.1", basic_server_port))
                 break
         except ConnectionRefusedError:
             time.sleep(0.1)
+            attempt += 1
     else:
         raise RuntimeError(f"Server failed to start after {max_attempts} attempts")
 
-    try:
-        yield
-    finally:
-        proc.kill()
-        proc.join(timeout=2)
+    yield
+
+    # Clean up
+    proc.kill()
+    proc.join(timeout=2)
 
 
 @pytest.fixture