perf(agent): segment mixed tool batches to recover lost concurrency (#64460)

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 a single barrier call (terminal, clarify, unknown
tool, malformed args) 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 on the existing
concurrent path, barrier calls on the sequential path, strictly in the
model's emission order. Invariants preserved:

- one tool result per call, appended in emission order (segments are
  contiguous, so no result reordering across a barrier)
- side-effect boundaries: no call starts before an earlier barrier ends
- overlapping file targets split into separate ordered parallel runs
- turn-end budget enforcement + /steer injection run exactly once per
  batch (segment executors run with finalize=False; the segmented
  dispatcher owns the whole-turn finalize)
- interrupt during segment k drains segments k+1..n with cancelled
  results, keeping one result per tool_call_id

Homogeneous batches keep their original single-path dispatch (zero
behavior delta); _should_parallelize_tool_batch remains as a thin view
over the planner for existing callers and tests.

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",

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",
 ]

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

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