Codex/stream tool args truncation error

flocks/session/lifecycle/retry.py

@@ -135,6 +135,9 @@ def retryable(error: Dict[str, Any]) -> Optional[str]:
         """
         error_name = error.get("name", "")
         error_data = error.get("data", {})
+
+        if error_name == "StreamToolArgumentsTruncatedError":
+            return "Model output was truncated while generating tool arguments"
 
         # Check if it's an APIError with isRetryable flag
         if error_name == "APIError":

flocks/session/runner.py

@@ -1338,6 +1338,8 @@ def classify_failover_error(error: Dict[str, Any]) -> FailoverDecision:
             "policy violation",
         )):
             return FailoverDecision(True, "content_policy")
+        if error_name == "StreamToolArgumentsTruncatedError":
+            return FailoverDecision(True, "stream_truncated")
         if error_name == "JSONDecodeError" or any(
             pattern in lowered for pattern in (
                 "malformed response", "invalid response", "empty choices",
@@ -1360,12 +1362,18 @@ def _deferred_failure_result(
         assistant_message_id: Optional[str],
         decision: FailoverDecision,
         attempts: int,
+        allow_fallback_override: Optional[bool] = None,
     ) -> StepResult:
         state = LlmAttemptState(
             received_chunk=self._attempt_state.received_chunk,
             observable_output_started=self._attempt_state.observable_output_started,
             tool_execution_started=self._attempt_state.tool_execution_started,
         )
+        allow_fallback = (
+            allow_fallback_override
+            if allow_fallback_override is not None
+            else decision.eligible and state.replay_safe
+        )
         return StepResult(
             action="stop",
             error=message,
@@ -1374,11 +1382,15 @@ def _deferred_failure_result(
                 error_data=error_data,
                 assistant_message_id=assistant_message_id,
                 reason=decision.reason,
-                allow_fallback=decision.eligible and state.replay_safe,
+                allow_fallback=allow_fallback,
                 attempt_state=state,
                 attempts=attempts,
             ),
         )
+
+    @staticmethod
+    def _is_stream_tool_arguments_truncated_error(error: Dict[str, Any]) -> bool:
+        return error.get("name") == "StreamToolArgumentsTruncatedError"
 
     async def _process_step(
         self,
@@ -1613,34 +1625,96 @@ async def _process_step(
             # Disable tools when max steps reached
             tools = []
 
-        # Create assistant message (will be reused across retries)
-        assistant_msg = await Message.create(
-            session_id=self.session.id,
-            role=MessageRole.ASSISTANT,
-            content="",
-            agent=agent.name,
-            model_id=self.model_id,
-            provider_id=self.provider_id,
-            parent_id=last_user.id,
-        )
-
-        # Publish assistant message SSE event so frontends can show the message card
-        if self.callbacks.event_publish_callback:
-            import time as _time
+        async def _publish_assistant_created(msg: MessageInfo) -> None:
+            if not self.callbacks.event_publish_callback:
+                return
             await self.callbacks.event_publish_callback("message.updated", {
                 "info": {
-                    "id": assistant_msg.id,
+                    "id": msg.id,
                     "sessionID": self.session.id,
                     "role": "assistant",
-                    "time": {"created": int(_time.time() * 1000)},
+                    "time": {"created": int(time.time() * 1000)},
                     "parentID": last_user.id,
                     "modelID": self.model_id,
                     "providerID": self.provider_id,
                     "agent": agent.name,
                     "mode": agent.name,
-                    "tokens": {"input": 0, "output": 0, "reasoning": 0, "cache": {"read": 0, "write": 0}},
+                    "tokens": {
+                        "input": 0,
+                        "output": 0,
+                        "reasoning": 0,
+                        "cache": {"read": 0, "write": 0},
+                    },
                 }
             })
+
+        async def _create_attempt_assistant_message(*, publish: bool = True) -> MessageInfo:
+            msg = await Message.create(
+                session_id=self.session.id,
+                role=MessageRole.ASSISTANT,
+                content="",
+                agent=agent.name,
+                model_id=self.model_id,
+                provider_id=self.provider_id,
+                parent_id=last_user.id,
+            )
+            if publish:
+                await _publish_assistant_created(msg)
+            return msg
+
+        async def _replace_assistant_message_for_replay(reason: str) -> bool:
+            nonlocal assistant_msg
+            previous_msg = assistant_msg
+            try:
+                next_msg = await _create_attempt_assistant_message(publish=False)
+            except Exception as exc:
+                log.error("runner.step.replay_message_create_failed", {
+                    "session_id": self.session.id,
+                    "previous_message_id": previous_msg.id,
+                    "reason": reason,
+                    "error": str(exc),
+                })
+                return False
+
+            try:
+                deleted = await Message.delete(self.session.id, previous_msg.id)
+            except Exception as exc:
+                deleted = False
+                log.error("runner.step.replay_message_delete_failed", {
+                    "session_id": self.session.id,
+                    "previous_message_id": previous_msg.id,
+                    "next_message_id": next_msg.id,
+                    "reason": reason,
+                    "error": str(exc),
+                })
+            if not deleted:
+                try:
+                    await Message.delete(self.session.id, next_msg.id)
+                except Exception as exc:
+                    log.debug("runner.step.replay_message_cleanup_failed", {
+                        "session_id": self.session.id,
+                        "message_id": next_msg.id,
+                        "error": str(exc),
+                    })
+                return False
+
+            if self.callbacks.event_publish_callback:
+                await self.callbacks.event_publish_callback("message.removed", {
+                    "sessionID": self.session.id,
+                    "messageID": previous_msg.id,
+                })
+                await _publish_assistant_created(next_msg)
+            assistant_msg = next_msg
+            log.info("runner.step.replay_message_replaced", {
+                "session_id": self.session.id,
+                "previous_message_id": previous_msg.id,
+                "next_message_id": next_msg.id,
+                "reason": reason,
+            })
+            return True
+
+        # Create assistant message for the first attempt.
+        assistant_msg = await _create_attempt_assistant_message()
 
         # Retry loop matching Flocks' SessionProcessor.process()
         # MAX_ERROR_RETRIES caps exception-based retries so a permanently-failing
@@ -1824,6 +1898,9 @@ async def _process_step(
                 # Check if retryable
                 retry_message = SessionRetry.retryable(error_dict)
                 failover_decision = self.classify_failover_error(error_dict)
+                is_stream_tool_args_truncated = (
+                    self._is_stream_tool_arguments_truncated_error(error_dict)
+                )
                 retry_limit = MAX_ERROR_RETRIES
                 will_retry = retry_message is not None and error_attempt <= retry_limit
                 retry_blocked_by_tool_execution = (
@@ -1837,7 +1914,18 @@ async def _process_step(
                 elif self._defer_step_errors and not self._attempt_state.replay_safe:
                     # Retrying after text/reasoning/tool activity can duplicate
                     # visible output or execute a tool twice.
-                    will_retry = False
+                    will_retry = will_retry and is_stream_tool_args_truncated
+
+                if will_retry and is_stream_tool_args_truncated:
+                    # A truncated tool-argument stream already created a
+                    # partial assistant message (and usually a tool part). The
+                    # retry is only safe if that partial attempt can be removed
+                    # before the next provider call.
+                    replaced = await _replace_assistant_message_for_replay(
+                        reason="stream_tool_arguments_truncated"
+                    )
+                    if not replaced:
+                        will_retry = False
 
                 if will_retry:
                     # Error is retryable and we have budget left
@@ -1869,6 +1957,8 @@ async def _process_step(
 
                     # Wait before retry
                     await SessionRetry.sleep(delay_ms, self._abort)
+
+                    self._attempt_state = LlmAttemptState()
 
                     # Continue to next retry attempt
                     continue
@@ -1898,12 +1988,19 @@ async def _process_step(
                         error_dict["data"]["displayMessage"] = CONNECTION_ERROR_DISPLAY_MESSAGE
 
                     if self._defer_step_errors:
+                        allow_fallback_override = None
+                        if is_stream_tool_args_truncated:
+                            allow_fallback_override = (
+                                failover_decision.eligible
+                                and not retry_blocked_by_tool_execution
+                            )
                         return self._deferred_failure_result(
                             message=final_error_message,
                             error_data=error_dict,
                             assistant_message_id=assistant_msg.id,
                             decision=failover_decision,
                             attempts=error_attempt,
+                            allow_fallback_override=allow_fallback_override,
                         )
 
                     if self.callbacks.on_error:
@@ -2509,6 +2606,17 @@ def _exception_to_error_dict(self, exception: Exception) -> Dict[str, Any]:
             }
         }
 
+        if type(exception).__name__ == "StreamToolArgumentsTruncatedError":
+            error_dict["data"].update({
+                "isRetryable": True,
+                "streamToolArgumentsTruncated": True,
+                "toolCallID": getattr(exception, "tool_call_id", None),
+                "toolName": getattr(exception, "tool_name", None),
+                "finishReason": getattr(exception, "finish_reason", None),
+                "argumentsLength": getattr(exception, "arguments_len", None),
+                "argumentsPreview": getattr(exception, "arguments_preview", None),
+            })
+
         transport_exception = _find_retryable_transport_exception(exception)
         if transport_exception is not None:
             transport_type = type(transport_exception).__name__
@@ -3749,7 +3857,11 @@ async def _flush_reasoning_rewriter() -> None:
             "agent": agent.name,
         })
 
-        await tool_accumulator.flush_remaining(stream_finish_reason)
+        try:
+            await tool_accumulator.flush_remaining(stream_finish_reason)
+        except Exception:
+            await processor.drain_parallel_tool_calls()
+            raise
 
         if stream_text_rewriter is not None:
             trailing_text = stream_text_rewriter.flush()

flocks/session/streaming/stream_events.py

@@ -18,6 +18,7 @@
     "tool-input-start",
     "tool-input-delta",
     "tool-input-end",
+    "tool-input-error",
     "tool-call",
     "tool-result",
     "tool-error",
@@ -82,6 +83,15 @@ class ToolInputEndEvent(BaseStreamEvent):
     id: str
 
 
+class ToolInputErrorEvent(BaseStreamEvent):
+    """Tool input failed before a runnable tool call was produced."""
+    type: Literal["tool-input-error"] = "tool-input-error"
+    id: str
+    tool_name: str
+    input: Dict[str, Any] = Field(default_factory=dict)
+    error: str
+
+
 class ToolCallEvent(BaseStreamEvent):
     """Tool call request (ready to execute)"""
     type: Literal["tool-call"] = "tool-call"
@@ -152,6 +162,7 @@ class FinishEvent(BaseStreamEvent):
     ToolInputStartEvent |
     ToolInputDeltaEvent |
     ToolInputEndEvent |
+    ToolInputErrorEvent |
     ToolCallEvent |
     ToolResultEvent |
     ToolErrorEvent |
@@ -184,6 +195,7 @@ def event_from_dict(data: Dict[str, Any]) -> StreamEvent:
         "tool-input-start": ToolInputStartEvent,
         "tool-input-delta": ToolInputDeltaEvent,
         "tool-input-end": ToolInputEndEvent,
+        "tool-input-error": ToolInputErrorEvent,
         "tool-call": ToolCallEvent,
         "tool-result": ToolResultEvent,
         "tool-error": ToolErrorEvent,

flocks/session/streaming/stream_processor.py

@@ -38,6 +38,7 @@
     TextDeltaEvent,
     TextEndEvent,
     ToolInputStartEvent,
+    ToolInputErrorEvent,
 )
 from flocks.tool.registry import ToolRegistry, ToolContext, ToolResult
 from flocks.permission import PermissionNext
@@ -218,6 +219,9 @@ async def process_event(self, event: StreamEvent) -> None:
 
         elif event_type == "tool-input-end":
             pass  # Input is complete
+
+        elif event_type == "tool-input-error":
+            await self._handle_tool_input_error(event)
 
         elif event_type == "tool-call":
             if self._should_run_tool_call_parallel(event):
@@ -471,6 +475,69 @@ async def _handle_tool_input_start(self, event: ToolInputStartEvent) -> None:
             })
         except Exception as e:
             log.error("stream.tool_input_start.store_part_failed", {"error": str(e)})
+
+    async def _handle_tool_input_error(self, event: ToolInputErrorEvent) -> None:
+        """Mark an input-generation failure without executing a tool."""
+        if event.id in self.tool_calls:
+            tool_state = self.tool_calls[event.id]
+            part_id = tool_state.part_id
+        else:
+            part_id = Identifier.create("part")
+            tool_state = ToolCallState(
+                id=event.id,
+                name=event.tool_name,
+                input=event.input,
+                part_id=part_id,
+                status="pending",
+            )
+            self.tool_calls[event.id] = tool_state
+
+        tool_state.name = event.tool_name
+        tool_state.input = event.input
+        tool_state.status = "error"
+        tool_state.error = event.error
+
+        tool_error_time = int(datetime.now().timestamp() * 1000)
+        error_state = ToolStateError(
+            status="error",
+            input=event.input,
+            error=event.error,
+            time={"start": tool_error_time, "end": tool_error_time},
+        )
+        error_part = ToolPart(
+            id=part_id,
+            sessionID=self.session_id,
+            messageID=self.assistant_message.id,
+            type="tool",
+            callID=event.id,
+            tool=event.tool_name,
+            state=error_state,
+        )
+        await Message.store_part(self.session_id, self.assistant_message.id, error_part)
+
+        if self.event_publish_callback:
+            await self.event_publish_callback("message.part.updated", {
+                "part": {
+                    "id": part_id,
+                    "messageID": self.assistant_message.id,
+                    "sessionID": self.session_id,
+                    "type": "tool",
+                    "callID": event.id,
+                    "tool": event.tool_name,
+                    "state": {
+                        "status": "error",
+                        "input": event.input,
+                        "error": event.error,
+                        "time": {"start": tool_error_time, "end": tool_error_time},
+                    },
+                },
+            })
+
+        log.warn("stream.tool_input.error", {
+            "tool_call_id": event.id,
+            "tool_name": event.tool_name,
+            "error": event.error,
+        })
 
     async def _handle_tool_call(self, event: ToolCallEvent) -> None:
         """
@@ -1778,7 +1845,7 @@ def _parse_dsml_text_tool_calls(self, text: str) -> list[dict]:
             re.DOTALL | re.IGNORECASE,
         ):
             body = match.group(1).strip()
-            if not body or not body[:1] in "{[":
+            if not body or body[:1] not in "{[":
                 continue
 
             try: