fix: stream errors to client before breaking generator loop

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
 

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
+