fix: bound model stream aggregation and close failed responses

artemis/services/llm.py

@@ -47,6 +47,7 @@ from artemis.config import (
 from artemis.context import ArtemisContext
 from artemis.data_engine.trace import CURRENT_TRACE_ID, DataEngineCallbackHandler
 from artemis.llm.diagnostics import ACTIVE_ATTEMPT, RequestAttempt, error_text, input_images
+from artemis.llm.chunk_accumulator import ChunkAccumulator
 from artemis.llm.final_review import (
     ACTIVE_FINAL_REVIEW,
     FinalReviewBudget,
@@ -425,8 +426,13 @@ async def _run_with_recovery[T](
                 observation_token = ACTIVE_ATTEMPT.set(observation)
                 request_timer = asyncio.timeout(timeout)
                 try:
-                    async with request_timer:
-                        result = await call_once()
+                    try:
+                        async with request_timer:
+                            result = await call_once()
+                    finally:
+                        # Release a failed stream before any retry backoff, even
+                        # when the SDK consumer stopped between yielded chunks.
+                        await observation.close_streams()
                 except (KeyboardInterrupt, asyncio.CancelledError) as e:
                     observation.cancel(e)
                     raise
@@ -818,24 +824,46 @@ class RobustChatModelWrapper:
 
         stream_exec_id = uuid4()
         full_response = None
+        accumulator = ChunkAccumulator()
         argument_guard = ToolArgumentGuard()
         try:
             async with aclosing(self.base_model.astream(*args, **kwargs)) as stream:
+                wait_started = time.monotonic()
                 async for chunk in stream:
                     observation = ACTIVE_ATTEMPT.get()
                     if observation is not None:
+                        observation.stream_work("await_sdk_chunk", time.monotonic() - wait_started)
+                        observe_started = time.monotonic()
                         observation.chunk(chunk)
+                        observation.stream_work("sdk_observe", time.monotonic() - observe_started)
                     argument_guard.observe(chunk)
-                    if full_response is None:
-                        full_response = chunk
-                    else:
-                        full_response = full_response + chunk
+                    accumulator.append(chunk)
+                    if accumulator.ready:
+                        # Buffered transport input may yield without suspending.
+                        # Give stop/deadline signals a turn before a batch merge.
+                        await asyncio.sleep(0)
+                        merge_started = time.monotonic()
+                        full_response = accumulator.flush()
+                        if observation is not None:
+                            observation.stream_work("merge", time.monotonic() - merge_started)
                     if emit_deltas:
+                        emit_started = time.monotonic()
                         self._emit_stream_delta(stream_exec_id, chunk)
+                        if observation is not None:
+                            observation.stream_work("emit", time.monotonic() - emit_started)
+                    wait_started = time.monotonic()
+            merge_started = time.monotonic()
+            full_response = accumulator.flush()
+            observation = ACTIVE_ATTEMPT.get()
+            if observation is not None:
+                observation.stream_work("merge", time.monotonic() - merge_started)
             if full_response is not None:
                 ensure_response_complete(full_response)
         except (KeyboardInterrupt, asyncio.CancelledError):
-            if full_response is not None:
+            observation = ACTIVE_ATTEMPT.get()
+            if observation is not None:
+                observation.stream_cancelled()
+            if accumulator.has_output:
                 _record_llm_event(
                     "llm_stream_reset",
                     {
@@ -847,7 +875,7 @@ class RobustChatModelWrapper:
                 )
             raise
         except Exception as stream_error:
-            if full_response is not None or isinstance(
+            if accumulator.has_output or isinstance(
                 stream_error, (ToolArgumentOverflow, ResponseIntegrityError)
             ):
                 if isinstance(stream_error, ResponseIntegrityError):