feat(sdk-python): add header propagation for X-AIMock-Strict

ContextVar-based ambient state + httpx event hook. Incoming
x-aimock-* headers from AG-UI requests are forwarded to outgoing
LLM API calls. Keys normalized to lowercase. Warning emitted when
install_httpx_hook receives an unrecognized client type.

Author: jpr5

sdk-python/copilotkit/__init__.py

@@ -7,6 +7,7 @@
 from .langgraph_agui_agent import LangGraphAGUIAgent
 from .copilotkit_lg_middleware import CopilotKitMiddleware
 from ag_ui_langgraph.middlewares.state_streaming import StateStreamingMiddleware, StateItem
+from .header_propagation import set_forwarded_headers, get_forwarded_headers, install_httpx_hook
 
 
 
@@ -24,4 +25,7 @@
     "CopilotKitMiddleware",
     "StateStreamingMiddleware",
     "StateItem",
+    "set_forwarded_headers",
+    "get_forwarded_headers",
+    "install_httpx_hook",
 ]

sdk-python/copilotkit/header_propagation.py

@@ -0,0 +1,57 @@
+"""Header propagation for forwarding x-aimock-* headers to outgoing LLM calls.
+
+Uses Python contextvars for per-request ambient state in async FastAPI handlers.
+An httpx event hook reads the ContextVar and injects headers on outgoing requests.
+"""
+
+import contextvars
+import warnings
+from typing import Dict
+
+# Ambient per-request state for headers to forward to LLM calls
+_forwarded_headers: contextvars.ContextVar[Dict[str, str]] = contextvars.ContextVar(
+    'copilotkit_forwarded_headers'
+)
+
+
+def set_forwarded_headers(headers: Dict[str, str]) -> None:
+    """Store headers to forward to outgoing LLM calls.
+    Filters to x-aimock-* headers only."""
+    filtered = {k.lower(): v for k, v in headers.items() if k.lower().startswith('x-aimock-')}
+    _forwarded_headers.set(filtered)
+
+
+def get_forwarded_headers() -> Dict[str, str]:
+    """Get headers that should be forwarded to outgoing LLM calls."""
+    return _forwarded_headers.get({})
+
+
+def install_httpx_hook(client) -> None:
+    """Append an event hook to an httpx client that injects forwarded headers.
+
+    Works with OpenAI and Anthropic Python SDKs (both use httpx internally).
+    No-op when no headers are set (demo traffic).
+
+    Parameters
+    ----------
+    client : object
+        An OpenAI/Anthropic client instance, or a raw httpx.Client/AsyncClient.
+        For SDK clients the underlying transport lives at ``client._client``.
+    """
+    def _inject_headers(request):
+        headers = get_forwarded_headers()
+        for key, value in headers.items():
+            request.headers[key] = value
+
+    # OpenAI / Anthropic SDKs wrap an httpx client at client._client
+    if hasattr(client, '_client') and hasattr(client._client, 'event_hooks'):
+        client._client.event_hooks['request'].append(_inject_headers)
+    elif hasattr(client, 'event_hooks'):
+        # Raw httpx.Client / httpx.AsyncClient
+        client.event_hooks['request'].append(_inject_headers)
+    else:
+        warnings.warn(
+            f"install_httpx_hook: client of type {type(client).__name__} has no "
+            "recognized event_hooks attribute; x-aimock-* headers will not be forwarded",
+            stacklevel=2,
+        )

sdk-python/copilotkit/integrations/fastapi.py

@@ -20,6 +20,7 @@
 )
 from ..action import ActionDict
 from ..html import generate_info_html
+from ..header_propagation import set_forwarded_headers
 logging.basicConfig(level=logging.ERROR)
 logger = logging.getLogger(__name__)
 
@@ -79,6 +80,9 @@ async def handler(request: Request, sdk: CopilotKitRemoteEndpoint):
     except: # pylint: disable=bare-except
         body = None
 
+    # Propagate x-aimock-* headers to outgoing LLM calls via ContextVar
+    set_forwarded_headers(dict(request.headers))
+
     path = request.path_params.get('path')
     method = request.method
     context = cast(

sdk-python/tests/test_header_propagation.py

@@ -0,0 +1,201 @@
+"""Tests for header propagation (x-aimock-* headers to outgoing LLM calls)."""
+
+import contextvars
+import warnings
+
+import pytest
+
+from copilotkit.header_propagation import (
+    get_forwarded_headers,
+    install_httpx_hook,
+    set_forwarded_headers,
+)
+
+
+class TestSetForwardedHeaders:
+    """set_forwarded_headers filters to x-aimock-* only."""
+
+    def test_filters_to_aimock_headers(self):
+        set_forwarded_headers({
+            "x-aimock-strict": "true",
+            "x-aimock-session": "abc123",
+            "authorization": "Bearer token",
+            "content-type": "application/json",
+        })
+        result = get_forwarded_headers()
+        assert result == {
+            "x-aimock-strict": "true",
+            "x-aimock-session": "abc123",
+        }
+
+    def test_case_insensitive_prefix_match(self):
+        set_forwarded_headers({
+            "X-AIMock-Strict": "true",
+            "X-AIMOCK-SESSION": "xyz",
+        })
+        result = get_forwarded_headers()
+        assert result == {
+            "x-aimock-strict": "true",
+            "x-aimock-session": "xyz",
+        }
+
+    def test_empty_when_no_aimock_headers(self):
+        set_forwarded_headers({
+            "authorization": "Bearer token",
+            "content-type": "application/json",
+        })
+        result = get_forwarded_headers()
+        assert result == {}
+
+    def test_empty_input(self):
+        set_forwarded_headers({})
+        result = get_forwarded_headers()
+        assert result == {}
+
+
+class TestGetForwardedHeaders:
+    """get_forwarded_headers returns empty dict by default."""
+
+    def test_default_is_empty_dict(self):
+        # Reset to default by running in a fresh context
+        ctx = contextvars.copy_context()
+        result = ctx.run(get_forwarded_headers)
+        assert result == {}
+
+
+class TestRoundTrip:
+    """set + get round-trip."""
+
+    def test_round_trip(self):
+        headers = {"x-aimock-strict": "true", "x-aimock-foo": "bar"}
+        set_forwarded_headers(headers)
+        assert get_forwarded_headers() == headers
+
+    def test_overwrite(self):
+        set_forwarded_headers({"x-aimock-a": "1"})
+        set_forwarded_headers({"x-aimock-b": "2"})
+        assert get_forwarded_headers() == {"x-aimock-b": "2"}
+
+
+class TestInstallHttpxHook:
+    """install_httpx_hook appends to event hooks."""
+
+    def test_appends_to_raw_httpx_client(self):
+        """Mock a raw httpx client with event_hooks dict."""
+        class FakeClient:
+            def __init__(self):
+                self.event_hooks = {"request": []}
+
+        client = FakeClient()
+        install_httpx_hook(client)
+        assert len(client.event_hooks["request"]) == 1
+
+    def test_appends_to_sdk_wrapped_client(self):
+        """Mock an OpenAI/Anthropic SDK client with _client attribute."""
+        class FakeTransport:
+            def __init__(self):
+                self.event_hooks = {"request": []}
+
+        class FakeSDKClient:
+            def __init__(self):
+                self._client = FakeTransport()
+
+        client = FakeSDKClient()
+        install_httpx_hook(client)
+        assert len(client._client.event_hooks["request"]) == 1
+
+    def test_hook_injects_headers(self):
+        """The installed hook reads from ContextVar and injects headers."""
+        class FakeHeaders(dict):
+            """Dict that also supports item assignment like httpx Headers."""
+            pass
+
+        class FakeRequest:
+            def __init__(self):
+                self.headers = FakeHeaders()
+
+        class FakeClient:
+            def __init__(self):
+                self.event_hooks = {"request": []}
+
+        client = FakeClient()
+        install_httpx_hook(client)
+
+        # Set headers in ContextVar
+        set_forwarded_headers({"x-aimock-strict": "true"})
+
+        # Simulate httpx calling the hook
+        request = FakeRequest()
+        client.event_hooks["request"][0](request)
+
+        assert request.headers["x-aimock-strict"] == "true"
+
+    def test_hook_noop_when_no_headers(self):
+        """Hook is a no-op when ContextVar is empty (demo traffic)."""
+        class FakeRequest:
+            def __init__(self):
+                self.headers = {}
+
+        class FakeClient:
+            def __init__(self):
+                self.event_hooks = {"request": []}
+
+        # Reset ContextVar to simulate a fresh request with no aimock headers
+        set_forwarded_headers({})
+
+        client = FakeClient()
+        install_httpx_hook(client)
+        request = FakeRequest()
+        client.event_hooks["request"][0](request)
+        assert request.headers == {}
+
+    def test_no_event_hooks_warns(self):
+        """Client without event_hooks emits a warning."""
+        with warnings.catch_warnings(record=True) as w:
+            warnings.simplefilter("always")
+            install_httpx_hook(object())
+            assert len(w) == 1
+            assert "event_hooks" in str(w[0].message)
+
+
+class TestContextVarIsolation:
+    """ContextVar provides proper isolation across contexts."""
+
+    def test_context_isolation(self):
+        """Headers set in one context don't leak to another."""
+        results = {}
+
+        def task_a():
+            set_forwarded_headers({"x-aimock-task": "a"})
+            results["a"] = get_forwarded_headers()
+
+        def task_b():
+            set_forwarded_headers({"x-aimock-task": "b"})
+            results["b"] = get_forwarded_headers()
+
+        # Run in separate contexts to verify isolation
+        ctx_a = contextvars.copy_context()
+        ctx_b = contextvars.copy_context()
+
+        ctx_a.run(task_a)
+        ctx_b.run(task_b)
+
+        assert results["a"] == {"x-aimock-task": "a"}
+        assert results["b"] == {"x-aimock-task": "b"}
+
+    def test_child_context_does_not_pollute_parent(self):
+        """Setting headers in a child context does not affect the parent."""
+        # Ensure clean state in a fresh context
+        def _run():
+            parent_before = get_forwarded_headers()
+
+            def child():
+                set_forwarded_headers({"x-aimock-child": "yes"})
+
+            ctx = contextvars.copy_context()
+            ctx.run(child)
+
+            parent_after = get_forwarded_headers()
+            assert parent_before == parent_after
+
+        contextvars.copy_context().run(_run)