test: add coverage for extra_headers in HTTP transport

- Add test_streamablehttp_client_tool_invocation_with_extra_headers for POST requests
- Add test_streamablehttp_client_resumption_with_extra_headers for resumption with extra headers
- Refactor common resumption setup code into _setup_resumption_test helper
- Achieve 100% coverage for streamable_http.py

Author: damianoneill

tests/client/test_extra_headers.py

@@ -557,7 +557,7 @@ async def mock_server():
                     from mcp.types import ListToolsResult
 
                     result = ServerResult(ListToolsResult(tools=[]))
-                else:
+                else:  # pragma: no cover
                     continue
 
                 await server_to_client_send.send(
@@ -648,7 +648,7 @@ async def mock_server():
         assert jsonrpc_request.root.method == "logging/setLevel"
 
         # Capture the metadata that was passed with the request
-        if isinstance(session_message.metadata, ClientMessageMetadata):
+        if isinstance(session_message.metadata, ClientMessageMetadata):  # pragma: no branch
             captured_metadata.append(session_message.metadata)
 
         # Send response

tests/shared/test_streamable_http.py

@@ -1032,7 +1032,20 @@ async def test_streamable_http_client_tool_invocation(initialized_client_session
 
 
 @pytest.mark.anyio
-async def test_streamable_http_client_error_handling(initialized_client_session: ClientSession):
+async def test_streamablehttp_client_tool_invocation_with_extra_headers(initialized_client_session: ClientSession):
+    """Test HTTP POST request with extra headers."""
+    result = await initialized_client_session.call_tool(
+        "test_tool",
+        {},
+        extra_headers={"X-Custom-Header": "test-value"},
+    )
+    assert len(result.content) == 1
+    assert result.content[0].type == "text"
+    assert result.content[0].text == "Called test_tool"
+
+
+@pytest.mark.anyio
+async def test_streamablehttp_client_error_handling(initialized_client_session: ClientSession):
     """Test error handling in client."""
     with pytest.raises(McpError) as exc_info:
         await initialized_client_session.read_resource(uri=AnyUrl("unknown://test-error"))
@@ -1245,26 +1258,27 @@ async def mock_delete(self: httpx.AsyncClient, *args: Any, **kwargs: Any) -> htt
                     await session.list_tools()
 
 
-@pytest.mark.anyio
-async def test_streamable_http_client_resumption(event_server: tuple[SimpleEventStore, str]):
-    """Test client session resumption using sync primitives for reliable coordination."""
-    _, server_url = event_server
+async def _setup_resumption_test(
+    server_url: str,
+) -> tuple[str | None, str | None, str | int | None, list[types.ServerNotification]]:
+    """Helper function to set up a resumption test by starting a session and capturing resumption state.
 
-    # Variables to track the state
-    captured_session_id = None
-    captured_resumption_token = None
-    captured_notifications: list[types.ServerNotification] = []
-    captured_protocol_version = None
-    first_notification_received = False
+    Returns:
+        Tuple of (session_id, resumption_token, protocol_version, notifications)
+    """
+    captured_session_id = None  # pragma: no cover
+    captured_resumption_token = None  # pragma: no cover
+    captured_notifications: list[types.ServerNotification] = []  # pragma: no cover
+    captured_protocol_version = None  # pragma: no cover
+    first_notification_received = False  # pragma: no cover
 
     async def message_handler(  # pragma: no branch
         message: RequestResponder[types.ServerRequest, types.ClientResult] | types.ServerNotification | Exception,
-    ) -> None:
+    ) -> None:  # pragma: no cover
         if isinstance(message, types.ServerNotification):  # pragma: no branch
             captured_notifications.append(message)
-            # Look for our first notification
             if isinstance(message.root, types.LoggingMessageNotification):  # pragma: no branch
-                if message.root.params.data == "First notification before lock":
+                if message.root.params.data == "First notification before lock":  # pragma: no branch
                     nonlocal first_notification_received
                     first_notification_received = True
 
@@ -1320,8 +1334,32 @@ async def run_tool():
     assert isinstance(captured_notifications[0].root, types.LoggingMessageNotification)  # pragma: no cover
     assert captured_notifications[0].root.params.data == "First notification before lock"  # pragma: no cover
 
-    # Clear notifications for the second phase
-    captured_notifications = []  # pragma: no cover
+    return (
+        captured_session_id,
+        captured_resumption_token,
+        captured_protocol_version,
+        captured_notifications,
+    )  # pragma: no cover  # noqa: E501
+
+
+@pytest.mark.anyio
+async def test_streamablehttp_client_resumption(event_server: tuple[SimpleEventStore, str]):
+    """Test client session resumption using sync primitives for reliable coordination."""
+    _, server_url = event_server
+
+    # Set up the initial session and capture resumption state
+    captured_session_id, captured_resumption_token, captured_protocol_version, _ = await _setup_resumption_test(
+        server_url
+    )
+
+    # Track notifications for the resumed session
+    captured_notifications: list[types.ServerNotification] = []  # pragma: no cover
+
+    async def message_handler(  # pragma: no branch
+        message: RequestResponder[types.ServerRequest, types.ClientResult] | types.ServerNotification | Exception,
+    ) -> None:  # pragma: no cover
+        if isinstance(message, types.ServerNotification):  # pragma: no branch
+            captured_notifications.append(message)
 
     # Now resume the session with the same mcp-session-id and protocol version
     headers: dict[str, Any] = {}  # pragma: no cover
@@ -1369,6 +1407,73 @@ async def run_tool():
             assert captured_notifications[0].root.params.data == "Second notification after lock"  # pragma: no cover
 
 
+@pytest.mark.anyio
+async def test_streamablehttp_client_resumption_with_extra_headers(event_server: tuple[SimpleEventStore, str]):
+    """Test client session resumption with extra headers."""
+    _, server_url = event_server
+
+    # Set up the initial session and capture resumption state
+    captured_session_id, captured_resumption_token, captured_protocol_version, _ = await _setup_resumption_test(
+        server_url
+    )
+
+    # Track notifications for the resumed session
+    captured_notifications: list[types.ServerNotification] = []  # pragma: no cover
+
+    async def message_handler(  # pragma: no branch
+        message: RequestResponder[types.ServerRequest, types.ClientResult] | types.ServerNotification | Exception,
+    ) -> None:  # pragma: no cover
+        if isinstance(message, types.ServerNotification):  # pragma: no branch
+            captured_notifications.append(message)
+
+    # Now resume the session with the same mcp-session-id and protocol version
+    headers: dict[str, Any] = {}  # pragma: no cover
+    if captured_session_id:  # pragma: no cover
+        headers[MCP_SESSION_ID_HEADER] = captured_session_id
+    if captured_protocol_version:  # pragma: no cover
+        headers[MCP_PROTOCOL_VERSION_HEADER] = captured_protocol_version
+
+    async with create_mcp_http_client(headers=headers) as httpx_client:
+        async with streamable_http_client(f"{server_url}/mcp", http_client=httpx_client) as (
+            read_stream,
+            write_stream,
+            _,
+        ):
+            async with ClientSession(read_stream, write_stream, message_handler=message_handler) as session:
+                result = await session.send_request(
+                    types.ClientRequest(
+                        types.CallToolRequest(
+                            params=types.CallToolRequestParams(name="release_lock", arguments={}),
+                        )
+                    ),
+                    types.CallToolResult,
+                )
+                # Test resumption WITH extra_headers
+                metadata = ClientMessageMetadata(
+                    resumption_token=captured_resumption_token,
+                    extra_headers={"X-Resumption-Test": "test-value"},
+                )
+
+                result = await session.send_request(
+                    types.ClientRequest(
+                        types.CallToolRequest(
+                            params=types.CallToolRequestParams(name="wait_for_lock_with_notification", arguments={}),
+                        )
+                    ),
+                    types.CallToolResult,
+                    metadata=metadata,
+                )
+                assert len(result.content) == 1
+                assert result.content[0].type == "text"
+                assert result.content[0].text == "Completed"
+
+                # We should have received the remaining notifications
+                assert len(captured_notifications) == 1
+
+                assert isinstance(captured_notifications[0].root, types.LoggingMessageNotification)
+                assert captured_notifications[0].root.params.data == "Second notification after lock"
+
+
 @pytest.mark.anyio
 async def test_streamablehttp_server_sampling(basic_server: None, basic_server_url: str):
     """Test server-initiated sampling request through streamable HTTP transport."""