From f11737e772319d2316bb6f6a8350d735e4bcabe3 Mon Sep 17 00:00:00 2001 From: "Paul S. Schweigert" Date: Wed, 4 Mar 2026 20:38:19 -0500 Subject: [PATCH 1/4] fix: do not post_process before finally in ModelOutputThunk.astream Signed-off-by: Paul S. Schweigert --- mellea/core/base.py | 146 ++++++++++++++++++++------------------------ 1 file changed, 66 insertions(+), 80 deletions(-) diff --git a/mellea/core/base.py b/mellea/core/base.py index e7b40c7cd..34d7e3b9c 100644 --- a/mellea/core/base.py +++ b/mellea/core/base.py @@ -282,88 +282,74 @@ async def astream(self) -> str: 0 if self._underlying_value is None else len(str(self._underlying_value)) ) # type: ignore - exception_to_raise = None - try: - # Type of the chunk depends on the backend. - chunks: list[Any | None] = [] - while True: - try: - item = self._async_queue.get_nowait() - chunks.append(item) - except asyncio.QueueEmpty: - # We've exhausted the current items in the queue. - break - - # Make sure we always get the minimum chunk size. - while len(chunks) <= self._chunk_size: - if len(chunks) > 0: - if chunks[-1] is None or isinstance(chunks[-1], Exception): - break # Hit sentinel value or an error. - # We could switch to relying on the `done` / `finish_reason` field of chunks, - # but that forces us to know about the chunk type here. Prefer sentinel values - # for now. - - item = await self._async_queue.get() + # Type of the chunk depends on the backend. + chunks: list[Any | None] = [] + while True: + try: + item = self._async_queue.get_nowait() chunks.append(item) - - # Process the sentinel value if it's there. - if chunks[-1] is None: - chunks.pop() # Remove the sentinel value. - do_set_computed = True - - # Shouldn't be needed, but cancel the Tasks this ModelOutputThunk relied on. - if self._generate is not None: - self._generate.cancel() - if self._generate_extra is not None: - # Covers an hf edge case. The task is done generating anything useful but isn't `done` yet. - await self._generate_extra - self._generate_extra.cancel() - - # If ModelOutputThunks get too bulky, we can do additional cleanup here - # and set fields to None. - - elif isinstance(chunks[-1], Exception): - # Mark as computed so post_process runs in finally block - self._computed = True - # Store exception to re-raise after cleanup - exception_to_raise = chunks[-1] - - for chunk in chunks: - assert self._process is not None - await self._process(self, chunk) - - if do_set_computed: - assert self._underlying_value is not None - self._computed = True - finally: - # Always call post_process if computed, even on exception - # This ensures telemetry spans are properly closed - if self._computed: - assert self._post_process is not None - await self._post_process(self) - - # Only parse if no exception occurred - if exception_to_raise is None: - match self._action: - case Component(): - self.parsed_repr = self._action._parse(self) - case CBlock(): - assert self.value is not None, ( - "value must be non-None since this thunk is computed" - ) - self.parsed_repr = self.value # type: ignore - case _: - raise ValueError( - "attempted to astream from a model output thunk with no ._action set" - ) - assert self.parsed_repr is not None, ( - "enforce constraint that a computed ModelOutputThunk has a non-None parsed_repr" + except asyncio.QueueEmpty: + # We've exhausted the current items in the queue. + break + + # Make sure we always get the minimum chunk size. + while len(chunks) <= self._chunk_size: + if len(chunks) > 0: + if chunks[-1] is None or isinstance(chunks[-1], Exception): + break # Hit sentinel value or an error. + # We could switch to relying on the `done` / `finish_reason` field of chunks, + # but that forces us to know about the chunk type here. Prefer sentinel values + # for now. + + item = await self._async_queue.get() + chunks.append(item) + + # Process the sentinel value if it's there. + if chunks[-1] is None: + chunks.pop() # Remove the sentinel value. + do_set_computed = True + + # Shouldn't be needed, but cancel the Tasks this ModelOutputThunk relied on. + if self._generate is not None: + self._generate.cancel() + if self._generate_extra is not None: + # Covers an hf edge case. The task is done generating anything useful but isn't `done` yet. + await self._generate_extra + self._generate_extra.cancel() + + # If ModelOutputThunks get too bulky, we can do additional cleanup here + # and set fields to None. + + elif isinstance(chunks[-1], Exception): + raise chunks[-1] + + for chunk in chunks: + assert self._process is not None + await self._process(self, chunk) + + if do_set_computed: + assert self._underlying_value is not None + self._computed = True + + assert self._post_process is not None + await self._post_process(self) + + match self._action: + case Component(): + self.parsed_repr = self._action._parse(self) + case CBlock(): + assert self.value is not None, ( + "value must be non-None since this thunk is computed" ) - return self._underlying_value # type: ignore - - # Re-raise exception after cleanup if one occurred - if exception_to_raise is not None: - raise exception_to_raise + self.parsed_repr = self.value # type: ignore + case _: + raise ValueError( + "attempted to astream from a model output thunk with no ._action set" + ) + assert self.parsed_repr is not None, ( + "enforce constraint that a computed ModelOutputThunk has a non-None parsed_repr" + ) + return self._underlying_value # type: ignore return ( self._underlying_value From e0cdbb9148a2b5b3e9dc1cb51d19a7937fad33d2 Mon Sep 17 00:00:00 2001 From: "Paul S. Schweigert" Date: Wed, 4 Mar 2026 20:47:38 -0500 Subject: [PATCH 2/4] add test Signed-off-by: Paul S. Schweigert --- .../test_astream_exception_propagation.py | 103 ++++++++++++++++++ 1 file changed, 103 insertions(+) create mode 100644 test/core/test_astream_exception_propagation.py diff --git a/test/core/test_astream_exception_propagation.py b/test/core/test_astream_exception_propagation.py new file mode 100644 index 000000000..883b223a1 --- /dev/null +++ b/test/core/test_astream_exception_propagation.py @@ -0,0 +1,103 @@ +"""Tests that exceptions during generation propagate correctly through ModelOutputThunk.astream(). + +Regression test for issue #577: post_process in a finally block was swallowing +the original generation exception by raising a secondary error from post_process +(which assumes system invariants that don't hold during failures). +""" + +import asyncio + +import pytest + +from mellea.core.base import CBlock, GenerateType, ModelOutputThunk + + +async def _noop_process(mot, chunk): + """Minimal process callback that appends chunk text to the thunk's value.""" + if mot._underlying_value is None: + mot._underlying_value = "" + mot._underlying_value += str(chunk) + + +async def _failing_post_process(mot): + """A post_process that fails, simulating real backends which assume invariants.""" + raise RuntimeError("post_process failed due to broken invariants") + + +def _make_streaming_thunk(post_process=None): + """Create a ModelOutputThunk wired up for async streaming without a real backend.""" + mot = ModelOutputThunk(value=None) + mot._generate_type = GenerateType.ASYNC + mot._process = _noop_process + mot._post_process = post_process or _failing_post_process + mot._action = CBlock("test") + mot._chunk_size = 0 # Don't require minimum chunks + return mot + + +@pytest.mark.asyncio +async def test_astream_propagates_generation_exception(): + """When the backend puts an Exception on the queue, astream must raise that exact exception. + + Before the fix for #577, a finally block called post_process on error, which + would itself fail and swallow the original generation error. + """ + original_error = ValueError("connection reset by peer") + mot = _make_streaming_thunk() + + # Simulate backend putting an error on the queue + await mot._async_queue.put(original_error) + + with pytest.raises(ValueError, match="connection reset by peer"): + await mot.astream() + + +@pytest.mark.asyncio +async def test_astream_exception_is_not_from_post_process(): + """Ensure the raised exception is the generation error, not a post_process error. + + This is the core of issue #577: post_process failures must not mask generation errors. + """ + generation_error = ConnectionError("server unavailable") + mot = _make_streaming_thunk(post_process=_failing_post_process) + + await mot._async_queue.put(generation_error) + + # Must get ConnectionError, NOT RuntimeError from _failing_post_process + with pytest.raises(ConnectionError, match="server unavailable"): + await mot.astream() + + +@pytest.mark.asyncio +async def test_astream_post_process_only_called_on_success(): + """Verify post_process is called on successful completion, not on error.""" + post_process_called = False + + async def _tracking_post_process(mot): + nonlocal post_process_called + post_process_called = True + + # Error path: post_process should NOT be called + mot = _make_streaming_thunk(post_process=_tracking_post_process) + await mot._async_queue.put(RuntimeError("generation failed")) + + with pytest.raises(RuntimeError, match="generation failed"): + await mot.astream() + + assert not post_process_called, ( + "post_process should not be called when generation fails" + ) + + # Success path: post_process SHOULD be called + post_process_called = False + mot = _make_streaming_thunk(post_process=_tracking_post_process) + await mot._async_queue.put("hello") + await mot._async_queue.put(None) # Sentinel for completion + + await mot.astream() + + assert post_process_called, "post_process should be called on successful completion" + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) From fb0bf89824595eeab2c63355d70ea7097386d955 Mon Sep 17 00:00:00 2001 From: "Paul S. Schweigert" Date: Wed, 4 Mar 2026 20:53:12 -0500 Subject: [PATCH 3/4] cleanup Signed-off-by: Paul S. Schweigert --- .../test_astream_exception_propagation.py | 62 +++++-------------- 1 file changed, 16 insertions(+), 46 deletions(-) diff --git a/test/core/test_astream_exception_propagation.py b/test/core/test_astream_exception_propagation.py index 883b223a1..06d5f1ab4 100644 --- a/test/core/test_astream_exception_propagation.py +++ b/test/core/test_astream_exception_propagation.py @@ -5,72 +5,46 @@ (which assumes system invariants that don't hold during failures). """ -import asyncio - import pytest from mellea.core.base import CBlock, GenerateType, ModelOutputThunk async def _noop_process(mot, chunk): - """Minimal process callback that appends chunk text to the thunk's value.""" if mot._underlying_value is None: mot._underlying_value = "" mot._underlying_value += str(chunk) async def _failing_post_process(mot): - """A post_process that fails, simulating real backends which assume invariants.""" raise RuntimeError("post_process failed due to broken invariants") -def _make_streaming_thunk(post_process=None): - """Create a ModelOutputThunk wired up for async streaming without a real backend.""" +def _make_thunk(post_process=_failing_post_process): mot = ModelOutputThunk(value=None) mot._generate_type = GenerateType.ASYNC mot._process = _noop_process - mot._post_process = post_process or _failing_post_process + mot._post_process = post_process mot._action = CBlock("test") - mot._chunk_size = 0 # Don't require minimum chunks + mot._chunk_size = 0 return mot -@pytest.mark.asyncio -async def test_astream_propagates_generation_exception(): - """When the backend puts an Exception on the queue, astream must raise that exact exception. - - Before the fix for #577, a finally block called post_process on error, which - would itself fail and swallow the original generation error. - """ - original_error = ValueError("connection reset by peer") - mot = _make_streaming_thunk() - - # Simulate backend putting an error on the queue - await mot._async_queue.put(original_error) +@pytest.mark.parametrize( + "error", + [ValueError("connection reset by peer"), ConnectionError("server unavailable")], +) +async def test_astream_propagates_generation_exception(error): + """The original generation error must propagate, not a secondary error from post_process.""" + mot = _make_thunk() + await mot._async_queue.put(error) - with pytest.raises(ValueError, match="connection reset by peer"): + with pytest.raises(type(error), match=str(error)): await mot.astream() -@pytest.mark.asyncio -async def test_astream_exception_is_not_from_post_process(): - """Ensure the raised exception is the generation error, not a post_process error. - - This is the core of issue #577: post_process failures must not mask generation errors. - """ - generation_error = ConnectionError("server unavailable") - mot = _make_streaming_thunk(post_process=_failing_post_process) - - await mot._async_queue.put(generation_error) - - # Must get ConnectionError, NOT RuntimeError from _failing_post_process - with pytest.raises(ConnectionError, match="server unavailable"): - await mot.astream() - - -@pytest.mark.asyncio async def test_astream_post_process_only_called_on_success(): - """Verify post_process is called on successful completion, not on error.""" + """post_process must be called on success but not on error.""" post_process_called = False async def _tracking_post_process(mot): @@ -78,7 +52,7 @@ async def _tracking_post_process(mot): post_process_called = True # Error path: post_process should NOT be called - mot = _make_streaming_thunk(post_process=_tracking_post_process) + mot = _make_thunk(post_process=_tracking_post_process) await mot._async_queue.put(RuntimeError("generation failed")) with pytest.raises(RuntimeError, match="generation failed"): @@ -90,14 +64,10 @@ async def _tracking_post_process(mot): # Success path: post_process SHOULD be called post_process_called = False - mot = _make_streaming_thunk(post_process=_tracking_post_process) + mot = _make_thunk(post_process=_tracking_post_process) await mot._async_queue.put("hello") - await mot._async_queue.put(None) # Sentinel for completion + await mot._async_queue.put(None) # sentinel for completion await mot.astream() assert post_process_called, "post_process should be called on successful completion" - - -if __name__ == "__main__": - pytest.main([__file__, "-v"]) From 0613ea0fb8f0f1564eee12787891d8aca3333a78 Mon Sep 17 00:00:00 2001 From: "Paul S. Schweigert" Date: Thu, 5 Mar 2026 12:01:21 -0500 Subject: [PATCH 4/4] handle spans Signed-off-by: Paul S. Schweigert --- mellea/core/base.py | 10 ++++++ .../test_astream_exception_propagation.py | 33 +++++++++++++++++++ 2 files changed, 43 insertions(+) diff --git a/mellea/core/base.py b/mellea/core/base.py index 34d7e3b9c..4ec601e0a 100644 --- a/mellea/core/base.py +++ b/mellea/core/base.py @@ -321,6 +321,16 @@ async def astream(self) -> str: # and set fields to None. elif isinstance(chunks[-1], Exception): + # Close any open telemetry span before propagating the error. + # We can't call full post_process here (it assumes success invariants), + # but we must not leak the span. + span = self._meta.get("_telemetry_span") + if span is not None: + from ..telemetry import end_backend_span, set_span_error + + set_span_error(span, chunks[-1]) + end_backend_span(span) + del self._meta["_telemetry_span"] raise chunks[-1] for chunk in chunks: diff --git a/test/core/test_astream_exception_propagation.py b/test/core/test_astream_exception_propagation.py index 06d5f1ab4..490cd0c9b 100644 --- a/test/core/test_astream_exception_propagation.py +++ b/test/core/test_astream_exception_propagation.py @@ -71,3 +71,36 @@ async def _tracking_post_process(mot): await mot.astream() assert post_process_called, "post_process should be called on successful completion" + + +async def test_astream_closes_telemetry_span_on_error(): + """Telemetry span must be ended and error recorded when generation fails.""" + from unittest.mock import MagicMock + + mock_span = MagicMock() + mot = _make_thunk() + mot._meta["_telemetry_span"] = mock_span + + error = ConnectionError("server unavailable") + await mot._async_queue.put(error) + + with pytest.raises(ConnectionError, match="server unavailable"): + await mot.astream() + + # Span should have been ended and cleaned up + mock_span.record_exception.assert_called_once_with(error) + mock_span.set_status.assert_called_once() + mock_span.end.assert_called_once() + assert "_telemetry_span" not in mot._meta + + +async def test_astream_no_span_leak_when_no_telemetry(): + """When no telemetry span is present, error propagation still works.""" + mot = _make_thunk() + assert "_telemetry_span" not in mot._meta + + error = ValueError("test error") + await mot._async_queue.put(error) + + with pytest.raises(ValueError, match="test error"): + await mot.astream()