Mirror x-mcp-header tool arguments into Mcp-Param-* request headers (SEP-2243) (#2990)
This commit is contained in:
parent
3a8da8c0c3
commit
5b2713d40c
2 changed files with 393 additions and 14 deletions
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -42,6 +42,7 @@ __all__ = [
|
|||
"InboundModernRoute",
|
||||
"MCP_METHOD_HEADER",
|
||||
"MCP_NAME_HEADER",
|
||||
"MCP_PARAM_HEADER_PREFIX",
|
||||
"MCP_PROTOCOL_VERSION_HEADER",
|
||||
"NAME_BEARING_METHODS",
|
||||
"X_MCP_HEADER_KEY",
|
||||
|
|
@ -49,6 +50,8 @@ __all__ = [
|
|||
"decode_header_value",
|
||||
"encode_header_value",
|
||||
"find_invalid_x_mcp_header",
|
||||
"mcp_param_headers",
|
||||
"x_mcp_header_map",
|
||||
]
|
||||
|
||||
MCP_PROTOCOL_VERSION_HEADER: Final = "mcp-protocol-version"
|
||||
|
|
@ -216,6 +219,56 @@ def find_invalid_x_mcp_header(input_schema: Any) -> str | None:
|
|||
return None
|
||||
|
||||
|
||||
MCP_PARAM_HEADER_PREFIX: Final = "Mcp-Param-"
|
||||
"""Prefix the `x-mcp-header` token is joined to, forming the per-parameter HTTP header name."""
|
||||
|
||||
|
||||
def x_mcp_header_map(input_schema: Any) -> dict[tuple[str, ...], str]:
|
||||
"""Map each property carrying a valid `x-mcp-header` to its annotation token, keyed by property path.
|
||||
|
||||
The key is the chain of `properties` keys from the schema root to the
|
||||
annotated property; a top-level property has a one-element path, a nested
|
||||
one a longer path. Call only on a schema that
|
||||
:func:`find_invalid_x_mcp_header` accepts; an invalid schema yields an
|
||||
undefined subset.
|
||||
"""
|
||||
mapping: dict[tuple[str, ...], str] = {}
|
||||
for path, schema in _walk_schema_positions(input_schema):
|
||||
if path and isinstance(header := schema.get(X_MCP_HEADER_KEY), str):
|
||||
mapping[path] = header
|
||||
return mapping
|
||||
|
||||
|
||||
def mcp_param_headers(header_map: Mapping[tuple[str, ...], str], arguments: Mapping[str, Any]) -> dict[str, str]:
|
||||
"""Build the `Mcp-Param-*` headers a `tools/call` mirrors from its arguments.
|
||||
|
||||
For each `(path, token)` in `header_map`, read the value at that property
|
||||
path in `arguments` and, when it is present and not `None`, emit
|
||||
`Mcp-Param-<token>` carrying it: `bool` as `true`/`false`, other scalars via
|
||||
`str`, each passed through :func:`encode_header_value` so a non-token value
|
||||
is base64-wrapped. A path that hits a missing key or a non-mapping node is
|
||||
skipped, matching the spec's "omit the header when no value is present".
|
||||
"""
|
||||
headers: dict[str, str] = {}
|
||||
for path, token in header_map.items():
|
||||
value = _value_at_path(arguments, path)
|
||||
if value is None:
|
||||
continue
|
||||
rendered = ("true" if value else "false") if isinstance(value, bool) else str(value)
|
||||
headers[f"{MCP_PARAM_HEADER_PREFIX}{token}"] = encode_header_value(rendered)
|
||||
return headers
|
||||
|
||||
|
||||
def _value_at_path(arguments: Mapping[str, Any], path: tuple[str, ...]) -> Any:
|
||||
"""Read the value at a `properties`-key path in `arguments`, or `None` if any step is missing or non-mapping."""
|
||||
node: Any = arguments
|
||||
for key in path:
|
||||
if not isinstance(node, Mapping):
|
||||
return None
|
||||
node = cast("Mapping[str, Any]", node).get(key)
|
||||
return node
|
||||
|
||||
|
||||
# INTERNAL_ERROR is deliberately unmapped (→ HTTP 200): the spec assigns no status to
|
||||
# -32603, and whether handler-origin errors get 5xx is an open S4 question — see TODO(L66).
|
||||
ERROR_CODE_HTTP_STATUS: Final[Mapping[int, int]] = MappingProxyType(
|
||||
|
|
|
|||
Loading…
Reference in a new issue