Skip to content

Instantly share code, notes, and snippets.

@shamashel
Created March 16, 2026 13:28
Show Gist options
  • Select an option

  • Save shamashel/51e70376e70e5464778f49b996b90cb0 to your computer and use it in GitHub Desktop.

Select an option

Save shamashel/51e70376e70e5464778f49b996b90cb0 to your computer and use it in GitHub Desktop.
PoC: FastMCP proxy drops all auth headers on the downstream hop (LLMX-421)
"""
PoC: FastMCP proxy header propagation — master vs proposed fix.
Monkey-patches ProxyProvider._get_client to capture the outbound transport
headers at the exact moment the backend client is created, proving whether
inbound request headers are carried to the downstream MCP call.
Usage:
uv run --package platform_api python poc_header_propagation.py
"""
from __future__ import annotations
import asyncio
import json
from typing import Any
from unittest.mock import patch
from fastmcp import Client, FastMCP
from fastmcp.client.client import Client as ClientClass
from fastmcp.server.providers.proxy import (
FastMCPProxy,
ProxyClient,
ProxyProvider,
ProxyTool,
)
HEADERS_WE_CARE_ABOUT = [
"authorization",
"cookie",
"x-user-id",
"x-organization-id",
"x-organization-ids",
"x-custom-header",
]
TEST_HEADERS = {
"Authorization": "Bearer test-token-12345",
"Cookie": ".mxa.auth=fake-cookie-value",
"X-User-Id": "48593",
"X-Organization-Id": "10646",
"X-Organization-Ids": "10646,10647",
"X-Custom-Header": "should-forward",
}
def make_downstream_server() -> FastMCP:
"""A real downstream FastMCP server (simulates WunderGraph)."""
server = FastMCP("Downstream")
@server.tool()
def echo_headers(dummy: str = "x") -> str:
"""A simple tool so the proxy has something to discover and call."""
return "ok"
return server
def create_proxy_master(target_url: str) -> FastMCP:
"""Production code: FastMCP.as_proxy()."""
return FastMCP.as_proxy(target_url, name="Proxy-master")
def create_proxy_fixed(target_url: str) -> FastMCP:
"""Proposed fix: client_factory that injects inbound request headers."""
from fastmcp.server.dependencies import get_http_headers
base_client = ProxyClient(target_url)
def client_factory() -> ClientClass:
fresh = base_client.new()
inbound_headers = get_http_headers()
fresh.transport.headers = {**fresh.transport.headers, **inbound_headers}
return fresh
return FastMCPProxy(client_factory=client_factory, name="Proxy-fixed")
captured_outbound_headers: list[dict[str, str]] = []
original_get_client = ProxyTool._get_client
async def instrumented_get_client(self):
"""Monkey-patched _get_client that captures the outbound transport headers."""
client = self._client_factory()
import inspect as _inspect
if _inspect.isawaitable(client):
client = await client
transport_headers = {}
if hasattr(client, "transport") and hasattr(client.transport, "headers"):
transport_headers = dict(client.transport.headers)
captured_outbound_headers.append(transport_headers)
return client
async def call_tool_via_proxy(proxy_url: str, headers: dict[str, str]) -> None:
"""Connect to proxy over HTTP with test headers, call a tool."""
from fastmcp.client.transports import StreamableHttpTransport
transport = StreamableHttpTransport(proxy_url, headers=headers)
async with Client(transport=transport) as client:
await client.call_tool("echo_headers", {"dummy": "test"})
async def run_test(label: str, proxy_factory, downstream_url: str) -> dict[str, str]:
"""Capture what headers the proxy puts on its outbound transport."""
from fastmcp.utilities.tests import run_server_async
captured_outbound_headers.clear()
proxy = proxy_factory(downstream_url)
with patch.object(ProxyTool, "_get_client", instrumented_get_client):
async with run_server_async(proxy, transport="http") as proxy_url:
await call_tool_via_proxy(proxy_url, TEST_HEADERS)
# The last _get_client call is the tool execution (not listing)
if captured_outbound_headers:
return captured_outbound_headers[-1]
return {}
async def main():
from fastmcp.utilities.tests import run_server_async
downstream = make_downstream_server()
async with run_server_async(downstream, transport="http") as downstream_url:
print(f"Downstream: {downstream_url}")
print("=" * 72)
# --- Test A ---
print("\n[A] MASTER — FastMCP.as_proxy()")
print("-" * 72)
master_hdrs = await run_test("master", create_proxy_master, downstream_url)
master_fwd = {}
master_drop = []
for h in HEADERS_WE_CARE_ABOUT:
val = master_hdrs.get(h) or master_hdrs.get(h.lower())
if val:
master_fwd[h] = val
else:
master_drop.append(h)
print(f" Transport headers on outbound client:")
if master_hdrs:
for k, v in sorted(master_hdrs.items()):
marker = " <-- auth" if k.lower() in {h.lower() for h in HEADERS_WE_CARE_ABOUT} else ""
print(f" {k}: {v}{marker}")
else:
print(f" (empty — no headers set on transport)")
print(f"\n Auth/org headers:")
for h in HEADERS_WE_CARE_ABOUT:
val = master_fwd.get(h)
print(f" {'✓' if val else '✗'} {h}: {val or 'MISSING'}")
# --- Test B ---
print("\n[B] PROPOSED FIX — FastMCPProxy + get_http_headers()")
print("-" * 72)
fixed_hdrs = await run_test("fixed", create_proxy_fixed, downstream_url)
fixed_fwd = {}
fixed_drop = []
for h in HEADERS_WE_CARE_ABOUT:
val = fixed_hdrs.get(h) or fixed_hdrs.get(h.lower())
if val:
fixed_fwd[h] = val
else:
fixed_drop.append(h)
print(f" Transport headers on outbound client:")
if fixed_hdrs:
for k, v in sorted(fixed_hdrs.items()):
marker = " <-- auth" if k.lower() in {h.lower() for h in HEADERS_WE_CARE_ABOUT} else ""
print(f" {k}: {v}{marker}")
else:
print(f" (empty — no headers set on transport)")
print(f"\n Auth/org headers:")
for h in HEADERS_WE_CARE_ABOUT:
val = fixed_fwd.get(h)
print(f" {'✓' if val else '✗'} {h}: {val or 'MISSING'}")
# --- Summary ---
print("\n" + "=" * 72)
print("SUMMARY")
print("=" * 72)
newly_forwarded = set(fixed_fwd) - set(master_fwd)
print(f"\n Headers tested: {HEADERS_WE_CARE_ABOUT}")
print(f" Master forwarded: {sorted(master_fwd) or '(none)'}")
print(f" Master dropped: {sorted(master_drop) or '(none)'}")
print(f" Fixed forwarded: {sorted(fixed_fwd) or '(none)'}")
print(f" Fixed dropped: {sorted(fixed_drop) or '(none)'}")
print(f" Recovered by fix: {sorted(newly_forwarded) or '(none)'}")
if master_drop and newly_forwarded:
print(f"\n ⚠ PRODUCTION BUG CONFIRMED + FIX VALIDATED")
print(f" Master drops: {sorted(master_drop)}")
print(f" Fix recovers: {sorted(newly_forwarded)}")
elif master_drop:
print(f"\n ⚠ Bug confirmed but fix doesn't recover all headers.")
elif not master_drop:
print(f"\n Master already forwards all tested headers on the transport.")
print(f"\n _get_client calls captured: {len(captured_outbound_headers)}")
if __name__ == "__main__":
asyncio.run(main())
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment