Python: fix(core): ResponseStream Transform Hook & Mapper Error Boundary

python/packages/core/agent_framework/_types.py

@@ -3504,10 +3504,10 @@ async def _record_update(self, update: UpdateT) -> UpdateT:
 
     async def __anext__(self) -> UpdateT:
         while True:
-            if self._pending_mapped_updates:
-                return await self._record_update(self._pending_mapped_updates.pop(0))
-
             try:
+                if self._pending_mapped_updates:
+                    return await self._record_update(self._pending_mapped_updates.pop(0))
+
                 with contextlib.ExitStack() as stack:
                     for factory in self._pull_context_manager_factories:
                         stack.enter_context(factory())
@@ -3519,6 +3519,18 @@ async def __anext__(self) -> UpdateT:
                         stream = await self._get_stream()
                         self._iterator = stream.__aiter__()
                     update: UpdateT = await self._iterator.__anext__()
+
+                if self._flat_map_update is not None:
+                    mapped_updates = self._flat_map_update(update)
+                    if isawaitable(mapped_updates):
+                        mapped_updates = await mapped_updates
+                    self._pending_mapped_updates.extend(mapped_updates)
+                    continue
+                if self._map_update is not None:
+                    update = self._map_update(update)  # type: ignore[assignment]
+                    if isawaitable(update):
+                        update = await update
+                return await self._record_update(update)
             except StopAsyncIteration:
                 self._consumed = True
                 await self._run_cleanup_hooks()
@@ -3531,17 +3543,6 @@ async def __anext__(self) -> UpdateT:
                 finally:
                     self._stream_error = None
                 raise
-            if self._flat_map_update is not None:
-                mapped_updates = self._flat_map_update(update)
-                if isawaitable(mapped_updates):
-                    mapped_updates = await mapped_updates
-                self._pending_mapped_updates.extend(mapped_updates)
-                continue
-            if self._map_update is not None:
-                update = self._map_update(update)  # type: ignore[assignment]
-                if isawaitable(update):
-                    update = await update
-            return await self._record_update(update)
 
     async def close(self) -> None:
         """Close the active iterator and run cleanup hooks.

python/packages/core/tests/core/test_types.py

@@ -4450,6 +4450,48 @@ async def async_cleanup() -> None:
 
         assert cleanup_called["value"] is True
 
+    async def test_transform_hook_raising_runs_cleanup(self) -> None:
+        """Cleanup hook is called when a transform hook raises mid-stream."""
+        cleanup_called = {"value": False}
+
+        def cleanup_hook() -> None:
+            cleanup_called["value"] = True
+
+        def failing_hook(update: ChatResponseUpdate) -> ChatResponseUpdate:
+            raise RuntimeError("hook error")
+
+        stream = ResponseStream(
+            _generate_updates(1),
+            finalizer=_combine_updates,
+            transform_hooks=[failing_hook],  # type: ignore[arg-type]
+            cleanup_hooks=[cleanup_hook],
+        )
+
+        with pytest.raises(RuntimeError, match="hook error"):
+            async for _ in stream:
+                pass
+
+        assert cleanup_called["value"] is True
+
+    async def test_mapper_raising_runs_cleanup(self) -> None:
+        """Cleanup hook on a mapped stream is called when the mapper raises mid-stream."""
+        cleanup_called = {"value": False}
+
+        def cleanup_hook() -> None:
+            cleanup_called["value"] = True
+
+        def failing_mapper(update: ChatResponseUpdate) -> ChatResponseUpdate:
+            raise ValueError("mapper error")
+
+        inner = ResponseStream(_generate_updates(1), finalizer=_combine_updates)
+        outer = inner.map(failing_mapper, _combine_updates).with_cleanup_hook(cleanup_hook)
+
+        with pytest.raises(ValueError, match="mapper error"):
+            async for _ in outer:
+                pass
+
+        assert cleanup_called["value"] is True
+
 
 class TestResponseStreamResultHooks:
     """Tests for result hooks (after finalizer)."""