From 5e34fa2d5c32a4c1c45d658d7bb6039acd018cc9 Mon Sep 17 00:00:00 2001
From: Alex Fournier <afournier@nvidia.com>
Date: Tue, 28 Jul 2026 09:58:43 -0700
Subject: [PATCH] fix(relay): preserve provider stream errors

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

diff --git a/agent/relay_llm.py b/agent/relay_llm.py
index 4508fce94144..3481ca072376 100644
--- a/agent/relay_llm.py
+++ b/agent/relay_llm.py
@@ -438,6 +438,11 @@ class ManagedLlmStream(Iterator[Any]):
                 run_callback(self._on_chunk, _jsonable(chunk))
 
         def relay_finalizer() -> Any:
+            # Relay can invoke the finalizer while unwinding a provider-stream
+            # failure. Preserve that original callback error instead of
+            # replacing it with a secondary "missing terminal response" error.
+            if self._callback_error is not None:
+                return None
             try:
                 if self.final_response is not None:
                     return _jsonable(self.final_response)
diff --git a/tests/agent/test_relay_llm.py b/tests/agent/test_relay_llm.py
index ae4eac15826a..2e09923f327e 100644
--- a/tests/agent/test_relay_llm.py
+++ b/tests/agent/test_relay_llm.py
@@ -181,6 +181,49 @@ def test_deferred_stream_preserves_provider_error_and_logical_scope_for_retry(
     assert "request-2" in turn.logical_llm_calls
 
 
+def test_stream_provider_error_is_not_replaced_by_finalizer_error(relay_turn):
+    _relay, turn = relay_turn
+
+    class ProviderError(Exception):
+        pass
+
+    provider_error = ProviderError("provider failed before first chunk")
+    finalizer_called = False
+
+    def failing_stream(_request):
+        def generate():
+            raise provider_error
+            yield  # pragma: no cover
+
+        return generate()
+
+    def failing_finalizer():
+        nonlocal finalizer_called
+        finalizer_called = True
+        raise RuntimeError("missing terminal response")
+
+    stream = relay_llm.stream(
+        {"model": "test-model", "input": "hi"},
+        failing_stream,
+        session_id="session-1",
+        name="test-provider",
+        model_name="test-model",
+        finalizer=failing_finalizer,
+        metadata={
+            "api_mode": "codex_responses",
+            "api_request_id": "request-provider-before-finalizer",
+        },
+        defer_logical_completion=True,
+    )
+
+    with pytest.raises(ProviderError) as caught:
+        list(stream)
+
+    assert caught.value is provider_error
+    assert finalizer_called is False
+    assert "request-provider-before-finalizer" in turn.logical_llm_calls
+
+
 def test_non_deferred_partial_stream_close_cancels_logical_call(
     relay_turn,
     monkeypatch,