From 7ee68cca4553fa4b947ea9c3bdbef2f8d48a2f65 Mon Sep 17 00:00:00 2001
From: Alex Fournier <afournier@nvidia.com>
Date: Mon, 17 Aug 2026 10:21:32 -0700
Subject: [PATCH] test(relay): cover lazy completion with interceptor

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

diff --git a/agent/relay_llm.py b/agent/relay_llm.py
index 6a5669d6e5..270960fd26 100644
--- a/agent/relay_llm.py
+++ b/agent/relay_llm.py
@@ -287,10 +287,11 @@ def stream_current(
     without the unwrap the response stays trapped as ``final_response`` on the
     inner ManagedLlmStream and the outer consumer sees an empty stream.
 
-    Determining that return shape requires starting a lazy managed pipeline.
-    The first pull therefore happens before this function returns: a genuine
-    first chunk is buffered, while provider latency and errors surface at the
-    same call boundary as the unmanaged ``stream_factory(request)`` path.
+    Determining that return shape requires starting the lazy managed pipeline,
+    and Relay may read ahead internally while satisfying that first pull. A
+    genuine first returned chunk remains buffered, while provider work,
+    latency, and pre-first-yield errors may surface before this function
+    returns.
     """
     turn = relay_runtime.active_turn()
     if turn is None:
diff --git a/tests/agent/test_relay_llm.py b/tests/agent/test_relay_llm.py
index 836eeedcfb..e28f28ea67 100644
--- a/tests/agent/test_relay_llm.py
+++ b/tests/agent/test_relay_llm.py
@@ -822,6 +822,36 @@ def test_stream_current_primes_lazy_completed_response(relay_turn, monkeypatch):
     assert result is completed
 
 
+def test_stream_current_unwraps_completed_response_with_real_interceptor(relay_turn):
+    """A real stream interceptor makes Relay lazy; completion still unwraps."""
+    relay, _turn = relay_turn
+    completed = _completed_response()
+
+    async def identity_stream(request, next_call):
+        return await next_call(request)
+
+    relay.intercepts.register_llm_stream_execution(
+        "hermes-test-prime-completed",
+        1,
+        identity_stream,
+    )
+    try:
+        result = relay_llm.stream_current(
+            {"model": "test-model", "messages": [], "stream": True},
+            lambda _request: completed,
+            name="test-provider",
+            model_name="test-model",
+            finalizer=lambda: completed,
+            completed_response_predicate=_choices_predicate,
+        )
+
+        assert result is completed
+    finally:
+        relay.intercepts.deregister_llm_stream_execution(
+            "hermes-test-prime-completed"
+        )
+
+
 def test_stream_current_preserves_real_relay_interceptor_chunks(relay_turn):
     """Priming a real managed pipeline must retain its transformed first chunk."""
     relay, _turn = relay_turn