From 998a999ce8cf31e15d6e8a32e69c95d7660b2010 Mon Sep 17 00:00:00 2001
From: Radu Swigler <radu.porumba@gmail.com>
Date: Thu, 3 Sep 2026 03:34:35 -0400
Subject: [PATCH] fix: stream errors to client before breaking generator loop
 (#397)

Since v1.7.7 (PR #384), when a generator handler yields an error the
code breaks out of the streaming loop before calling stream_result(),
so the error is never delivered to clients listening on /stream. The
error is saved for /status but the stream goes silent.

Additionally, the error extraction stripped all extra fields (context,
refresh_worker, etc.) by replacing the dict with {"error": msg}.

- Call stream_result() with the error BEFORE breaking out of the loop
- Preserve extra fields from the error dict alongside the error message
- Add tests for both streaming delivery and field preservation

Closes #397
---
 runpod/serverless/modules/rp_job.py           | 11 ++-
 .../test_serverless/test_modules/test_job.py  | 68 +++++++++++++++++++
 2 files changed, 78 insertions(+), 1 deletion(-)

diff --git a/runpod/serverless/modules/rp_job.py b/runpod/serverless/modules/rp_job.py
index a45cebc68..01613a2e9 100644
--- a/runpod/serverless/modules/rp_job.py
+++ b/runpod/serverless/modules/rp_job.py
@@ -204,9 +204,18 @@ async def handle_job(session: ClientSession, config: Dict[str, Any], job) -> dic
 
             if isinstance(stream_output.get("output"), dict):
                 if stream_output["output"].get("error"):
-                    stream_output = {"error": str(stream_output["output"]["error"])}
+                    error_fields = {
+                        k: v
+                        for k, v in stream_output["output"].items()
+                        if k != "error"
+                    }
+                    stream_output = {
+                        "error": str(stream_output["output"]["error"]),
+                        **error_fields,
+                    }
 
             if stream_output.get("error"):
+                await stream_result(session, stream_output, job)
                 job_result = stream_output
                 break
 
diff --git a/tests/test_serverless/test_modules/test_job.py b/tests/test_serverless/test_modules/test_job.py
index 1ec5ce353..0aef10393 100644
--- a/tests/test_serverless/test_modules/test_job.py
+++ b/tests/test_serverless/test_modules/test_job.py
@@ -440,3 +440,71 @@ async def test_run_job_generator_exception(self):
         assert mock_log.info.call_count == 1
         mock_log.info.assert_called_with("Finished running generator.", "123")
 
+
+class TestHandleJobStreamError(IsolatedAsyncioTestCase):
+    """Tests that generator errors are streamed to the client. (Fixes #397)"""
+
+    @staticmethod
+    async def _noop(*args, **kwargs):
+        pass
+
+    async def test_error_is_streamed_before_break(self):
+        """When a generator yields an error, stream_result must be called with it."""
+
+        def error_gen(job):
+            yield "chunk_1"
+            yield {"error": "GPU OOM"}
+
+        config = {"handler": error_gen}
+        job = {"id": "test-123"}
+        session = Mock()
+
+        stream_calls = []
+
+        async def capture_stream(*args, **kwargs):
+            stream_calls.append(args)
+
+        with patch(
+            "runpod.serverless.modules.rp_job.stream_result", side_effect=capture_stream,
+        ), patch(
+            "runpod.serverless.modules.rp_job.send_result", side_effect=self._noop,
+        ), patch(
+            "runpod.serverless.modules.rp_job.log", new_callable=Mock
+        ):
+            await rp_job.handle_job(session, config, job)
+
+        # stream_result should be called twice: once for "chunk_1", once for the error
+        assert len(stream_calls) == 2
+        streamed_payload = stream_calls[1][1]  # second call, second positional arg
+        assert "error" in streamed_payload
+
+    async def test_nested_error_preserves_extra_fields(self):
+        """Extra fields like refresh_worker should survive error extraction."""
+
+        def error_gen(job):
+            yield {"error": "bad things", "context": "step 47", "refresh_worker": True}
+
+        config = {"handler": error_gen}
+        job = {"id": "test-456"}
+        session = Mock()
+
+        stream_calls = []
+
+        async def capture_stream(*args, **kwargs):
+            stream_calls.append(args)
+
+        with patch(
+            "runpod.serverless.modules.rp_job.stream_result", side_effect=capture_stream,
+        ), patch(
+            "runpod.serverless.modules.rp_job.send_result", side_effect=self._noop,
+        ), patch(
+            "runpod.serverless.modules.rp_job.log", new_callable=Mock
+        ):
+            await rp_job.handle_job(session, config, job)
+
+        assert len(stream_calls) == 1
+        streamed_payload = stream_calls[0][1]
+        assert streamed_payload["error"] == "bad things"
+        assert streamed_payload["context"] == "step 47"
+        assert streamed_payload["refresh_worker"] is True
+