fix(sdk): propagate custom API key header

Author: abhinav-galileo

sdks/python/src/agent_control/__init__.py

@@ -159,6 +159,7 @@ class _RefreshContext:
     agent_name: str
     server_url: str
     api_key: str | None
+    api_key_header: str | None
     target_type: str | None
     target_id: str | None
 
@@ -220,6 +221,7 @@ def _snapshot_refresh_context() -> _RefreshContext:
         agent = state.current_agent
         server_url = state.server_url
         api_key = state.api_key
+        api_key_header = state.api_key_header
         target_type = state.target_type
         target_id = state.target_id
 
@@ -233,6 +235,7 @@ def _snapshot_refresh_context() -> _RefreshContext:
         agent_name=agent.agent_name,
         server_url=server_url,
         api_key=api_key,
+        api_key_header=api_key_header,
         target_type=target_type,
         target_id=target_id,
     )
@@ -243,6 +246,7 @@ async def _fetch_controls_for_context_async(context: _RefreshContext) -> list[di
     async with AgentControlClient(
         base_url=context.server_url,
         api_key=context.api_key,
+        api_key_header=context.api_key_header,
     ) as client:
         response = await agents.list_agent_controls(
             client,
@@ -429,6 +433,7 @@ def init(
     agent_version: str | None = None,
     server_url: str | None = None,
     api_key: str | None = None,
+    api_key_header: str | None = None,
     controls_file: str | None = None,
     steps: list[StepSchemaDict] | None = None,
     conflict_mode: Literal["strict", "overwrite"] = "overwrite",
@@ -467,6 +472,8 @@ def init(
         server_url: Optional server URL (defaults to AGENT_CONTROL_URL env var
                    or http://localhost:8000)
         api_key: Optional API key for authentication (defaults to AGENT_CONTROL_API_KEY env var)
+        api_key_header: Optional HTTP header name for API key authentication
+            (defaults to AGENT_CONTROL_API_KEY_HEADER env var or X-API-Key)
         controls_file: Optional explicit path to controls.yaml (auto-discovered if not provided)
         steps: Optional list of step schemas for registration:
                [{"type": "tool", "name": "search", "input_schema": {...}, "output_schema": {...}}]
@@ -561,6 +568,7 @@ async def handle(message: str):
         state.current_agent = next_agent
         state.server_url = server_url or os.getenv('AGENT_CONTROL_URL') or 'http://localhost:8000'
         state.api_key = api_key
+        state.api_key_header = api_key_header
         state.runtime_token_cache.clear()
         state.target_type = target_type
         state.target_id = target_id
@@ -599,6 +607,7 @@ async def register() -> list[dict[str, Any]] | None:
             async with AgentControlClient(
                 base_url=state.server_url,
                 api_key=state.api_key,
+                api_key_header=state.api_key_header,
             ) as client:
                 # Check server health first
                 try:
@@ -685,6 +694,7 @@ def run_in_thread() -> None:
     batcher = init_observability(
         server_url=state.server_url,
         api_key=state.api_key,
+        api_key_header=state.api_key_header,
         enabled=observability_enabled,
         sink_name=observability_sink_name,
         sink_config=observability_sink_config,
@@ -716,6 +726,7 @@ def _reset_state() -> None:
         state.server_controls = None
         state.server_url = None
         state.api_key = None
+        state.api_key_header = None
         state.runtime_token_cache.clear()
         state.target_type = None
         state.target_id = None

sdks/python/src/agent_control/_state.py

@@ -26,6 +26,7 @@ def __init__(self) -> None:
         self.server_controls: list[dict[str, Any]] | None = None
         self.server_url: str | None = None
         self.api_key: str | None = None
+        self.api_key_header: str | None = None
         self.runtime_token_cache = RuntimeTokenCache()
         # Optional target context fixed at init() time; both fields are set
         # together or both remain None.

sdks/python/src/agent_control/evaluation.py

@@ -557,6 +557,7 @@ async def evaluate_controls(
     async with AgentControlClient(
         base_url=state.server_url,
         api_key=state.api_key,
+        api_key_header=state.api_key_header,
         runtime_token_cache=state.runtime_token_cache,
     ) as client:
         return await check_evaluation_with_local(

sdks/python/src/agent_control/integrations/google_adk/plugin.py

@@ -853,7 +853,11 @@ async def _sync_steps_async(self, steps: list[StepSchemaDict]) -> None:
                 "with the same agent_name as AgentControlPlugin."
             )
 
-        async with AgentControlClient(base_url=state.server_url, api_key=state.api_key) as client:
+        async with AgentControlClient(
+            base_url=state.server_url,
+            api_key=state.api_key,
+            api_key_header=state.api_key_header,
+        ) as client:
             response = await agents.get_agent(client, self.agent_name)
             existing = GetAgentResponse.model_validate(response)
             existing_keys = {(step.type, step.name) for step in existing.steps}

sdks/python/tests/test_evaluation.py

@@ -4,10 +4,9 @@
 from uuid import UUID
 
 import pytest
-from pydantic import ValidationError
-
 from agent_control import evaluation
 from agent_control.evaluation import EvaluationResult
+from pydantic import ValidationError
 
 
 @pytest.mark.asyncio
@@ -126,6 +125,27 @@ async def test_evaluate_controls_with_context(monkeypatch):
     assert mock_check.call_args is not None
 
 
+@pytest.mark.asyncio
+async def test_evaluate_controls_uses_session_api_key_header(monkeypatch):
+    """evaluate_controls should pass init's API-key header into the client."""
+    mock_result = EvaluationResult(is_safe=True, confidence=1.0)
+    mock_check = AsyncMock(return_value=mock_result)
+    monkeypatch.setattr(evaluation, "check_evaluation_with_local", mock_check)
+
+    with patch("agent_control.state.server_url", "http://localhost:8000"), patch(
+        "agent_control.state.api_key", "test-key"
+    ), patch("agent_control.state.api_key_header", "Galileo-API-Key"):
+        await evaluation.evaluate_controls(
+            step_name="chat",
+            input="hello",
+            stage="pre",
+            agent_name="test-bot",
+        )
+
+    client = mock_check.call_args.kwargs["client"]
+    assert client.api_key_header == "Galileo-API-Key"
+
+
 @pytest.mark.asyncio
 async def test_check_evaluation_forwards_target_context():
     """When target_type and target_id are supplied, they are forwarded to the server."""

sdks/python/tests/test_init_step_merge.py

@@ -4,7 +4,7 @@
 
 import logging
 from collections.abc import Generator
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 from unittest.mock import ANY, AsyncMock, patch
 from uuid import uuid4
 
@@ -219,6 +219,47 @@ def test_init_omits_merge_events_from_public_signature() -> None:
     assert "merge_events" not in signature.parameters
 
 
+def test_init_passes_api_key_header_to_client_and_state() -> None:
+    register_agent_mock = AsyncMock(return_value={"created": True, "controls": []})
+    health_check_mock = AsyncMock(return_value={"status": "healthy"})
+    client_init_kwargs: list[dict[str, Any]] = []
+    original_init = agent_control.AgentControlClient.__init__
+
+    def recording_init(
+        self: agent_control.AgentControlClient,
+        *args: Any,
+        **kwargs: Any,
+    ) -> None:
+        client_init_kwargs.append(dict(kwargs))
+        original_init(self, *args, **kwargs)
+
+    with patch.object(
+        agent_control.AgentControlClient,
+        "__init__",
+        new=recording_init,
+    ), patch(
+        "agent_control.__init__.AgentControlClient.health_check",
+        new=health_check_mock,
+    ), patch(
+        "agent_control.__init__.agents.register_agent",
+        new=register_agent_mock,
+    ), patch.object(
+        agent_control,
+        "init_observability",
+        return_value=None,
+    ) as observability_mock:
+        agent_control.init(
+            agent_name=f"agent-{uuid4().hex[:12]}",
+            api_key="test-key",
+            api_key_header="Galileo-API-Key",
+            policy_refresh_interval_seconds=0,
+        )
+
+    assert agent_control.state.api_key_header == "Galileo-API-Key"
+    assert client_init_kwargs[0]["api_key_header"] == "Galileo-API-Key"
+    assert observability_mock.call_args.kwargs["api_key_header"] == "Galileo-API-Key"
+
+
 @pytest.mark.asyncio
 async def test_refresh_controls_calls_agent_controls_endpoint() -> None:
     # Given: an initialized SDK agent session with network-facing calls mocked.
@@ -238,6 +279,7 @@ async def test_refresh_controls_calls_agent_controls_endpoint() -> None:
     ):
         agent_control.init(
             agent_name=f"agent-{uuid4().hex[:12]}",
+            api_key_header="Galileo-API-Key",
             policy_refresh_interval_seconds=0,
         )
 
@@ -255,4 +297,5 @@ async def test_refresh_controls_calls_agent_controls_endpoint() -> None:
         target_type=None,
         target_id=None,
     )
+    assert agent_control.state.api_key_header == "Galileo-API-Key"
     assert register_agent_mock.await_count == 0

sdks/python/tests/test_shutdown.py

@@ -9,10 +9,9 @@
 from typing import Any
 from unittest.mock import AsyncMock, MagicMock, patch
 
-import pytest
-
 import agent_control
 import agent_control.observability as obs_mod
+import pytest
 from agent_control._state import state
 from agent_control.observability import EventBatcher
 
@@ -64,6 +63,7 @@ def test_shutdown_resets_state(self):
         state.server_controls = [{"name": "test"}]
         state.server_url = "http://localhost:8000"
         state.api_key = "key"
+        state.api_key_header = "X-Custom-API-Key"
 
         agent_control.shutdown()
 
@@ -72,6 +72,7 @@ def test_shutdown_resets_state(self):
         assert state.server_controls is None
         assert state.server_url is None
         assert state.api_key is None
+        assert state.api_key_header is None
 
     def test_shutdown_idempotent(self):
         agent_control.shutdown()