diff --git a/agent/tool_dispatch_helpers.py b/agent/tool_dispatch_helpers.py index 1b6cae98ac66..4360f60f7eb9 100644 --- a/agent/tool_dispatch_helpers.py +++ b/agent/tool_dispatch_helpers.py @@ -102,50 +102,118 @@ def _is_mcp_tool_parallel_safe(tool_name: str) -> bool: return False -def _should_parallelize_tool_batch(tool_calls) -> bool: - """Return True when a tool-call batch is safe to run concurrently.""" - if len(tool_calls) <= 1: - return False +def _plan_tool_batch_segments(tool_calls) -> List[tuple]: + """Split a tool-call batch into ordered ``(kind, calls)`` segments. - tool_names = [tc.function.name for tc in tool_calls] - if any(name in _NEVER_PARALLEL_TOOLS for name in tool_names): - return False + ``kind`` is ``"parallel"`` (a maximal contiguous run of parallel-safe + calls) or ``"sequential"`` (one or more barrier calls that must run + in-order on the sequential path). Segments preserve the model's + original call order exactly — a later call never crosses an earlier + barrier — so tool-result ordering and side-effect boundaries are + identical to fully-sequential execution. The per-call safety rules + are the same ones the old all-or-nothing gate applied to the whole + batch: + * ``_NEVER_PARALLEL_TOOLS`` (interactive tools) → barrier. + * Unparseable / non-dict arguments → barrier. + * Path-scoped tools (``read_file``/``write_file``/``patch``) join a + parallel run only when their target path does not overlap another + path already reserved in the same run; an overlap closes the run so + the conflicting call starts a NEW run after the first completes. + * Anything not in ``_PARALLEL_SAFE_TOOLS`` and not an opted-in MCP + tool → barrier. + + Parallel runs shorter than two calls are demoted to sequential (no + concurrency win, and the sequential executor owns the richer inline + dispatch), and adjacent sequential segments are merged. + """ + segments: list[list] = [] # [kind, calls] pairs, normalized to tuples on return + current: list = [] reserved_paths: list[Path] = [] + + def _close_parallel() -> None: + nonlocal current, reserved_paths + if current: + segments.append(["parallel", current]) + current = [] + reserved_paths = [] + + def _add_sequential(tc) -> None: + _close_parallel() + if segments and segments[-1][0] == "sequential": + segments[-1][1].append(tc) + else: + segments.append(["sequential", [tc]]) + for tool_call in tool_calls: tool_name = tool_call.function.name + + if tool_name in _NEVER_PARALLEL_TOOLS: + _add_sequential(tool_call) + continue + try: function_args = json.loads(tool_call.function.arguments) except Exception: logging.debug( - "Could not parse args for %s — defaulting to sequential; raw=%s", + "Could not parse args for %s — treating as sequential barrier; raw=%s", tool_name, tool_call.function.arguments[:200], ) - return False + _add_sequential(tool_call) + continue if not isinstance(function_args, dict): logging.debug( - "Non-dict args for %s (%s) — defaulting to sequential", + "Non-dict args for %s (%s) — treating as sequential barrier", tool_name, type(function_args).__name__, ) - return False + _add_sequential(tool_call) + continue if tool_name in _PATH_SCOPED_TOOLS: scoped_path = _extract_parallel_scope_path(tool_name, function_args) if scoped_path is None: - return False + _add_sequential(tool_call) + continue if any(_paths_overlap(scoped_path, existing) for existing in reserved_paths): - return False + # Same-subtree conflict inside this run: close it so this + # call starts a fresh run AFTER the conflicting one lands. + _close_parallel() reserved_paths.append(scoped_path) + current.append(tool_call) continue - if tool_name not in _PARALLEL_SAFE_TOOLS: - # Check if it's an MCP tool from a server that opted into parallel calls. - if not _is_mcp_tool_parallel_safe(tool_name): - return False + if tool_name in _PARALLEL_SAFE_TOOLS or _is_mcp_tool_parallel_safe(tool_name): + current.append(tool_call) + continue - return True + _add_sequential(tool_call) + + _close_parallel() + + normalized: list[list] = [] + for kind, calls in segments: + if kind == "parallel" and len(calls) < 2: + kind = "sequential" + if normalized and normalized[-1][0] == "sequential" and kind == "sequential": + normalized[-1][1].extend(calls) + else: + normalized.append([kind, calls]) + return [(kind, calls) for kind, calls in normalized] + + +def _should_parallelize_tool_batch(tool_calls) -> bool: + """Return True when the WHOLE tool-call batch is safe to run concurrently. + + Thin view over ``_plan_tool_batch_segments`` kept for callers/tests that + only care about the homogeneous case: True iff the planner produces a + single all-parallel segment. + """ + if len(tool_calls) <= 1: + return False + segments = _plan_tool_batch_segments(tool_calls) + return len(segments) == 1 and segments[0][0] == "parallel" def _extract_parallel_scope_path(tool_name: str, function_args: dict) -> Optional[Path]: @@ -542,6 +610,7 @@ __all__ = [ "_DESTRUCTIVE_PATTERNS", "_REDIRECT_OVERWRITE", "_is_destructive_command", + "_plan_tool_batch_segments", "_should_parallelize_tool_batch", "_extract_parallel_scope_path", "_paths_overlap", diff --git a/agent/tool_executor.py b/agent/tool_executor.py index ac505c6d8299..a5871442b449 100644 --- a/agent/tool_executor.py +++ b/agent/tool_executor.py @@ -36,6 +36,7 @@ from agent.tool_dispatch_helpers import ( _is_multimodal_tool_result, _multimodal_text_summary, _append_subdir_hint_to_multimodal, + _plan_tool_batch_segments, make_tool_result_message, ) from tools.terminal_tool import ( @@ -322,11 +323,15 @@ def _run_agent_tool_execution_middleware( return result, observed_args -def execute_tool_calls_concurrent(agent, assistant_message, messages: list, effective_task_id: str, api_call_count: int = 0) -> None: +def execute_tool_calls_concurrent(agent, assistant_message, messages: list, effective_task_id: str, api_call_count: int = 0, *, finalize: bool = True) -> None: """Execute multiple tool calls concurrently using a thread pool. Results are collected in the original tool-call order and appended to messages so the API sees them in the expected sequence. + + ``finalize=False`` skips the end-of-batch aggregate budget enforcement + and /steer injection — used when this call is one segment of a larger + mixed batch and the segmented dispatcher owns the turn-end work. """ tool_calls = assistant_message.tool_calls num_tools = len(tool_calls) @@ -1006,7 +1011,7 @@ def execute_tool_calls_concurrent(agent, assistant_message, messages: list, effe # ── Per-turn aggregate budget enforcement ───────────────────────── num_tools = len(parsed_calls) - if num_tools > 0: + if finalize and num_tools > 0: turn_tool_msgs = messages[-num_tools:] enforce_turn_budget(turn_tool_msgs, env=get_active_env(effective_task_id), config=_tool_budget) @@ -1014,13 +1019,18 @@ def execute_tool_calls_concurrent(agent, assistant_message, messages: list, effe # Append any pending user steer text to the last tool result so the # agent sees it on its next iteration. Runs AFTER budget enforcement # so the steer marker is never truncated. See steer() for details. - if num_tools > 0: + if finalize and num_tools > 0: agent._apply_pending_steer_to_tool_results(messages, num_tools) -def execute_tool_calls_sequential(agent, assistant_message, messages: list, effective_task_id: str, api_call_count: int = 0) -> None: - """Execute tool calls sequentially (original behavior). Used for single calls or interactive tools.""" +def execute_tool_calls_sequential(agent, assistant_message, messages: list, effective_task_id: str, api_call_count: int = 0, *, finalize: bool = True) -> None: + """Execute tool calls sequentially (original behavior). Used for single calls or interactive tools. + + ``finalize=False`` skips the end-of-batch aggregate budget enforcement + and /steer injection — used when this call is one segment of a larger + mixed batch and the segmented dispatcher owns the turn-end work. + """ # Resolve the context-scaled tool-output budget once per turn. _tool_budget = _budget_for_agent(agent) for i, tool_call in enumerate(assistant_message.tool_calls, 1): @@ -1716,19 +1726,73 @@ def execute_tool_calls_sequential(agent, assistant_message, messages: list, effe # ── Per-turn aggregate budget enforcement ───────────────────────── num_tools_seq = len(assistant_message.tool_calls) - if num_tools_seq > 0: + if finalize and num_tools_seq > 0: enforce_turn_budget(messages[-num_tools_seq:], env=get_active_env(effective_task_id), config=_tool_budget) # ── /steer injection ────────────────────────────────────────────── # See _execute_tool_calls_parallel for the rationale. Same hook, # applied to sequential execution as well. - if num_tools_seq > 0: + if finalize and num_tools_seq > 0: agent._apply_pending_steer_to_tool_results(messages, num_tools_seq) +def execute_tool_calls_segmented(agent, assistant_message, messages: list, effective_task_id: str, api_call_count: int = 0, segments=None) -> None: + """Execute a mixed tool-call batch as ordered parallel/sequential segments. + + ``segments`` is the ``(kind, calls)`` plan from + ``_plan_tool_batch_segments``: maximal contiguous runs of parallel-safe + calls execute on the concurrent path, barrier calls on the sequential + path, strictly in the model's original call order. Because segments are + contiguous, every tool result is still appended one-per-call in emission + order and no call ever starts before an earlier barrier finishes — + identical ordering and side-effect boundaries to fully-sequential + execution, with I/O parallelism recovered inside the safe runs. + + Turn-end work (aggregate budget enforcement + /steer injection) is done + once here for the WHOLE batch; the per-segment executor calls run with + ``finalize=False`` so a multi-segment turn cannot multiply the budget or + truncate a steer marker. + + Interrupt semantics: each segment executor already checks + ``agent._interrupt_requested`` up front and appends a cancelled/skipped + result per call, so an interrupt during segment *k* drains segments + *k+1..n* without executing them while preserving one result per + tool_call_id. + """ + from types import SimpleNamespace + + if segments is None: + segments = _plan_tool_batch_segments(assistant_message.tool_calls) + + for kind, calls in segments: + segment_message = SimpleNamespace(tool_calls=list(calls)) + if kind == "parallel": + execute_tool_calls_concurrent( + agent, segment_message, messages, effective_task_id, api_call_count, + finalize=False, + ) + else: + execute_tool_calls_sequential( + agent, segment_message, messages, effective_task_id, api_call_count, + finalize=False, + ) + + # ── Whole-turn finalize (budget + /steer) ───────────────────────── + total_tools = len(assistant_message.tool_calls) + if total_tools > 0: + _tool_budget = _budget_for_agent(agent) + enforce_turn_budget( + messages[-total_tools:], + env=get_active_env(effective_task_id), + config=_tool_budget, + ) + agent._apply_pending_steer_to_tool_results(messages, total_tools) + + __all__ = [ "execute_tool_calls_concurrent", "execute_tool_calls_sequential", + "execute_tool_calls_segmented", ] diff --git a/run_agent.py b/run_agent.py index 31cca895592f..504402459c21 100644 --- a/run_agent.py +++ b/run_agent.py @@ -198,7 +198,7 @@ from agent.trajectory import ( save_trajectory as _save_trajectory_to_file, ) from agent.tool_dispatch_helpers import ( - _should_parallelize_tool_batch, + _should_parallelize_tool_batch, # noqa: F401 # re-exported for tests that `from run_agent import _should_parallelize_tool_batch` _is_destructive_command, # noqa: F401 # re-exported for tests that access `run_agent._is_destructive_command` _extract_parallel_scope_path, # noqa: F401 # re-exported for tests that `from run_agent import _extract_parallel_scope_path` _paths_overlap, # noqa: F401 # re-exported for tests that `from run_agent import _paths_overlap` @@ -5714,22 +5714,41 @@ class AIAgent: def _execute_tool_calls(self, assistant_message, messages: list, effective_task_id: str, api_call_count: int = 0) -> None: """Execute tool calls from the assistant message and append results to messages. - Dispatches to concurrent execution only for batches that look - independent: read-only tools may always share the parallel path, while - file reads/writes may do so only when their target paths do not overlap. + The segment planner splits the batch into maximal contiguous runs of + parallel-safe calls (read-only tools, non-overlapping file targets, + opted-in MCP tools) separated by sequential barriers (interactive, + unsafe, or unrecognized tools). Homogeneous batches keep their + original single-path dispatch; mixed batches execute segment by + segment in emission order so safe subsets still run concurrently + while side-effect ordering is preserved. """ tool_calls = assistant_message.tool_calls # Allow _vprint during tool execution even with stream consumers self._executing_tools = True try: - if not _should_parallelize_tool_batch(tool_calls): + if len(tool_calls) <= 1: return self._execute_tool_calls_sequential( assistant_message, messages, effective_task_id, api_call_count ) - return self._execute_tool_calls_concurrent( - assistant_message, messages, effective_task_id, api_call_count + from agent.tool_dispatch_helpers import _plan_tool_batch_segments + segments = _plan_tool_batch_segments(tool_calls) + + if len(segments) == 1: + kind = segments[0][0] + if kind == "parallel": + return self._execute_tool_calls_concurrent( + assistant_message, messages, effective_task_id, api_call_count + ) + return self._execute_tool_calls_sequential( + assistant_message, messages, effective_task_id, api_call_count + ) + + from agent.tool_executor import execute_tool_calls_segmented + return execute_tool_calls_segmented( + self, assistant_message, messages, effective_task_id, api_call_count, + segments=segments, ) finally: self._executing_tools = False diff --git a/tests/run_agent/test_tool_batch_segmentation.py b/tests/run_agent/test_tool_batch_segmentation.py new file mode 100644 index 000000000000..9cde31479d00 --- /dev/null +++ b/tests/run_agent/test_tool_batch_segmentation.py @@ -0,0 +1,400 @@ +"""Segment-aware mixed tool-batch dispatch. + +A model response containing several parallel-safe reads plus one unsafe +tool used to lose ALL concurrency: `_should_parallelize_tool_batch` was +all-or-nothing, so one barrier call forced the entire batch onto the +sequential path. `_plan_tool_batch_segments` now splits the batch into +ordered segments — maximal contiguous runs of parallel-safe calls execute +concurrently, barrier calls sequentially — while preserving: + + * model tool-result ordering (one result per call, in emission order), + * side-effect boundaries (no call starts before an earlier barrier ends). +""" + +import json +import threading +import time +import uuid +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest + +from run_agent import AIAgent +from agent.tool_dispatch_helpers import ( + _plan_tool_batch_segments, + _should_parallelize_tool_batch, +) + + +def _tc(name="web_search", arguments="{}", call_id=None): + return SimpleNamespace( + id=call_id or f"call_{uuid.uuid4().hex[:8]}", + type="function", + function=SimpleNamespace(name=name, arguments=arguments), + ) + + +def _kinds(segments): + return [kind for kind, _ in segments] + + +def _flatten_ids(segments): + return [tc.id for _, calls in segments for tc in calls] + + +# --------------------------------------------------------------------------- +# Planner unit tests +# --------------------------------------------------------------------------- + + +class TestPlanToolBatchSegments: + def test_all_safe_batch_is_single_parallel_segment(self): + calls = [_tc("web_search"), _tc("read_file", '{"path":"a.py"}'), _tc("web_extract")] + segments = _plan_tool_batch_segments(calls) + assert _kinds(segments) == ["parallel"] + assert _flatten_ids(segments) == [c.id for c in calls] + + def test_three_safe_reads_plus_trailing_unsafe_keeps_reads_parallel(self): + """The headline case: 3 safe reads + 1 unsafe tool must NOT go fully sequential.""" + calls = [ + _tc("web_search", call_id="r1"), + _tc("web_search", call_id="r2"), + _tc("read_file", '{"path":"a.py"}', call_id="r3"), + _tc("terminal", '{"command":"echo hi"}', call_id="b1"), + ] + segments = _plan_tool_batch_segments(calls) + assert _kinds(segments) == ["parallel", "sequential"] + assert [tc.id for tc in segments[0][1]] == ["r1", "r2", "r3"] + assert [tc.id for tc in segments[1][1]] == ["b1"] + + def test_barrier_in_middle_splits_runs_and_preserves_order(self): + calls = [ + _tc("web_search", call_id="r1"), + _tc("web_search", call_id="r2"), + _tc("terminal", '{"command":"make"}', call_id="b1"), + _tc("web_search", call_id="r3"), + _tc("web_search", call_id="r4"), + ] + segments = _plan_tool_batch_segments(calls) + assert _kinds(segments) == ["parallel", "sequential", "parallel"] + assert _flatten_ids(segments) == ["r1", "r2", "b1", "r3", "r4"] + + def test_single_safe_call_after_barrier_is_demoted_and_merged(self): + # parallel run of 1 gains nothing — demote to sequential and merge + # with the adjacent barrier segment. + calls = [ + _tc("web_search", call_id="r1"), + _tc("web_search", call_id="r2"), + _tc("terminal", '{"command":"make"}', call_id="b1"), + _tc("web_search", call_id="r3"), + ] + segments = _plan_tool_batch_segments(calls) + assert _kinds(segments) == ["parallel", "sequential"] + assert [tc.id for tc in segments[1][1]] == ["b1", "r3"] + + def test_adjacent_barriers_merge_into_one_sequential_segment(self): + calls = [ + _tc("terminal", '{"command":"a"}', call_id="b1"), + _tc("terminal", '{"command":"b"}', call_id="b2"), + _tc("web_search", call_id="r1"), + _tc("web_search", call_id="r2"), + ] + segments = _plan_tool_batch_segments(calls) + assert _kinds(segments) == ["sequential", "parallel"] + assert [tc.id for tc in segments[0][1]] == ["b1", "b2"] + + def test_never_parallel_tool_is_a_barrier(self): + calls = [ + _tc("web_search", call_id="r1"), + _tc("web_search", call_id="r2"), + _tc("clarify", '{"question":"?"}', call_id="c1"), + ] + segments = _plan_tool_batch_segments(calls) + assert _kinds(segments) == ["parallel", "sequential"] + assert [tc.id for tc in segments[1][1]] == ["c1"] + + def test_malformed_args_call_is_a_barrier_not_a_batch_poison(self): + calls = [ + _tc("web_search", call_id="r1"), + _tc("web_search", call_id="r2"), + _tc("web_search", "{not json", call_id="bad"), + _tc("web_search", call_id="r3"), + _tc("web_search", call_id="r4"), + ] + segments = _plan_tool_batch_segments(calls) + assert _kinds(segments) == ["parallel", "sequential", "parallel"] + assert [tc.id for tc in segments[1][1]] == ["bad"] + + def test_non_dict_args_call_is_a_barrier(self): + calls = [ + _tc("web_search", call_id="r1"), + _tc("web_search", call_id="r2"), + _tc("web_search", '"just a string"', call_id="bad"), + ] + segments = _plan_tool_batch_segments(calls) + assert _kinds(segments) == ["parallel", "sequential"] + + def test_overlapping_paths_split_across_segments(self, tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) + calls = [ + _tc("read_file", '{"path":"a.py"}', call_id="w1"), + _tc("web_search", call_id="r1"), + _tc("write_file", '{"path":"a.py","content":"x"}', call_id="w2"), + _tc("web_search", call_id="r2"), + ] + segments = _plan_tool_batch_segments(calls) + # w2 conflicts with w1 → closes the first run; w2+r2 form the second. + assert _kinds(segments) == ["parallel", "parallel"] + assert [tc.id for tc in segments[0][1]] == ["w1", "r1"] + assert [tc.id for tc in segments[1][1]] == ["w2", "r2"] + # Order and completeness preserved. + assert _flatten_ids(segments) == ["w1", "r1", "w2", "r2"] + + def test_path_scoped_tool_without_path_is_a_barrier(self): + calls = [ + _tc("read_file", "{}", call_id="nopath"), + _tc("web_search", call_id="r1"), + _tc("web_search", call_id="r2"), + ] + segments = _plan_tool_batch_segments(calls) + assert _kinds(segments) == ["sequential", "parallel"] + + def test_flattened_segments_always_preserve_emission_order(self): + calls = [ + _tc("terminal", '{"command":"x"}', call_id="b1"), + _tc("web_search", call_id="r1"), + _tc("clarify", '{"question":"?"}', call_id="c1"), + _tc("read_file", '{"path":"a.py"}', call_id="r2"), + _tc("read_file", '{"path":"b.py"}', call_id="r3"), + ] + segments = _plan_tool_batch_segments(calls) + assert _flatten_ids(segments) == ["b1", "r1", "c1", "r2", "r3"] + + +class TestShouldParallelizeBackwardCompat: + """The boolean gate is now a view over the planner — same answers as before.""" + + def test_single_call_is_sequential(self): + assert not _should_parallelize_tool_batch([_tc("web_search")]) + + def test_all_safe_batch_is_parallel(self): + assert _should_parallelize_tool_batch([_tc("web_search"), _tc("web_extract")]) + + def test_mixed_batch_is_not_wholly_parallel(self): + assert not _should_parallelize_tool_batch( + [_tc("web_search"), _tc("terminal", '{"command":"ls"}')] + ) + + def test_clarify_anywhere_blocks_whole_batch_parallelism(self): + assert not _should_parallelize_tool_batch( + [_tc("web_search"), _tc("clarify", '{"question":"?"}')] + ) + + +# --------------------------------------------------------------------------- +# Dispatcher integration +# --------------------------------------------------------------------------- + + +def _make_tool_defs(*names: str) -> list: + return [ + { + "type": "function", + "function": { + "name": n, + "description": f"{n} tool", + "parameters": {"type": "object", "properties": {}}, + }, + } + for n in names + ] + + +@pytest.fixture() +def agent(): + with ( + patch( + "run_agent.get_tool_definitions", + return_value=_make_tool_defs("web_search", "terminal"), + ), + patch("run_agent.check_toolset_requirements", return_value={}), + patch("run_agent.OpenAI"), + ): + a = AIAgent( + api_key="test-key-1234567890", + base_url="https://openrouter.ai/api/v1", + quiet_mode=True, + skip_context_files=True, + skip_memory=True, + ) + a.client = MagicMock() + return a + + +class TestSegmentedDispatchIntegration: + def test_mixed_batch_runs_safe_prefix_concurrently_and_barrier_after(self, agent): + """Two web_search calls must overlap in time; terminal must start only + after both finish; results land in the model's emission order.""" + calls = [ + _tc("web_search", '{"query":"a"}', call_id="s1"), + _tc("web_search", '{"query":"b"}', call_id="s2"), + _tc("terminal", '{"command":"echo done"}', call_id="t1"), + ] + msg = SimpleNamespace(content="", tool_calls=calls) + messages = [] + + rendezvous = threading.Barrier(2, timeout=10) + events = [] + events_lock = threading.Lock() + + def fake_handle(name, args, task_id, **kwargs): + with events_lock: + events.append(("start", name, kwargs["tool_call_id"])) + if name == "web_search": + # Both searches must be in flight at once to pass this + # barrier — proves genuine concurrency for the safe prefix. + rendezvous.wait() + with events_lock: + events.append(("end", name, kwargs["tool_call_id"])) + return json.dumps({"ok": name}) + + with patch("run_agent.handle_function_call", side_effect=fake_handle): + agent._execute_tool_calls(msg, messages, "task-1") + + # One result per call, in emission order. + assert [m["tool_call_id"] for m in messages] == ["s1", "s2", "t1"] + assert all(m["role"] == "tool" for m in messages) + + # The barrier (terminal) started only after BOTH searches ended. + terminal_start = events.index(("start", "terminal", "t1")) + search_ends = [ + i for i, e in enumerate(events) if e[0] == "end" and e[1] == "web_search" + ] + assert len(search_ends) == 2 + assert all(i < terminal_start for i in search_ends) + + def test_mixed_batch_preserves_order_with_barrier_in_middle(self, agent): + calls = [ + _tc("web_search", '{"query":"a"}', call_id="s1"), + _tc("web_search", '{"query":"b"}', call_id="s2"), + _tc("terminal", '{"command":"touch x"}', call_id="t1"), + _tc("web_search", '{"query":"c"}', call_id="s3"), + _tc("web_search", '{"query":"d"}', call_id="s4"), + ] + msg = SimpleNamespace(content="", tool_calls=calls) + messages = [] + executed = [] + lock = threading.Lock() + + def fake_handle(name, args, task_id, **kwargs): + with lock: + executed.append(kwargs["tool_call_id"]) + return json.dumps({"ok": True}) + + with patch("run_agent.handle_function_call", side_effect=fake_handle): + agent._execute_tool_calls(msg, messages, "task-1") + + assert [m["tool_call_id"] for m in messages] == ["s1", "s2", "t1", "s3", "s4"] + # Barrier ordering: t1 executed after {s1,s2} and before {s3,s4}. + t1_pos = executed.index("t1") + assert {"s1", "s2"} == set(executed[:t1_pos]) + assert {"s3", "s4"} == set(executed[t1_pos + 1:]) + + def test_homogeneous_safe_batch_still_uses_plain_concurrent_path(self, agent): + calls = [_tc("web_search", '{"query":"a"}'), _tc("web_search", '{"query":"b"}')] + msg = SimpleNamespace(content="", tool_calls=calls) + + with ( + patch.object(agent, "_execute_tool_calls_concurrent") as conc, + patch.object(agent, "_execute_tool_calls_sequential") as seq, + ): + agent._execute_tool_calls(msg, [], "task-1") + + conc.assert_called_once() + seq.assert_not_called() + + def test_homogeneous_unsafe_batch_still_uses_plain_sequential_path(self, agent): + calls = [ + _tc("terminal", '{"command":"a"}'), + _tc("terminal", '{"command":"b"}'), + ] + msg = SimpleNamespace(content="", tool_calls=calls) + + with ( + patch.object(agent, "_execute_tool_calls_concurrent") as conc, + patch.object(agent, "_execute_tool_calls_sequential") as seq, + ): + agent._execute_tool_calls(msg, [], "task-1") + + seq.assert_called_once() + conc.assert_not_called() + + def test_single_call_uses_sequential_path(self, agent): + msg = SimpleNamespace(content="", tool_calls=[_tc("web_search", '{"query":"a"}')]) + + with ( + patch.object(agent, "_execute_tool_calls_concurrent") as conc, + patch.object(agent, "_execute_tool_calls_sequential") as seq, + ): + agent._execute_tool_calls(msg, [], "task-1") + + seq.assert_called_once() + conc.assert_not_called() + + def test_interrupt_during_barrier_drains_later_segments(self, agent): + """Interrupt raised while the barrier tool runs: the trailing parallel + segment must be drained with cancelled results — one per call — + without executing.""" + calls = [ + _tc("web_search", '{"query":"a"}', call_id="s1"), + _tc("web_search", '{"query":"b"}', call_id="s2"), + _tc("terminal", '{"command":"long"}', call_id="t1"), + _tc("web_search", '{"query":"c"}', call_id="s3"), + _tc("web_search", '{"query":"d"}', call_id="s4"), + ] + msg = SimpleNamespace(content="", tool_calls=calls) + messages = [] + executed = [] + lock = threading.Lock() + + def fake_handle(name, args, task_id, **kwargs): + with lock: + executed.append(kwargs["tool_call_id"]) + if kwargs["tool_call_id"] == "t1": + agent._interrupt_requested = True + return json.dumps({"ok": True}) + + with patch("run_agent.handle_function_call", side_effect=fake_handle): + agent._execute_tool_calls(msg, messages, "task-1") + + # Every call still gets exactly one result, in order. + assert [m["tool_call_id"] for m in messages] == ["s1", "s2", "t1", "s3", "s4"] + # s3/s4 were never executed. + assert "s3" not in executed and "s4" not in executed + for m in messages[-2:]: + assert "cancelled" in m["content"] or "skipped" in m["content"] + + def test_steer_lands_exactly_once_in_mixed_batch(self, agent): + """Steer is drained once (per-tool drains + one dispatcher-level + finalize) — the marker must appear exactly once across the batch, + never duplicated by segment boundaries.""" + calls = [ + _tc("web_search", '{"query":"a"}', call_id="s1"), + _tc("web_search", '{"query":"b"}', call_id="s2"), + _tc("terminal", '{"command":"echo hi"}', call_id="t1"), + ] + msg = SimpleNamespace(content="", tool_calls=calls) + messages = [] + + def fake_handle(name, args, task_id, **kwargs): + return json.dumps({"ok": True}) + + agent.steer("focus on the tests") + with patch("run_agent.handle_function_call", side_effect=fake_handle): + agent._execute_tool_calls(msg, messages, "task-1") + + contents = [m["content"] for m in messages] + hits = [c for c in contents if "focus on the tests" in c] + assert len(hits) == 1