From afd0b3ecd4e6f66b1cab6da70778bcfebd13e887 Mon Sep 17 00:00:00 2001
From: Alex Fournier <afournier@nvidia.com>
Date: Wed, 22 Jul 2026 14:35:55 -0700
Subject: [PATCH] fix(observability): keep Relay session headers local

Signed-off-by: Alex Fournier <afournier@nvidia.com>
---
 agent/relay_llm.py            | 11 ++++++++++-
 tests/agent/test_relay_llm.py | 31 ++++++++++++++++++++++++++++++-
 2 files changed, 40 insertions(+), 2 deletions(-)

diff --git a/agent/relay_llm.py b/agent/relay_llm.py
index 1dc80b0ba11c..142ba070b186 100644
--- a/agent/relay_llm.py
+++ b/agent/relay_llm.py
@@ -18,6 +18,9 @@ logger = logging.getLogger(__name__)
 _PROVIDER_MESSAGE_EXTENSION_KEYS = frozenset(
     {"reasoning_content", "reasoning_details"}
 )
+_RELAY_INTERNAL_PROVIDER_HEADERS = frozenset(
+    {"x-dynamo-parent-session-id", "x-dynamo-session-id"}
+)
 
 
 def execute(
@@ -694,7 +697,13 @@ def _provider_request(
                 final[key] = value
         _restore_provider_message_extensions(original, final)
     headers = getattr(request, "headers", None)
-    if isinstance(headers, dict) and headers:
+    if isinstance(headers, dict):
+        headers = {
+            key: value
+            for key, value in headers.items()
+            if str(key).lower() not in _RELAY_INTERNAL_PROVIDER_HEADERS
+        }
+    if headers:
         final["extra_headers"] = {
             **dict(final.get("extra_headers") or {}),
             **headers,
diff --git a/tests/agent/test_relay_llm.py b/tests/agent/test_relay_llm.py
index b3b62cb7ab87..4cd9484d2bdb 100644
--- a/tests/agent/test_relay_llm.py
+++ b/tests/agent/test_relay_llm.py
@@ -101,7 +101,11 @@ def test_stream_uses_rewritten_request_and_post_intercept_chunks(relay_turn):
     )
     try:
         stream = relay_llm.stream(
-            {"model": "test-model", "messages": []},
+            {
+                "model": "test-model",
+                "messages": [],
+                "extra_headers": {"authorization": "Bearer provider-token"},
+            },
             raw_stream,
             session_id="session-1",
             name="test-provider",
@@ -127,6 +131,9 @@ def test_stream_uses_rewritten_request_and_post_intercept_chunks(relay_turn):
         relay.intercepts.deregister_llm_request("hermes-test-request")
 
     assert captured_requests[0]["temperature"] == 0.25
+    assert captured_requests[0]["extra_headers"] == {
+        "authorization": "Bearer provider-token"
+    }
     assert chunks[0].choices[0].delta.content == "HELLO"
     assert stream.output_modified is True
     assert turn.logical_llm_calls == {}
@@ -221,6 +228,28 @@ def test_non_stream_preserves_raw_provider_response_identity(relay_turn):
     assert result is raw_response
 
 
+def test_non_stream_does_not_forward_relay_session_headers(relay_turn):
+    _relay, _turn = relay_turn
+    captured_requests = []
+
+    relay_llm.execute(
+        {
+            "model": "test-model",
+            "messages": [],
+            "extra_headers": {"x-provider-header": "provider-value"},
+        },
+        lambda request: captured_requests.append(request) or {"content": "ok"},
+        session_id="session-1",
+        name="test-provider",
+        model_name="test-model",
+        metadata={"api_mode": "custom", "api_request_id": "request-headers"},
+    )
+
+    assert captured_requests[0]["extra_headers"] == {
+        "x-provider-header": "provider-value"
+    }
+
+
 def test_non_stream_defers_logical_success_and_reuses_scope_for_retry(relay_turn):
     _relay, turn = relay_turn
     metadata = {"api_mode": "custom", "api_request_id": "request-retry"}