fix(observability): keep Relay session headers local

Signed-off-by: Alex Fournier <afournier@nvidia.com>
This commit is contained in:
Alex Fournier 2026-07-22 14:35:55 -07:00
parent 4add7d0c12
commit afd0b3ecd4
2 changed files with 40 additions and 2 deletions

View file

@ -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,

View file

@ -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"}