fix(sdk): honor API key header for observability

abhinav-galileo · abhinav-galileo · commit 7ac987cee81e · 2026-05-12T15:20:04.000+05:30
diff --git a/sdks/python/src/agent_control/observability.py b/sdks/python/src/agent_control/observability.py
@@ -26,6 +26,9 @@
     await shutdown_observability()
 
 Configuration (Environment Variables):
+    # Server connection
+    AGENT_CONTROL_API_KEY_HEADER: API key header name (default: X-API-Key)
+
     # Observability (event batching)
     AGENT_CONTROL_OBSERVABILITY_ENABLED: Enable observability (default: true)
     AGENT_CONTROL_OBSERVABILITY_SINK_NAME: Selected control-event sink (default: default)
@@ -286,6 +289,7 @@ class EventBatcher:
     Attributes:
         server_url: Base URL of the Agent Control server
         api_key: API key for authentication
+        api_key_header: HTTP header used to send the API key
         batch_size: Maximum events per batch
         flush_interval: Seconds between automatic flushes
     """
@@ -294,6 +298,7 @@ def __init__(
         self,
         server_url: str | None = None,
         api_key: str | None = None,
+        api_key_header: str | None = None,
         batch_size: int | None = None,
         flush_interval: float | None = None,
     ):
@@ -303,11 +308,13 @@ def __init__(
         Args:
             server_url: Server URL (defaults to get_settings().url)
             api_key: API key (defaults to get_settings().api_key)
+            api_key_header: API key header (defaults to get_settings().api_key_header)
             batch_size: Max events per batch (defaults to get_settings().batch_size)
             flush_interval: Seconds between flushes (defaults to get_settings().flush_interval)
         """
         self.server_url = server_url or get_settings().url
         self.api_key = api_key or get_settings().api_key
+        self.api_key_header = api_key_header or get_settings().api_key_header
         self.batch_size = batch_size if batch_size is not None else get_settings().batch_size
         if flush_interval is not None:
             self.flush_interval = flush_interval
@@ -424,7 +431,7 @@ def _build_batch_request(
         url = f"{self.server_url}/api/v1/observability/events"
         headers = {"Content-Type": "application/json"}
         if self.api_key:
-            headers["X-API-Key"] = self.api_key
+            headers[self.api_key_header] = self.api_key
         payload = {"events": [event.model_dump(mode="json") for event in events]}
         return url, headers, payload
 
@@ -1098,6 +1105,7 @@ def _get_custom_control_event_sinks_to_shutdown() -> tuple[ControlEventSink, ...
 def init_observability(
     server_url: str | None = None,
     api_key: str | None = None,
+    api_key_header: str | None = None,
     enabled: bool | None = None,
     sink_name: str | None = None,
     sink_config: JSONObject | None = None,
@@ -1110,6 +1118,7 @@ def init_observability(
     Args:
         server_url: Server URL for sending events
         api_key: API key for authentication
+        api_key_header: HTTP header used to send the API key
         enabled: Override AGENT_CONTROL_OBSERVABILITY_ENABLED
         sink_name: Override AGENT_CONTROL_OBSERVABILITY_SINK_NAME
         sink_config: Override AGENT_CONTROL_OBSERVABILITY_SINK_CONFIG
@@ -1157,7 +1166,11 @@ def init_observability(
         return _batcher
 
     # Create batcher
-    _batcher = EventBatcher(server_url=server_url, api_key=api_key)
+    _batcher = EventBatcher(
+        server_url=server_url,
+        api_key=api_key,
+        api_key_header=api_key_header,
+    )
     _batcher.start()
     _event_sink = _BatcherControlEventSink(_batcher)
 
diff --git a/sdks/python/src/agent_control/settings.py b/sdks/python/src/agent_control/settings.py
@@ -56,6 +56,10 @@ class SDKSettings(BaseSettings):
         default="",
         description="API key for server authentication",
     )
+    api_key_header: str = Field(
+        default="X-API-Key",
+        description="HTTP header used to send the API key",
+    )
 
     # Observability (event batching)
     observability_enabled: bool = Field(
diff --git a/sdks/python/tests/test_observability.py b/sdks/python/tests/test_observability.py
@@ -120,6 +120,7 @@ def reset_observability_state() -> None:
         observability_enabled=True,
         observability_sink_name=DEFAULT_CONTROL_EVENT_SINK_NAME,
         observability_sink_config={},
+        api_key_header="X-API-Key",
     )
     with obs._used_custom_event_sinks_lock:
         obs._used_custom_event_sinks.clear()
@@ -136,6 +137,7 @@ class TestEventBatcherInit:
     def test_init_default_values(self):
         """Test EventBatcher initializes with default values."""
         batcher = EventBatcher()
+        assert batcher.api_key_header == get_settings().api_key_header
         assert batcher.batch_size == get_settings().batch_size
         assert batcher.flush_interval == get_settings().flush_interval
         assert batcher.shutdown_join_timeout == get_settings().shutdown_join_timeout
@@ -149,32 +151,37 @@ def test_init_custom_values(self):
         batcher = EventBatcher(
             server_url="http://custom:9000",
             api_key="test-key",
+            api_key_header="X-Custom-API-Key",
             batch_size=50,
             flush_interval=5.0,
         )
         assert batcher.server_url == "http://custom:9000"
         assert batcher.api_key == "test-key"
+        assert batcher.api_key_header == "X-Custom-API-Key"
         assert batcher.batch_size == 50
         assert batcher.flush_interval == 5.0
 
     def test_init_from_settings(self):
         """Test EventBatcher reads from settings."""
         from agent_control.settings import configure_settings
 
-        # Save original values
-        original_url = get_settings().url
-        original_api_key = get_settings().api_key
+        original_settings = get_settings().model_dump()
 
         try:
             # Configure settings programmatically
-            configure_settings(url="http://configured-server:8080", api_key="configured-api-key")
+            configure_settings(
+                url="http://configured-server:8080",
+                api_key="configured-api-key",
+                api_key_header="X-Custom-API-Key",
+            )
 
             batcher = EventBatcher()
             assert batcher.server_url == "http://configured-server:8080"
             assert batcher.api_key == "configured-api-key"
+            assert batcher.api_key_header == "X-Custom-API-Key"
         finally:
             # Restore original settings
-            configure_settings(url=original_url, api_key=original_api_key)
+            configure_settings(**original_settings)
 
 
 class TestEventBatcherStartStop:
@@ -557,6 +564,46 @@ def test_send_batch_sync_returns_true_on_202(self):
         assert result is True
         client_ctor.assert_called_once_with(timeout=30.0)
         client.post.assert_called_once()
+        assert client.post.call_args.kwargs["headers"]["X-API-Key"] == "test-key"
+
+    def test_send_batch_sync_uses_configured_api_key_header(self):
+        batcher = EventBatcher(
+            server_url="http://test:8000",
+            api_key="test-key",
+            api_key_header="X-Custom-API-Key",
+        )
+        response = MagicMock(status_code=202, text="accepted")
+        client = MagicMock()
+        client.post.return_value = response
+        client_context = MagicMock()
+        client_context.__enter__.return_value = client
+
+        with patch(
+            "agent_control.observability.httpx.Client",
+            return_value=client_context,
+        ):
+            result = batcher._send_batch_sync([create_mock_event()])
+
+        assert result is True
+        headers = client.post.call_args.kwargs["headers"]
+        assert headers["X-Custom-API-Key"] == "test-key"
+        assert "X-API-Key" not in headers
+
+    def test_build_batch_request_uses_settings_api_key_header(self):
+        original_settings = get_settings().model_dump()
+        try:
+            configure_settings(
+                api_key="settings-key",
+                api_key_header="X-Custom-API-Key",
+            )
+            batcher = EventBatcher()
+
+            _, headers, _ = batcher._build_batch_request([create_mock_event()])
+
+            assert headers["X-Custom-API-Key"] == "settings-key"
+            assert "X-API-Key" not in headers
+        finally:
+            configure_settings(**original_settings)
 
     def test_send_batch_sync_returns_false_on_401_without_retry(self):
         batcher = EventBatcher()
@@ -1069,10 +1116,12 @@ def test_init_enabled_creates_batcher(self):
             result = init_observability(
                 server_url="http://test:8000",
                 api_key="test-key",
+                api_key_header="X-Custom-API-Key",
                 enabled=True,
             )
             assert result is not None
             assert isinstance(result, EventBatcher)
+            assert result.api_key_header == "X-Custom-API-Key"
             assert result._running is True
             assert get_event_sink() is not None
 
