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/shared/test_streamable_http.py

@@ -893,6 +893,19 @@ async def test_streamablehttp_client_tool_invocation(initialized_client_session:
     assert result.content[0].text == "Called test_tool"
 
 
+@pytest.mark.anyio
+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."""
@@ -1106,12 +1119,14 @@ async def mock_delete(self: httpx.AsyncClient, *args: Any, **kwargs: Any) -> htt
                 await session.list_tools()
 
 
-@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
+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
+    Returns:
+        Tuple of (session_id, resumption_token, protocol_version, notifications)
+    """
     captured_session_id = None
     captured_resumption_token = None
     captured_notifications: list[types.ServerNotification] = []
@@ -1123,7 +1138,6 @@ async def message_handler(  # pragma: no branch
     ) -> None:
         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":
                     nonlocal first_notification_received
@@ -1181,8 +1195,90 @@ 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
+
+
+@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] = []
+
+    async def message_handler(  # pragma: no branch
+        message: RequestResponder[types.ServerRequest, types.ClientResult] | types.ServerNotification | Exception,
+    ) -> None:
+        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 streamablehttp_client(f"{server_url}/mcp", headers=headers) 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,
+            )
+            metadata = ClientMessageMetadata(
+                resumption_token=captured_resumption_token,
+            )
+
+            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_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] = []
+
+    async def message_handler(  # pragma: no branch
+        message: RequestResponder[types.ServerRequest, types.ClientResult] | types.ServerNotification | Exception,
+    ) -> None:
+        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
@@ -1204,8 +1300,10 @@ async def run_tool():
                 ),
                 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(