From 4a479e66c23381c7a508a76b714c75b27e723db0 Mon Sep 17 00:00:00 2001 From: xiami762 <> Date: Tue, 4 Aug 2026 14:10:19 +0800 Subject: [PATCH 1/2] fix(subagent): preserve interruption state and prevent replay --- flocks/session/runner.py | 34 +++++++-- flocks/tool/agent/delegate_task.py | 2 +- flocks/tool/subagent_result.py | 10 +++ tests/session/test_runner_llm_hooks.py | 93 ++++++++++++++++++++++++ tests/session/test_runner_step.py | 54 ++++++++++++++ tests/tool/test_delegate_task_compat.py | 94 ++++++++++++++++++++++++- 6 files changed, 281 insertions(+), 6 deletions(-) diff --git a/flocks/session/runner.py b/flocks/session/runner.py index 44dfc52d9..0e74e4577 100644 --- a/flocks/session/runner.py +++ b/flocks/session/runner.py @@ -1678,8 +1678,11 @@ async def device_asset_prompt_factory() -> Optional[str]: and not result.content and not result.tool_calls): empty_attempt += 1 unsafe_auto_replay = ( - self._defer_step_errors - and not self._attempt_state.replay_safe + self._attempt_state.tool_execution_started + or ( + self._defer_step_errors + and not self._attempt_state.replay_safe + ) ) if empty_attempt <= MAX_EMPTY_RETRIES and not unsafe_auto_replay: # Record usage for this empty attempt even though we are @@ -1821,7 +1824,15 @@ async def device_asset_prompt_factory() -> Optional[str]: ) else: will_retry = retry_message is not None and error_attempt <= retry_limit - if self._defer_step_errors and not self._attempt_state.replay_safe: + retry_blocked_by_tool_execution = ( + self._attempt_state.tool_execution_started + ) + if retry_blocked_by_tool_execution: + # A tool may already have changed external state. Replaying + # the provider request can emit the same call again with a + # new call id, so no retry mode is safe past this boundary. + will_retry = False + 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 @@ -1861,7 +1872,13 @@ async def device_asset_prompt_factory() -> Optional[str]: continue else: # Error is not retryable, or retry budget exhausted - if retry_message is not None: + if retry_blocked_by_tool_execution: + log.error("runner.step.retry_suppressed", { + **error_log_context, + "attempt": error_attempt, + "reason": "tool_execution_started", + }) + elif retry_message is not None: log.error("runner.step.max_retries_exceeded", { **error_log_context, "attempt": error_attempt, @@ -3482,7 +3499,16 @@ async def _on_tool_execution_start( chunk_counts["tool"] += 1 for tc in chunk_tool_calls: await tool_accumulator.feed_chunk(tc) + except asyncio.CancelledError: + # Foreground delegate tasks own child sessions. Let their + # cancellation/finalization finish before unwinding this step. + await processor.drain_parallel_tool_calls() + raise except Exception as exc: + # A foreground delegate may already be running when the provider + # stream fails. Drain first so the retry layer observes the tool + # side-effect fence and cannot dispatch the same work twice. + await processor.drain_parallel_tool_calls() partial_response = _build_llm_response_payload( content=processor.get_text_content(), reasoning=processor.get_reasoning_content(), diff --git a/flocks/tool/agent/delegate_task.py b/flocks/tool/agent/delegate_task.py index 34280507a..0ee80c8df 100644 --- a/flocks/tool/agent/delegate_task.py +++ b/flocks/tool/agent/delegate_task.py @@ -559,6 +559,6 @@ async def delegate_task_tool( loop_result=result, metadata=forwarder.final_metadata, ) - result_status = "completed" if tool_result.success else "error" + result_status = str((tool_result.metadata or {}).get("status") or ("completed" if tool_result.success else "error")) ctx.metadata({"title": description, "metadata": {**forwarder.final_metadata, "status": result_status}}) return tool_result diff --git a/flocks/tool/subagent_result.py b/flocks/tool/subagent_result.py index 82da7e136..e6676273a 100644 --- a/flocks/tool/subagent_result.py +++ b/flocks/tool/subagent_result.py @@ -46,6 +46,16 @@ async def format_sync_subagent_result( metadata=final_metadata, ) + loop_metadata = getattr(loop_result, "metadata", None) + if isinstance(loop_metadata, dict) and loop_metadata.get("aborted") is True: + final_metadata["status"] = "interrupted" + return ToolResult( + success=False, + error=(f"Sub-agent execution was interrupted.\n\n{_task_metadata_block(session_id)}"), + title=description, + metadata=final_metadata, + ) + last_message = getattr(loop_result, "last_message", None) if not last_message: return ToolResult( diff --git a/tests/session/test_runner_llm_hooks.py b/tests/session/test_runner_llm_hooks.py index 7a3233e36..a2c7ba864 100644 --- a/tests/session/test_runner_llm_hooks.py +++ b/tests/session/test_runner_llm_hooks.py @@ -11,8 +11,10 @@ import flocks.session.runner as runner_mod from flocks.hooks.pipeline import HookBase, HookPipeline from flocks.provider.provider import ChatMessage +from flocks.session.streaming.stream_processor import StreamProcessor from flocks.session.runner import SessionRunner from flocks.session.session import SessionInfo +from flocks.tool.registry import ToolResult def _make_session(session_id: str = "ses_runner_llm_hooks") -> SessionInfo: @@ -317,3 +319,94 @@ async def _gen(): ) assert order == ["before", "provider", "after"] + + +@pytest.mark.asyncio +async def test_call_llm_drains_started_delegate_before_raising_provider_error( + monkeypatch: pytest.MonkeyPatch, +): + runner = _make_runner("ses_runner_delegate_provider_error") + assistant_msg = SimpleNamespace(id="msg_assistant_delegate_provider_error") + agent = SimpleNamespace(name="rex") + delegate_started = asyncio.Event() + release_delegate = asyncio.Event() + + async def _execute_delegate(tool_name, ctx, **_kwargs): + assert tool_name == "delegate_task" + assert ctx.call_id == "call-delegate" + delegate_started.set() + await release_delegate.wait() + return ToolResult(success=True, output="child done") + + monkeypatch.setattr(runner_mod, "langfuse_is_active", lambda: False) + monkeypatch.setattr(runner_mod.Message, "update", AsyncMock(return_value=None)) + monkeypatch.setattr( + runner_mod.HookPipeline, + "has_stage_handlers", + AsyncMock(return_value=False), + ) + monkeypatch.setattr( + "flocks.provider.options.build_provider_options", + lambda provider_id, model_id: {}, + ) + monkeypatch.setattr( + "flocks.session.streaming.stream_processor.Message.store_part", + AsyncMock(return_value=None), + ) + monkeypatch.setattr( + "flocks.session.streaming.stream_processor.Message.parts", + AsyncMock(return_value=[]), + ) + monkeypatch.setattr( + "flocks.session.streaming.stream_processor.ToolRegistry.execute", + _execute_delegate, + ) + monkeypatch.setattr( + StreamProcessor, + "_resolve_sandbox_meta", + AsyncMock(return_value={"blocked": False, "error": None, "extra": {}}), + ) + + class _Provider: + def chat_stream(self, **_kwargs): + async def _gen(): + yield SimpleNamespace( + delta="", + reasoning=None, + tool_calls=[ + { + "index": 0, + "id": "call-delegate", + "function": { + "name": "delegate_task", + "arguments": ('{"subagent_type":"explore","prompt":"inspect the failure"}'), + }, + } + ], + event_type=None, + finish_reason=None, + usage=None, + ) + await delegate_started.wait() + raise RuntimeError("provider stream failed after delegate start") + + return _gen() + + call_task = asyncio.create_task( + runner._call_llm( + provider=_Provider(), + messages=[ChatMessage(role="user", content="delegate the investigation")], + tools=[], + agent=agent, + assistant_msg=assistant_msg, + ) + ) + await asyncio.wait_for(delegate_started.wait(), timeout=1) + await asyncio.sleep(0) + provider_error_waited_for_delegate = not call_task.done() + + release_delegate.set() + with pytest.raises(RuntimeError, match="provider stream failed after delegate start"): + await call_task + + assert provider_error_waited_for_delegate is True diff --git a/tests/session/test_runner_step.py b/tests/session/test_runner_step.py index 82fb6aaec..a57e73b33 100644 --- a/tests/session/test_runner_step.py +++ b/tests/session/test_runner_step.py @@ -2931,6 +2931,60 @@ async def test_process_step_retries_empty_transport_exception(monkeypatch): runner.callbacks.on_error.assert_not_awaited() +@pytest.mark.asyncio +async def test_process_step_does_not_retry_after_tool_execution_started(monkeypatch): + runner = _make_runner("ses_runner_tool_side_effect_no_retry") + runner.callbacks = RunnerCallbacks(on_error=AsyncMock()) + + last_user = UserMessageInfo( + id="msg_user_tool_side_effect_no_retry", + sessionID=runner.session.id, + role="user", + time={"created": 1_000}, + agent="rex", + model={"providerID": "openai", "modelID": "gpt-5"}, + ) + agent = SimpleNamespace(name="rex", steps=None, mode="primary", prompt="", tools=[]) + provider = MagicMock() + provider.is_configured.return_value = True + assistant_msg = SimpleNamespace(id="msg_assistant_tool_side_effect_no_retry") + call_count = 0 + sleep_mock = AsyncMock(return_value=None) + + async def _call_llm(*_args, **_kwargs): + nonlocal call_count + call_count += 1 + runner._attempt_state.tool_execution_started = True + raise httpcore.ReadError() + + 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_prompts", + 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, "create", AsyncMock(return_value=assistant_msg)) + monkeypatch.setattr(runner_mod.Message, "update", AsyncMock(return_value=None)) + monkeypatch.setattr(runner_mod.SessionRetry, "sleep", sleep_mock) + monkeypatch.setattr(runner, "_call_llm", _call_llm) + + result = await runner._process_step([last_user], last_user) + + assert call_count == 1 + assert result.action == "stop" + assert result.error == runner_mod.CONNECTION_ERROR_DISPLAY_MESSAGE + sleep_mock.assert_not_awaited() + + @pytest.mark.asyncio async def test_process_step_uses_default_max_steps_when_agent_steps_missing(monkeypatch): runner = _make_runner("ses_runner_default_max_steps") diff --git a/tests/tool/test_delegate_task_compat.py b/tests/tool/test_delegate_task_compat.py index 59be3c16f..644718c10 100644 --- a/tests/tool/test_delegate_task_compat.py +++ b/tests/tool/test_delegate_task_compat.py @@ -1,5 +1,5 @@ from types import SimpleNamespace -from unittest.mock import AsyncMock, patch +from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -257,3 +257,95 @@ async def test_delegate_task_sync_continue_fails_when_last_message_missing(self) assert result.metadata["sessionId"] == "ses-child" assert result.metadata["emptyOutput"] is True assert "without producing a final assistant message" in (result.output or "") + + @pytest.mark.asyncio + async def test_delegate_task_sync_continue_reports_normalized_abort(self): + session = SimpleNamespace( + id="ses-child-aborted", + agent="asset-survey", + ) + + with ( + patch( + "flocks.tool.agent.delegate_task.Session.get_by_id", + AsyncMock(return_value=session), + ), + patch("flocks.tool.agent.delegate_task.Message.create", AsyncMock()), + patch( + "flocks.tool.agent.delegate_task.SessionLoop.run", + AsyncMock( + return_value=SimpleNamespace( + action="stop", + error=None, + last_message=None, + metadata={"aborted": True}, + ) + ), + ), + ): + result = await ToolRegistry.execute( + "delegate_task", + ctx=_make_ctx(), + session_id="ses-child-aborted", + prompt="Continue investigating", + ) + + assert result.success is False + assert result.metadata["sessionId"] == "ses-child-aborted" + assert result.metadata["status"] == "interrupted" + assert "Sub-agent execution was interrupted" in (result.error or "") + assert "session_id: ses-child-aborted" in (result.error or "") + + @pytest.mark.asyncio + async def test_new_delegate_reports_normalized_abort_to_parent_metadata(self): + ctx = _make_ctx() + metadata_callback = MagicMock() + ctx.metadata = metadata_callback + parent_session = SimpleNamespace( + id="test-session", + project_id="proj", + directory="/tmp/project", + provider=None, + model=None, + ) + child_session = SimpleNamespace(id="ses-new-child-aborted") + + with ( + patch( + "flocks.tool.agent.delegate_task._find_completed_delegate", + AsyncMock(return_value=None), + ), + patch("flocks.tool.agent.delegate_task.is_delegatable", return_value=True), + patch( + "flocks.tool.agent.delegate_task.Session.get_by_id", + AsyncMock(return_value=parent_session), + ), + patch( + "flocks.tool.agent.delegate_task.Session.create", + AsyncMock(return_value=child_session), + ), + patch("flocks.tool.agent.delegate_task.Message.create", AsyncMock()), + patch( + "flocks.tool.agent.delegate_task.SessionLoop.run", + AsyncMock( + return_value=SimpleNamespace( + action="stop", + error=None, + last_message=None, + metadata={"aborted": True}, + ) + ), + ), + ): + result = await ToolRegistry.execute( + "delegate_task", + ctx=ctx, + subagent_type="asset-survey", + prompt="Investigate the interruption", + ) + + assert result.success is False + assert result.metadata["sessionId"] == "ses-new-child-aborted" + assert result.metadata["status"] == "interrupted" + final_parent_metadata = metadata_callback.call_args_list[-1].args[0] + assert final_parent_metadata["metadata"]["status"] == "interrupted" From a7217cfa8808eafc5d519395f9485f7b6937005d Mon Sep 17 00:00:00 2001 From: xiami762 <> Date: Tue, 4 Aug 2026 14:23:21 +0800 Subject: [PATCH 2/2] fix(subagent): preserve session metadata on cancellation --- flocks/session/streaming/stream_processor.py | 14 +++++- tests/session/test_stream_processor.py | 53 +++++++++++++++++++- 2 files changed, 64 insertions(+), 3 deletions(-) diff --git a/flocks/session/streaming/stream_processor.py b/flocks/session/streaming/stream_processor.py index ca1666474..05397ab4a 100644 --- a/flocks/session/streaming/stream_processor.py +++ b/flocks/session/streaming/stream_processor.py @@ -11,7 +11,7 @@ import time as _time from datetime import datetime from typing import Dict, Any, Optional, List, AsyncIterator, Callable, Awaitable -from dataclasses import dataclass +from dataclasses import dataclass, field from flocks.utils.log import Log from flocks.utils.id import Identifier @@ -91,6 +91,7 @@ class ToolCallState: status: str = "pending" # "pending", "running", "completed", "error" output: Optional[str] = None error: Optional[str] = None + metadata: Dict[str, Any] = field(default_factory=dict) class StreamProcessor: @@ -773,6 +774,7 @@ def _cb(metadata: Dict[str, Any]): if _finished[0]: return snapshot = copy.deepcopy(metadata) + tool_state.metadata = snapshot state_dict = { "status": "running", "input": _input, @@ -1127,6 +1129,12 @@ async def _finalize_interrupted_tool_call( ) -> None: """Emit and persist the terminal state for an interrupted tool call.""" interrupt_msg = "Tool execution was interrupted" + interrupted_metadata = { + **tool_state.metadata, + "status": "interrupted", + "interrupted": True, + } + tool_state.metadata = interrupted_metadata log.info("stream.tool_call.cancelled", { "tool_call_id": tool_call_id, "tool_name": tool_name, @@ -1136,7 +1144,7 @@ async def _finalize_interrupted_tool_call( interrupted_result = ToolResult( success=False, error=interrupt_msg, - metadata={"interrupted": True}, + metadata=interrupted_metadata, ) await self._run_tool_after_hook( tool_name=tool_name, @@ -1169,6 +1177,7 @@ async def _finalize_interrupted_tool_call( status="error", input=tool_input, error=interrupt_msg, + metadata=interrupted_metadata, time={"start": tool_start_time, "end": tool_end_time}, ) error_part = ToolPart( @@ -1201,6 +1210,7 @@ async def _finalize_interrupted_tool_call( "status": "error", "input": tool_input, "error": interrupt_msg, + "metadata": interrupted_metadata, "time": { "start": tool_start_time, "end": tool_end_time, diff --git a/tests/session/test_stream_processor.py b/tests/session/test_stream_processor.py index 914f16d30..3776136ab 100644 --- a/tests/session/test_stream_processor.py +++ b/tests/session/test_stream_processor.py @@ -31,7 +31,7 @@ ToolCallEvent, ToolInputStartEvent, ) -from flocks.session.message import MessageRole +from flocks.session.message import MessageRole, ToolStateError # --------------------------------------------------------------------------- @@ -809,6 +809,57 @@ async def _cancelled_execute(*, tool_name, ctx, **kwargs): await asyncio.sleep(0) assert len(event_callback.await_args_list) == baseline_calls + @pytest.mark.asyncio + async def test_cancelled_tool_preserves_running_metadata_in_error_state(self): + proc = _make_processor() + store_part = AsyncMock() + + async def _cancelled_execute(*, tool_name, ctx, **kwargs): + ctx.metadata({ + "title": "Inspect child", + "metadata": { + "sessionId": "ses_child_cancelled", + "status": "running", + }, + }) + await asyncio.sleep(0) + raise asyncio.CancelledError() + + with ( + patch( + "flocks.session.streaming.stream_processor.Message.store_part", + new=store_part, + ), + patch( + "flocks.session.streaming.stream_processor.Message.update_part", + new=AsyncMock(), + ), + patch( + "flocks.session.streaming.stream_processor.ToolRegistry.execute", + new=AsyncMock(side_effect=_cancelled_execute), + ), + ): + await proc.process_event( + ToolInputStartEvent(id="tc_cancel_metadata", tool_name="run_workflow") + ) + with pytest.raises(asyncio.CancelledError): + await proc.process_event( + ToolCallEvent( + tool_call_id="tc_cancel_metadata", + tool_name="run_workflow", + input={"workflow": "wf.json"}, + ) + ) + + final_part = store_part.await_args_list[-1].args[2] + assert isinstance(final_part.state, ToolStateError) + assert final_part.state.metadata == { + "title": "Inspect child", + "sessionId": "ses_child_cancelled", + "status": "interrupted", + "interrupted": True, + } + @pytest.mark.asyncio async def test_completed_tool_cancels_pending_running_metadata_tasks(self): proc = _make_processor(event_callback=AsyncMock())