fix(runtime): preserve native streaming responses

Signed-off-by: Alex Fournier <afournier@nvidia.com>

agent/chat_completion_helpers.py

@@ -3282,6 +3282,8 @@ def interruptible_streaming_api_call(agent, api_kwargs: dict, *, on_first_delta=
 
         if agent._interrupt_requested:
             return None
+        if base_final_message is not None and not stream.output_modified:
+            return base_final_message
         final_message = accumulator.response(base_final_message)
         if (
             not getattr(final_message, "content", None)

agent/relay_llm.py

@@ -296,6 +296,7 @@ class ManagedLlmStream(Iterator[Any]):
         self._accept_chunk = accept_chunk
         self._relay_observes_chunks = False
         self._raw_chunks: list[tuple[Any, Any]] = []
+        self.output_modified = False
 
         runtime, session, parent = relay_runtime.resolve_execution_context(session_id)
         if runtime is None or session is None:
@@ -405,6 +406,8 @@ class ManagedLlmStream(Iterator[Any]):
         try:
             chunk = self._loop.run_until_complete(next_chunk())
         except StopAsyncIteration:
+            if self._raw_chunks:
+                self.output_modified = True
             if not self._defer_logical_completion:
                 _complete_logical(self._logical, outcome="success")
                 self._logical = None
@@ -423,8 +426,11 @@ class ManagedLlmStream(Iterator[Any]):
             self._on_chunk(chunk)
         for index, (encoded, raw) in enumerate(self._raw_chunks):
             if _json_equal(chunk, encoded):
+                if index > 0:
+                    self.output_modified = True
                 del self._raw_chunks[: index + 1]
                 return raw
+        self.output_modified = True
         return self._chunk_adapter(chunk)
 
     def close(self) -> None:

tests/agent/test_relay_llm.py

@@ -125,6 +125,7 @@ def test_stream_uses_rewritten_request_and_post_intercept_chunks(relay_turn):
 
     assert captured_requests[0]["temperature"] == 0.25
     assert chunks[0].choices[0].delta.content == "HELLO"
+    assert stream.output_modified is True
     assert turn.logical_llm_calls == {}
 
 
@@ -302,6 +303,7 @@ def test_stream_defers_logical_success_for_response_validation(relay_turn):
     )
 
     assert list(stream) == [{"delta": "pending-validation"}]
+    assert stream.output_modified is False
     assert "request-stream-defer" in turn.logical_llm_calls
 
     relay_llm.complete_logical_call("request-stream-defer", outcome="success")