Mirror x-mcp-header tool arguments into Mcp-Param-* request headers (SEP-2243) (#2990)

src/mcp/client/session.py

@@ -41,6 +41,8 @@ from mcp.shared.inbound import (
     NAME_BEARING_METHODS,
     encode_header_value,
     find_invalid_x_mcp_header,
+    mcp_param_headers,
+    x_mcp_header_map,
 )
 from mcp.shared.jsonrpc_dispatcher import JSONRPCDispatcher
 from mcp.shared.message import ClientMessageMetadata, SessionMessage
@@ -68,7 +70,10 @@ def _make_handshake_stamp(protocol_version: str) -> Callable[[dict[str, Any], Ca
 
 
 def _make_modern_stamp(
-    protocol_version: str, client_info: dict[str, Any], capabilities: dict[str, Any]
+    protocol_version: str,
+    client_info: dict[str, Any],
+    capabilities: dict[str, Any],
+    resolve_param_headers: Callable[[str, Mapping[str, Any]], dict[str, str]],
 ) -> Callable[[dict[str, Any], CallOptions], None]:
     def stamp(data: dict[str, Any], opts: CallOptions) -> None:
         params = data.setdefault("params", {})
@@ -83,6 +88,8 @@ def _make_modern_stamp(
         name_key = NAME_BEARING_METHODS.get(data["method"])
         if name_key is not None and isinstance(name := params.get(name_key), str):
             headers[MCP_NAME_HEADER] = encode_header_value(name)
+        if data["method"] == "tools/call" and isinstance(name := params.get("name"), str):
+            headers.update(resolve_param_headers(name, params.get("arguments") or {}))
 
     return stamp
 
@@ -215,6 +222,7 @@ class ClientSession:
         self._logging_callback = logging_callback or _default_logging_callback
         self._message_handler = message_handler or _default_message_handler
         self._tool_output_schemas: dict[str, dict[str, Any] | None] = {}
+        self._x_mcp_header_maps: dict[str, dict[tuple[str, ...], str]] = {}
         self._initialize_result: types.InitializeResult | None = None
         self._discover_result: types.DiscoverResult | None = None
         self._negotiated_version: str | None = None
@@ -393,7 +401,7 @@ class ClientSession:
                 )
             client_info = self._client_info.model_dump(by_alias=True, mode="json", exclude_none=True)
             capabilities = self._build_capabilities().model_dump(by_alias=True, mode="json", exclude_none=True)
-            self._stamp = _make_modern_stamp(mutual[-1], client_info, capabilities)
+            self._stamp = _make_modern_stamp(mutual[-1], client_info, capabilities, self._resolve_param_headers)
             self._discover_result = result
             self._initialize_result = None
             self._negotiated_version = mutual[-1]
@@ -646,6 +654,11 @@ class ClientSession:
     ) -> types.CallToolResult | types.InputRequiredResult:
         """Send a tools/call request with optional progress callback support.
 
+        On a modern (2026-07-28) connection, arguments annotated with `x-mcp-header`
+        in the tool's input schema are mirrored into `Mcp-Param-*` request headers.
+        The annotations are read from the tool's last `list_tools` entry, so list
+        the tool before calling it to enable header emission.
+
         Args:
             input_responses: Responses to a prior `InputRequiredResult.input_requests`.
             request_state: Opaque state echoed from a prior `InputRequiredResult`.
@@ -657,7 +670,6 @@ class ClientSession:
             RuntimeError: If the server returns an `InputRequiredResult` and
                 ``allow_input_required`` is ``False``.
         """
-
         result = await self.send_request(
             types.CallToolRequest(
                 params=types.CallToolRequestParams(
@@ -683,6 +695,13 @@ class ClientSession:
             )
         return result
 
+    def _resolve_param_headers(self, name: str, arguments: Mapping[str, Any]) -> dict[str, str]:
+        """`Mcp-Param-*` headers for a `tools/call`, or empty when the tool was never listed."""
+        header_map = self._x_mcp_header_maps.get(name)
+        if header_map is None:
+            return {}
+        return mcp_param_headers(header_map, arguments)
+
     async def _validate_tool_result(self, name: str, result: types.CallToolResult) -> None:
         """Validate the structured content of a tool result against its output schema."""
         if name not in self._tool_output_schemas:
@@ -767,7 +786,12 @@ class ClientSession:
             for tool in result.tools:
                 if (reason := find_invalid_x_mcp_header(tool.input_schema)) is not None:
                     logger.warning("dropping tool %r: invalid x-mcp-header (%s)", tool.name, reason)
+                    # Evict any map cached from a prior valid listing so a stale entry can't
+                    # mirror headers for a tool this listing dropped.
+                    self._x_mcp_header_maps.pop(tool.name, None)
                     continue
+                # Cache the arg→header map so a later tools/call mirrors it into Mcp-Param-* headers.
+                self._x_mcp_header_maps[tool.name] = x_mcp_header_map(tool.input_schema)
                 kept.append(tool)
             result.tools = kept