diff --git a/flocks/session/lifecycle/retry.py b/flocks/session/lifecycle/retry.py index ae7bf77f3..fdf048145 100644 --- a/flocks/session/lifecycle/retry.py +++ b/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": diff --git a/flocks/session/runner.py b/flocks/session/runner.py index 70fcff824..3cc4f8a7b 100644 --- a/flocks/session/runner.py +++ b/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() diff --git a/flocks/session/streaming/stream_events.py b/flocks/session/streaming/stream_events.py index 607054027..64d96b8d5 100644 --- a/flocks/session/streaming/stream_events.py +++ b/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, diff --git a/flocks/session/streaming/stream_processor.py b/flocks/session/streaming/stream_processor.py index dc4e4b967..c03940a5d 100644 --- a/flocks/session/streaming/stream_processor.py +++ b/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: diff --git a/flocks/session/streaming/tool_accumulator.py b/flocks/session/streaming/tool_accumulator.py index 6e266e95d..bfd6536de 100644 --- a/flocks/session/streaming/tool_accumulator.py +++ b/flocks/session/streaming/tool_accumulator.py @@ -20,12 +20,37 @@ from flocks.tool.registry import ToolRegistry from flocks.session.streaming.stream_events import ( ToolInputStartEvent, + ToolInputErrorEvent, ToolCallEvent, ) log = Log.create(service="tool_accumulator") +class StreamToolArgumentsTruncatedError(RuntimeError): + """Raised when the stream ends before tool arguments form valid JSON.""" + + def __init__( + self, + *, + tool_call_id: str, + tool_name: str, + finish_reason: str, + arguments_len: int, + arguments_preview: str, + ) -> None: + self.tool_call_id = tool_call_id + self.tool_name = tool_name + self.finish_reason = finish_reason + self.arguments_len = arguments_len + self.arguments_preview = arguments_preview + super().__init__( + "Model output was truncated while generating tool arguments " + f"for '{tool_name}' (finish_reason='{finish_reason}', " + f"{arguments_len} chars). The tool was not executed." + ) + + class ToolCallAccumulator: """Accumulates streamed tool-call JSON fragments and dispatches execution. @@ -121,13 +146,52 @@ async def flush_remaining( ) -> None: """Process any tool calls still in the accumulator after the stream ends.""" is_truncated = stream_finish_reason in ("length", "max_tokens") + truncation_error: StreamToolArgumentsTruncatedError | None = None for tc_id, tc_data in list(self._accumulator.items()): - if tc_data.get("completed"): + if tc_data.get("completed") or tc_data.get("failed"): continue accumulated_args = tc_data.get("arguments_str", "") tool_name = tc_data.get("name", "") - if not (accumulated_args and tool_name): + if not tool_name: + continue + + if is_truncated: + if accumulated_args: + detail = ( + f"Tool arguments for '{tool_name}' cut off at " + f"{len(accumulated_args)} chars." + ) + else: + detail = f"Tool arguments for '{tool_name}' were not completed." + error_msg = ( + f"Output was truncated (finish_reason='{stream_finish_reason}'). " + f"{detail} The tool was not executed." + ) + await self._processor.process_event( + ToolInputErrorEvent( + id=tc_id, + tool_name=tool_name, + input={ + "tool": tool_name, + "arguments_preview": accumulated_args[:500], + "finish_reason": stream_finish_reason, + }, + error=error_msg, + ) + ) + tc_data["failed"] = True + if truncation_error is None: + truncation_error = StreamToolArgumentsTruncatedError( + tool_call_id=tc_id, + tool_name=tool_name, + finish_reason=str(stream_finish_reason), + arguments_len=len(accumulated_args), + arguments_preview=accumulated_args[:500], + ) + continue + + if not accumulated_args: continue arguments, ok = _parse_json_robust(accumulated_args) @@ -151,17 +215,10 @@ async def flush_remaining( continue # All strategies failed — redirect to invalid tool - if is_truncated: - error_msg = ( - f"Output was truncated (finish_reason='{stream_finish_reason}'). " - f"Tool arguments for '{tool_name}' cut off at {len(accumulated_args)} chars. " - f"Please reduce the content size or split the operation." - ) - else: - error_msg = ( - f"Failed to parse tool arguments ({len(accumulated_args)} chars). " - f"Please ensure valid JSON with balanced braces/brackets." - ) + error_msg = ( + f"Failed to parse tool arguments ({len(accumulated_args)} chars). " + f"Please ensure valid JSON with balanced braces/brackets." + ) await self._processor.process_event( ToolCallEvent( @@ -175,6 +232,9 @@ async def flush_remaining( ) ) + if truncation_error is not None: + raise truncation_error + # ------------------------------------------------------------------ # Internal helpers # ------------------------------------------------------------------ diff --git a/tests/session/test_auto_model_failover.py b/tests/session/test_auto_model_failover.py index b1d655b1f..6cf432e65 100644 --- a/tests/session/test_auto_model_failover.py +++ b/tests/session/test_auto_model_failover.py @@ -192,6 +192,78 @@ async def test_auto_runner_uses_standard_retry_policy( assert result.failure is not None assert call_llm.await_count == expected_calls + assert result.failure.allow_fallback is True + + +@pytest.mark.asyncio +async def test_retry_exhausted_safe_api_error_switches_to_fallback(monkeypatch): + ctx = _ctx() + last_user = SimpleNamespace(id="msg_user", agent="rex", role="user") + events = [] + calls = [] + create_count = 0 + + provider = MagicMock() + provider.is_configured.return_value = True + + async def create_message(**_kwargs): + nonlocal create_count + create_count += 1 + return SimpleNamespace(id=f"msg_assistant_{create_count}") + + async def call_llm(runner, *_args, **_kwargs): + calls.append((runner.provider_id, runner.model_id)) + if runner.provider_id == "primary": + failure = RuntimeError("Provider HTTP 500") + failure.status_code = 500 + raise failure + return StepResult(action="stop", content="recovered") + + async def publish(event, payload): + events.append((event, payload)) + + monkeypatch.setattr( + "flocks.session.runner.Agent.get", + AsyncMock(return_value=SimpleNamespace( + name="rex", + steps=None, + mode="primary", + prompt="", + tools=[], + )), + ) + monkeypatch.setattr("flocks.session.runner.Provider.get", lambda _provider_id: provider) + monkeypatch.setattr("flocks.session.runner.Provider.apply_config", AsyncMock()) + monkeypatch.setattr( + "flocks.session.runner.SessionPrompt.build_system_prompts", + AsyncMock(return_value=[]), + ) + monkeypatch.setattr(SessionRunner, "_build_callable_tool_schema", AsyncMock(return_value=[])) + monkeypatch.setattr( + SessionRunner, + "_to_chat_messages", + AsyncMock(return_value=[SimpleNamespace(role="user", content="hello")]), + ) + monkeypatch.setattr(Message, "get_text_content", AsyncMock(return_value="hello")) + monkeypatch.setattr(Message, "parts", AsyncMock(return_value=[])) + monkeypatch.setattr(Message, "create", create_message) + monkeypatch.setattr(Message, "delete", AsyncMock(return_value=True)) + monkeypatch.setattr(Message, "update", AsyncMock()) + monkeypatch.setattr(SessionRunner, "_call_llm", call_llm) + monkeypatch.setattr("flocks.session.runner.SessionRetry.sleep", AsyncMock()) + + result = await SessionLoop._process_step_with_failover( + ctx, + LoopCallbacks(event_publish_callback=publish), + [last_user], + last_user, + ) + + assert result.content == "recovered" + assert calls == [("primary", "primary-model")] * 6 + [("fallback", "fallback-model")] + assert (ctx.provider_id, ctx.model_id) == ("fallback", "fallback-model") + assert any(event == "message.removed" for event, _ in events) + assert any(event == "session.model.fallback" for event, _ in events) @pytest.mark.asyncio diff --git a/tests/session/test_retry.py b/tests/session/test_retry.py index 8fc07d29e..78ac43abd 100644 --- a/tests/session/test_retry.py +++ b/tests/session/test_retry.py @@ -50,6 +50,14 @@ def test_api_error_retryable_generic(self): result = SessionRetry.retryable(error) assert result == "Internal server error" + def test_stream_tool_arguments_truncated_error_is_retryable(self): + error = { + "name": "StreamToolArgumentsTruncatedError", + "data": {"message": "tool arguments were truncated"}, + } + result = SessionRetry.retryable(error) + assert result == "Model output was truncated while generating tool arguments" + def test_json_message_too_many_requests(self): import json msg = json.dumps({"type": "error", "error": {"type": "too_many_requests"}}) diff --git a/tests/session/test_runner_step.py b/tests/session/test_runner_step.py index 15ae0d025..849bbccd0 100644 --- a/tests/session/test_runner_step.py +++ b/tests/session/test_runner_step.py @@ -33,6 +33,7 @@ StepResult, ToolCall, ) +from flocks.session.streaming.tool_accumulator import StreamToolArgumentsTruncatedError from flocks.session.prompt import ( SessionPrompt, SystemPromptBlock, @@ -266,6 +267,24 @@ def test_connection_error_exception_is_retryable(self): assert result["data"]["isRetryable"] is True assert result["data"]["displayMessage"] == runner_mod.CONNECTION_ERROR_DISPLAY_MESSAGE + def test_stream_tool_arguments_truncated_exception_is_retryable(self): + runner = _make_runner() + exc = StreamToolArgumentsTruncatedError( + tool_call_id="call_trunc", + tool_name="write", + finish_reason="length", + arguments_len=42, + arguments_preview='{"path":', + ) + + result = runner._exception_to_error_dict(exc) + + assert result["name"] == "StreamToolArgumentsTruncatedError" + assert result["data"]["isRetryable"] is True + assert result["data"]["streamToolArgumentsTruncated"] is True + assert result["data"]["toolCallID"] == "call_trunc" + assert result["data"]["toolName"] == "write" + def test_incomplete_chunked_read_exception_is_retryable_connection_error(self): runner = _make_runner() exc = Exception( @@ -2377,6 +2396,212 @@ async def fake_to_chat_messages(_messages, _system_prompts): # noqa: ANN001 assert result.content == "done" +@pytest.mark.asyncio +async def test_process_step_retries_truncated_tool_arguments_with_fresh_message(monkeypatch): + runner = _make_runner("ses_runner_stream_tool_args_retry") + events = [] + + async def capture_event(event_type, data): + events.append((event_type, data)) + + runner.callbacks = RunnerCallbacks( + on_error=AsyncMock(), + event_publish_callback=capture_event, + ) + + last_user = UserMessageInfo( + id="msg_user_stream_tool_args_retry", + sessionID=runner.session.id, + role="user", + time={"created": 1_000}, + agent="rex", + model={"providerID": "anthropic", "modelID": "claude-sonnet"}, + ) + agent = SimpleNamespace(name="rex", steps=None, mode="primary", prompt="", tools=[]) + provider = MagicMock() + provider.is_configured.return_value = True + assistant_1 = SimpleNamespace(id="msg_assistant_truncated_1") + assistant_2 = SimpleNamespace(id="msg_assistant_truncated_2") + create_mock = AsyncMock(side_effect=[assistant_1, assistant_2]) + delete_mock = AsyncMock(return_value=True) + update_mock = AsyncMock(return_value=None) + call_ids = [] + + async def fake_call_llm(*_args, **kwargs): + call_ids.append(kwargs["assistant_msg"].id) + if len(call_ids) == 1: + runner._attempt_state.observable_output_started = True + raise StreamToolArgumentsTruncatedError( + tool_call_id="call_trunc", + tool_name="write", + finish_reason="length", + arguments_len=42, + arguments_preview='{"path":', + ) + assert runner._attempt_state.observable_output_started is False + return StepResult(action="stop", content="done") + + monkeypatch.setattr(runner_mod.Agent, "get", AsyncMock(return_value=agent)) + monkeypatch.setattr(runner_mod.Provider, "get", lambda provider_id: provider) + monkeypatch.setattr(runner_mod.Provider, "apply_config", AsyncMock(return_value=None)) + monkeypatch.setattr(runner_mod.SessionPrompt, "build_system_prompt_blocks", AsyncMock(return_value=[])) + monkeypatch.setattr(runner, "_build_callable_tool_schema", AsyncMock(return_value=[])) + monkeypatch.setattr( + runner, + "_to_chat_messages", + AsyncMock(return_value=[SimpleNamespace(role="user", content="hi")]), + ) + monkeypatch.setattr(runner_mod.Message, "get_text_content", AsyncMock(return_value="hi")) + monkeypatch.setattr(runner_mod.Message, "parts", AsyncMock(return_value=[])) + monkeypatch.setattr(runner_mod.Message, "create", create_mock) + monkeypatch.setattr(runner_mod.Message, "delete", delete_mock) + monkeypatch.setattr(runner_mod.Message, "update", update_mock) + monkeypatch.setattr(runner_mod.SessionRetry, "sleep", AsyncMock(return_value=None)) + monkeypatch.setattr(runner, "_call_llm", fake_call_llm) + + result = await runner._process_step([last_user], last_user) + + assert result.action == "stop" + assert result.content == "done" + assert call_ids == [assistant_1.id, assistant_2.id] + delete_mock.assert_awaited_once_with(runner.session.id, assistant_1.id) + runner.callbacks.on_error.assert_not_awaited() + assert ("message.removed", {"sessionID": runner.session.id, "messageID": assistant_1.id}) in events + assert events[-1][0] == "message.updated" + assert events[-1][1]["info"]["id"] == assistant_2.id + assert update_mock.await_args_list[-1].args[1] == assistant_2.id + assert update_mock.await_args_list[-1].kwargs["finish"] == "stop" + + +@pytest.mark.asyncio +async def test_process_step_does_not_retry_truncated_tool_arguments_after_tool_started(monkeypatch): + runner = _make_runner("ses_runner_stream_tool_args_no_retry_after_tool") + runner.callbacks = RunnerCallbacks(on_error=AsyncMock()) + + last_user = UserMessageInfo( + id="msg_user_stream_tool_args_no_retry_after_tool", + sessionID=runner.session.id, + role="user", + time={"created": 1_000}, + agent="rex", + model={"providerID": "anthropic", "modelID": "claude-sonnet"}, + ) + agent = SimpleNamespace(name="rex", steps=None, mode="primary", prompt="", tools=[]) + provider = MagicMock() + provider.is_configured.return_value = True + assistant = SimpleNamespace(id="msg_assistant_no_retry_after_tool") + update_mock = AsyncMock(return_value=None) + delete_mock = AsyncMock(return_value=True) + call_count = 0 + + async def fake_call_llm(*_args, **_kwargs): + nonlocal call_count + call_count += 1 + runner._attempt_state.observable_output_started = True + runner._attempt_state.tool_execution_started = True + raise StreamToolArgumentsTruncatedError( + tool_call_id="call_trunc", + tool_name="write", + finish_reason="length", + arguments_len=42, + arguments_preview='{"path":', + ) + + monkeypatch.setattr(runner_mod.Agent, "get", AsyncMock(return_value=agent)) + monkeypatch.setattr(runner_mod.Provider, "get", lambda provider_id: provider) + monkeypatch.setattr(runner_mod.Provider, "apply_config", AsyncMock(return_value=None)) + monkeypatch.setattr(runner_mod.SessionPrompt, "build_system_prompt_blocks", AsyncMock(return_value=[])) + monkeypatch.setattr(runner, "_build_callable_tool_schema", AsyncMock(return_value=[])) + monkeypatch.setattr( + runner, + "_to_chat_messages", + AsyncMock(return_value=[SimpleNamespace(role="user", content="hi")]), + ) + monkeypatch.setattr(runner_mod.Message, "get_text_content", AsyncMock(return_value="hi")) + monkeypatch.setattr(runner_mod.Message, "parts", AsyncMock(return_value=[])) + monkeypatch.setattr(runner_mod.Message, "store_part", AsyncMock(return_value=None)) + monkeypatch.setattr(runner_mod.Message, "create", AsyncMock(return_value=assistant)) + monkeypatch.setattr(runner_mod.Message, "delete", delete_mock) + monkeypatch.setattr(runner_mod.Message, "update", update_mock) + monkeypatch.setattr(runner_mod.SessionRetry, "sleep", AsyncMock(return_value=None)) + monkeypatch.setattr(runner, "_call_llm", fake_call_llm) + + result = await runner._process_step([last_user], last_user) + + assert call_count == 1 + assert result.action == "stop" + assert "truncated" in result.error.lower() + delete_mock.assert_not_awaited() + runner.callbacks.on_error.assert_awaited_once() + assert update_mock.await_args_list[-1].args[1] == assistant.id + assert update_mock.await_args_list[-1].kwargs["finish"] == "error" + + +@pytest.mark.asyncio +async def test_process_step_allows_fallback_after_truncated_tool_argument_retries_exhausted(monkeypatch): + runner = SessionRunner( + session=_make_session("ses_runner_stream_tool_args_fallback"), + provider_id="primary", + model_id="primary-model", + defer_step_errors=True, + failover_available=True, + ) + + last_user = UserMessageInfo( + id="msg_user_stream_tool_args_fallback", + sessionID=runner.session.id, + role="user", + time={"created": 1_000}, + agent="rex", + model={"providerID": "primary", "modelID": "primary-model"}, + ) + agent = SimpleNamespace(name="rex", steps=None, mode="primary", prompt="", tools=[]) + provider = MagicMock() + provider.is_configured.return_value = True + create_count = 0 + + async def create_message(**_kwargs): + nonlocal create_count + create_count += 1 + return SimpleNamespace(id=f"msg_assistant_fallback_{create_count}") + + async def fake_call_llm(*_args, **_kwargs): + runner._attempt_state.observable_output_started = True + raise StreamToolArgumentsTruncatedError( + tool_call_id="call_trunc", + tool_name="write", + finish_reason="length", + arguments_len=42, + arguments_preview='{"path":', + ) + + monkeypatch.setattr(runner_mod.Agent, "get", AsyncMock(return_value=agent)) + monkeypatch.setattr(runner_mod.Provider, "get", lambda provider_id: provider) + monkeypatch.setattr(runner_mod.Provider, "apply_config", AsyncMock(return_value=None)) + monkeypatch.setattr(runner_mod.SessionPrompt, "build_system_prompt_blocks", AsyncMock(return_value=[])) + monkeypatch.setattr(runner, "_build_callable_tool_schema", AsyncMock(return_value=[])) + monkeypatch.setattr( + runner, + "_to_chat_messages", + AsyncMock(return_value=[SimpleNamespace(role="user", content="hi")]), + ) + monkeypatch.setattr(runner_mod.Message, "get_text_content", AsyncMock(return_value="hi")) + monkeypatch.setattr(runner_mod.Message, "parts", AsyncMock(return_value=[])) + monkeypatch.setattr(runner_mod.Message, "create", create_message) + monkeypatch.setattr(runner_mod.Message, "delete", AsyncMock(return_value=True)) + monkeypatch.setattr(runner_mod.Message, "update", AsyncMock(return_value=None)) + monkeypatch.setattr(runner_mod.SessionRetry, "sleep", AsyncMock(return_value=None)) + monkeypatch.setattr(runner, "_call_llm", fake_call_llm) + + result = await runner._process_step([last_user], last_user) + + assert result.failure is not None + assert result.failure.reason == "stream_truncated" + assert result.failure.allow_fallback is True + assert result.failure.attempt_state.observable_output_started is True + assert result.failure.attempt_state.tool_execution_started is False + + @pytest.mark.asyncio async def test_process_step_limits_connection_error_retries(monkeypatch): runner = _make_runner("ses_runner_connection_error") diff --git a/tests/session/test_stream_processor.py b/tests/session/test_stream_processor.py index 62b21e51b..b770f579c 100644 --- a/tests/session/test_stream_processor.py +++ b/tests/session/test_stream_processor.py @@ -28,6 +28,7 @@ TextEndEvent, TextStartEvent, ToolCallEvent, + ToolInputErrorEvent, ToolInputStartEvent, ) from flocks.session.message import MessageRole, ToolStateError @@ -495,6 +496,46 @@ async def test_tool_call_state_created(self): assert "tc_002" in proc.tool_calls assert proc.tool_calls["tc_002"].name == "read_file" + @pytest.mark.asyncio + async def test_tool_input_error_updates_pending_part_without_execution(self): + event_callback = AsyncMock() + proc = _make_processor(event_callback=event_callback) + execute_mock = AsyncMock(return_value=ToolResult(success=True, output="should not run")) + + with ( + patch("flocks.session.streaming.stream_processor.Message.store_part", new=AsyncMock()) as mock_store, + patch( + "flocks.session.streaming.stream_processor.ToolRegistry.execute", + new=execute_mock, + ), + ): + await proc.process_event(ToolInputStartEvent(id="tc_trunc", tool_name="write")) + await proc.process_event( + ToolInputErrorEvent( + id="tc_trunc", + tool_name="write", + input={"arguments_preview": '{"path": "/tmp/f"', "finish_reason": "length"}, + error="Output was truncated while generating tool arguments.", + ) + ) + + execute_mock.assert_not_awaited() + state = proc.tool_calls["tc_trunc"] + assert state.status == "error" + assert state.name == "write" + assert "truncated" in state.error + + completed_part = mock_store.await_args_list[-1].args[2] + assert completed_part.tool == "write" + assert completed_part.state.status == "error" + assert completed_part.state.input["finish_reason"] == "length" + assert "truncated" in completed_part.state.error + + published_part = event_callback.await_args_list[-1].args[1]["part"] + assert published_part["tool"] == "write" + assert published_part["state"]["status"] == "error" + assert published_part["state"]["input"]["finish_reason"] == "length" + # --------------------------------------------------------------------------- # ToolCall execution diff --git a/tests/session/test_tool_accumulator.py b/tests/session/test_tool_accumulator.py index 0ef761f6f..ad90eb162 100644 --- a/tests/session/test_tool_accumulator.py +++ b/tests/session/test_tool_accumulator.py @@ -15,8 +15,14 @@ from unittest.mock import AsyncMock, MagicMock, patch -from flocks.session.streaming.tool_accumulator import ToolCallAccumulator -from flocks.session.streaming.stream_events import ToolInputStartEvent, ToolCallEvent +from flocks.session.streaming.tool_accumulator import ( + StreamToolArgumentsTruncatedError, + ToolCallAccumulator, +) +from flocks.session.streaming.stream_events import ( + ToolInputErrorEvent, + ToolInputStartEvent, +) # --------------------------------------------------------------------------- @@ -236,7 +242,7 @@ async def test_flush_invalid_json_sends_invalid_tool(self): assert any(e.tool_name == "invalid" for e in tool_call_events) @pytest.mark.asyncio - async def test_flush_with_length_finish_reason_mentions_truncated(self): + async def test_flush_with_length_finish_reason_marks_input_error_and_raises(self): acc, proc = _make_accumulator() with patch("flocks.session.streaming.tool_accumulator.ToolRegistry") as mock_reg: mock_reg.get_schema.return_value = None @@ -248,15 +254,167 @@ async def test_flush_with_length_finish_reason_mentions_truncated(self): "arguments_str": '{"path": "/tmp/f", "content": "abc', "completed": False, } - await acc.flush_remaining(stream_finish_reason="length") + with pytest.raises(StreamToolArgumentsTruncatedError) as exc_info: + await acc.flush_remaining(stream_finish_reason="length") + + assert exc_info.value.tool_call_id == "call_trunc" + assert exc_info.value.tool_name == "write_file" + assert exc_info.value.finish_reason == "length" + + input_error_events = [ + c.args[0] for c in proc.process_event.call_args_list + if c.args[0].type == "tool-input-error" + ] + assert len(input_error_events) == 1 + assert isinstance(input_error_events[0], ToolInputErrorEvent) + assert input_error_events[0].tool_name == "write_file" + assert "truncated" in input_error_events[0].error.lower() tool_call_events = [ c.args[0] for c in proc.process_event.call_args_list - if c.args[0].type == "tool-call" and c.args[0].tool_name == "invalid" + if c.args[0].type == "tool-call" + ] + assert tool_call_events == [] + + @pytest.mark.asyncio + async def test_flush_with_truncated_repairable_json_does_not_execute_tool(self): + acc, proc = _make_accumulator() + with patch("flocks.session.streaming.tool_accumulator.ToolRegistry") as mock_reg: + mock_reg.get_schema.return_value = None + mock_reg.get.return_value = MagicMock() + acc._accumulator["call_repairable"] = { + "id": "call_repairable", + "name": "write_file", + "arguments_str": '{"path": "/tmp/f", "content": "abc', + "completed": False, + } + + with pytest.raises(StreamToolArgumentsTruncatedError): + await acc.flush_remaining(stream_finish_reason="max_tokens") + + event_types = [c.args[0].type for c in proc.process_event.call_args_list] + assert "tool-input-error" in event_types + assert "tool-call" not in event_types + + @pytest.mark.asyncio + async def test_flush_with_truncated_valid_json_missing_required_does_not_execute_tool(self): + acc, proc = _make_accumulator() + schema = MagicMock() + schema.required = ["path", "content"] + + with patch("flocks.session.streaming.tool_accumulator.ToolRegistry") as mock_reg: + mock_reg.get_schema.return_value = schema + mock_reg.get.return_value = MagicMock() + + await acc.feed_chunk(_make_chunk( + tc_id="call_missing_required", + name="write_file", + arguments='{"path": "/tmp/f"}', + )) + + with pytest.raises(StreamToolArgumentsTruncatedError) as exc_info: + await acc.flush_remaining(stream_finish_reason="length") + + assert exc_info.value.tool_call_id == "call_missing_required" + event_types = [c.args[0].type for c in proc.process_event.call_args_list] + assert event_types == ["tool-input-start", "tool-input-error"] + + @pytest.mark.asyncio + async def test_flush_with_truncated_name_only_tool_marks_input_error(self): + acc, proc = _make_accumulator() + with patch("flocks.session.streaming.tool_accumulator.ToolRegistry") as mock_reg: + mock_reg.get_schema.return_value = None + + await acc.feed_chunk(_make_chunk( + tc_id="call_name_only", + name="write_file", + )) + + with pytest.raises(StreamToolArgumentsTruncatedError) as exc_info: + await acc.flush_remaining(stream_finish_reason="max_tokens") + + assert exc_info.value.tool_call_id == "call_name_only" + events = [c.args[0] for c in proc.process_event.call_args_list] + assert [event.type for event in events] == [ + "tool-input-start", + "tool-input-error", + ] + assert events[-1].input["arguments_preview"] == "" + + @pytest.mark.asyncio + async def test_flush_with_multiple_truncated_tools_marks_all_errors(self): + acc, proc = _make_accumulator() + with patch("flocks.session.streaming.tool_accumulator.ToolRegistry") as mock_reg: + mock_reg.get_schema.return_value = None + mock_reg.get.return_value = MagicMock() + + await acc.feed_chunk(_make_chunk( + index=0, + tc_id="call_done", + name="tool_a", + arguments='{"value": 1}', + )) + await acc.feed_chunk(_make_chunk( + index=1, + tc_id="call_trunc_b", + name="tool_b", + arguments='{"value": "unfinished', + )) + await acc.feed_chunk(_make_chunk( + index=2, + tc_id="call_trunc_c", + name="tool_c", + arguments='{"value": "also unfinished', + )) + + with pytest.raises(StreamToolArgumentsTruncatedError) as exc_info: + await acc.flush_remaining(stream_finish_reason="length") + + assert exc_info.value.tool_call_id == "call_trunc_b" + + events = [c.args[0] for c in proc.process_event.call_args_list] + tool_call_events = [event for event in events if event.type == "tool-call"] + input_error_events = [ + event for event in events if event.type == "tool-input-error" + ] + + assert [event.tool_call_id for event in tool_call_events] == ["call_done"] + assert [event.id for event in input_error_events] == [ + "call_trunc_b", + "call_trunc_c", + ] + assert all("truncated" in event.error.lower() for event in input_error_events) + + @pytest.mark.asyncio + async def test_flush_with_truncation_does_not_execute_later_valid_pending_tool(self): + acc, proc = _make_accumulator() + with patch("flocks.session.streaming.tool_accumulator.ToolRegistry") as mock_reg: + mock_reg.get_schema.return_value = None + mock_reg.get.return_value = MagicMock() + acc._accumulator["call_trunc"] = { + "id": "call_trunc", + "name": "tool_b", + "arguments_str": '{"value": "unfinished', + "completed": False, + "input_started": True, + } + acc._accumulator["call_valid_pending"] = { + "id": "call_valid_pending", + "name": "tool_c", + "arguments_str": '{"value": 1}', + "completed": False, + "input_started": True, + } + + with pytest.raises(StreamToolArgumentsTruncatedError): + await acc.flush_remaining(stream_finish_reason="length") + + events = [c.args[0] for c in proc.process_event.call_args_list] + assert [event.type for event in events] == [ + "tool-input-error", + "tool-input-error", ] - if tool_call_events: - error_msg = tool_call_events[0].input.get("error", "") - assert "truncated" in error_msg.lower() or "length" in error_msg.lower() + assert [event.id for event in events] == ["call_trunc", "call_valid_pending"] @pytest.mark.asyncio async def test_flush_empty_accumulator_does_nothing(self):