fix(aws-strands): correct parallel tool-call history and defer frontend hand-off

integrations/aws-strands/python/src/ag_ui_strands/agent.py

@@ -293,6 +293,88 @@ def flush_tool_results() -> None:
             out.append({"role": "assistant", "content": blocks})
 
     flush_tool_results()
+    # Normalize so Bedrock's toolUse/toolResult pairing holds even when results
+    # arrive out of order, are wedged apart by other messages, or span multiple
+    # consecutive tool-call turns (parallel tool calls).
+    return _normalize_tool_turns(out)
+
+
+def _is_tooluse_only_assistant(m):
+    return (
+        m.get("role") == "assistant"
+        and m.get("content")
+        and all("toolUse" in b for b in m["content"])
+    )
+
+
+def _is_toolresult_only_user(m):
+    return (
+        m.get("role") == "user"
+        and m.get("content")
+        and all("toolResult" in b for b in m["content"])
+    )
+
+
+def _normalize_tool_turns(msgs):
+    """Merge same-turn toolUse into one assistant msg and their toolResults
+    into the immediately following user msg, dropping any messages wedged
+    between a toolUse turn and its toolResults so Bedrock accepts the history.
+
+    Messages that legitimately *follow* a completed toolUse/toolResult pair are
+    preserved in place; only messages wedged *between* the toolUse turn and its
+    results are dropped.
+    """
+    out = []
+    i = 0
+    n = len(msgs)
+    while i < n:
+        m = msgs[i]
+        if not _is_tooluse_only_assistant(m):
+            out.append(m)
+            i += 1
+            continue
+
+        # Collect consecutive toolUse-only assistant messages into one.
+        merged_tooluse = list(m["content"])
+        j = i + 1
+        while j < n and _is_tooluse_only_assistant(msgs[j]):
+            merged_tooluse.extend(msgs[j]["content"])
+            j += 1
+        # Preserve first-seen order and de-duplicate ids: a repeated toolUseId
+        # must not later emit a duplicate toolResult (Bedrock rejects that).
+        tooluse_ids = []
+        seen_ids = set()
+        for b in merged_tooluse:
+            rid = b["toolUse"]["toolUseId"]
+            if rid not in seen_ids:
+                seen_ids.add(rid)
+                tooluse_ids.append(rid)
+
+        # Scan forward for the matching toolResults. Anything that is not a
+        # matching result and appears *before* results are complete is "wedged"
+        # and dropped; once every result is collected, the remaining messages
+        # are left untouched to be processed in place by the outer loop.
+        results_by_id = {}
+        k = j
+        while k < n and len(results_by_id) < len(tooluse_ids):
+            mk = msgs[k]
+            if _is_toolresult_only_user(mk):
+                for b in mk["content"]:
+                    rid = b["toolResult"].get("toolUseId")
+                    if rid in seen_ids and rid not in results_by_id:
+                        results_by_id[rid] = b
+                    # non-matching / duplicate result blocks wedged in are dropped
+            # non-toolResult messages wedged before completion are dropped
+            k += 1
+
+        # Emit merged assistant(toolUse) + merged user(toolResult) adjacently.
+        out.append({"role": "assistant", "content": merged_tooluse})
+        ordered = [results_by_id[tid] for tid in tooluse_ids if tid in results_by_id]
+        if ordered:
+            out.append({"role": "user", "content": ordered})
+
+        # Continue with whatever legitimately follows, in place (no reordering).
+        i = k
     return out
 
 
@@ -871,6 +953,12 @@ async def run(self, input_data: RunAgentInput) -> AsyncIterator[Any]:
             stop_text_streaming = False
             halt_event_stream = False
             pending_halt = False
+            # Frontend-tool ToolCallEnd ids are buffered here so the client's
+            # "execute this frontend tool" signal is delayed until AFTER this
+            # turn's backend tool results have been emitted. This prevents the
+            # client dispatching its follow-up run before the backend results
+            # reach it, narrowing the ConcurrencyException race window.
+            deferred_frontend_tool_ends = []
 
             # Reasoning/thinking state tracking
             reasoning_started = False
@@ -1444,6 +1532,21 @@ async def run(self, input_data: RunAgentInput) -> AsyncIterator[Any]:
                                 # Break inner loop — no further results should be emitted
                                 break
 
+                        # Defer hand-off: now that this turn's backend
+                        # TOOL_CALL_RESULT(s) have been emitted above, flush the
+                        # buffered frontend-tool ToolCallEnd(s). Flushing here —
+                        # after the per-item loop and before the halt break below —
+                        # guarantees the wire order backend TOOL_CALL_RESULT ->
+                        # frontend TOOL_CALL_END, so the client only starts the
+                        # frontend tool once backend work has reached it.
+                        if deferred_frontend_tool_ends:
+                            for _fe_tool_use_id in deferred_frontend_tool_ends:
+                                yield ToolCallEndEvent(
+                                    type=EventType.TOOL_CALL_END,
+                                    tool_call_id=_fe_tool_use_id,
+                                )
+                            deferred_frontend_tool_ends = []
+
                         # The batch is fully emitted; stop before Strands runs
                         # another model cycle. Breaking HERE rather than relying
                         # on the check at the top of the loop means termination
@@ -1756,10 +1859,20 @@ async def run(self, input_data: RunAgentInput) -> AsyncIterator[Any]:
                                                 exc_info=True,
                                             )
 
-                                    yield ToolCallEndEvent(
-                                        type=EventType.TOOL_CALL_END,
-                                        tool_call_id=tool_use_id,
-                                    )
+                                    # Defer hand-off: for frontend tools, buffer the
+                                    # ToolCallEnd instead of emitting it now. It is
+                                    # flushed after this turn's backend results (see
+                                    # the pending_halt handler). Backend tools and
+                                    # continue_after_frontend_call tools emit now.
+                                    if is_frontend_tool and not (
+                                        behavior and behavior.continue_after_frontend_call
+                                    ):
+                                        deferred_frontend_tool_ends.append(tool_use_id)
+                                    else:
+                                        yield ToolCallEndEvent(
+                                            type=EventType.TOOL_CALL_END,
+                                            tool_call_id=tool_use_id,
+                                        )
 
                                     if self._will_emit_tool_snapshot(behavior, emit_snapshots):
                                         snapshot_messages.append(
@@ -1949,6 +2062,19 @@ async def run(self, input_data: RunAgentInput) -> AsyncIterator[Any]:
                                             f"Deferring halt after frontend tool call: tool_name={tool_name}, tool_call_id={tool_use_id}, thread_id={input_data.thread_id}"
                                         )
                                         pending_halt = True
+
+                # Defer hand-off (safety flush): if the stream ended without a
+                # backend tool-result message (e.g. a turn with ONLY frontend tool
+                # calls), the per-batch flush above never ran and the buffered
+                # frontend ToolCallEnd(s) would be lost — leaving TOOL_CALL_START
+                # events with no matching END. Flush any remainder here.
+                if deferred_frontend_tool_ends:
+                    for _fe_tool_use_id in deferred_frontend_tool_ends:
+                        yield ToolCallEndEvent(
+                            type=EventType.TOOL_CALL_END,
+                            tool_call_id=_fe_tool_use_id,
+                        )
+                    deferred_frontend_tool_ends = []
             finally:
                 # Properly close the async generator to avoid context detachment errors
                 # The generator should complete naturally when we consume all events,

integrations/aws-strands/python/tests/test_parallel_tool_call_handling.py

@@ -38,7 +38,11 @@
 )
 from strands.tools.registry import ToolRegistry
 
-from ag_ui_strands.agent import StrandsAgent, _build_strands_history
+from ag_ui_strands.agent import (
+    StrandsAgent,
+    _build_strands_history,
+    _normalize_tool_turns,
+)
 from ag_ui_strands.config import StrandsAgentConfig, ToolBehavior
 
 
@@ -143,6 +147,144 @@ def test_strands_history_bundles_parallel_tool_results():
     }
 
 
+def test_strands_history_reorders_out_of_order_tool_results():
+    """Tool results that arrive in a different order than their toolUse blocks
+    must be re-ordered to match, so Bedrock's positional pairing holds."""
+    messages = [
+        UserMessage(id="u1", role="user", content="hi"),
+        AssistantMessage(
+            id="a1",
+            role="assistant",
+            content="",
+            tool_calls=[
+                ToolCall(id="tooluse_1", type="function",
+                         function=FunctionCall(name="my_tool", arguments="{}")),
+                ToolCall(id="tooluse_2", type="function",
+                         function=FunctionCall(name="my_tool", arguments="{}")),
+                ToolCall(id="tooluse_3", type="function",
+                         function=FunctionCall(name="my_tool", arguments="{}")),
+            ],
+        ),
+        # Results arrive out of order: 3, 1, 2
+        ToolMessage(id="t3", role="tool", tool_call_id="tooluse_3", content='{"ok":3}'),
+        ToolMessage(id="t1", role="tool", tool_call_id="tooluse_1", content='{"ok":1}'),
+        ToolMessage(id="t2", role="tool", tool_call_id="tooluse_2", content='{"ok":2}'),
+    ]
+
+    native = _build_strands_history(messages)
+
+    assert len(native) == 3
+    assert native[1]["role"] == "assistant"
+    assert native[2]["role"] == "user"
+    tool_use_ids = [b["toolUse"]["toolUseId"] for b in native[1]["content"]]
+    result_ids = [b["toolResult"]["toolUseId"] for b in native[2]["content"]]
+    # toolResult order must follow the toolUse order, not arrival order.
+    assert result_ids == tool_use_ids == ["tooluse_1", "tooluse_2", "tooluse_3"]
+
+
+def test_strands_history_keeps_tooluse_and_results_adjacent():
+    """A non-tool message wedged between a toolUse turn and its results must be
+    moved out so the assistant(toolUse) message is immediately followed by the
+    user(toolResult) message, as Bedrock requires."""
+    messages = [
+        UserMessage(id="u1", role="user", content="hi"),
+        AssistantMessage(
+            id="a1",
+            role="assistant",
+            content="",
+            tool_calls=[
+                ToolCall(id="tooluse_1", type="function",
+                         function=FunctionCall(name="my_tool", arguments="{}")),
+                ToolCall(id="tooluse_2", type="function",
+                         function=FunctionCall(name="my_tool", arguments="{}")),
+            ],
+        ),
+        # A stray user message wedged between the toolUse turn and its results.
+        UserMessage(id="uX", role="user", content="wait, interrupting"),
+        ToolMessage(id="t1", role="tool", tool_call_id="tooluse_1", content='{"ok":1}'),
+        ToolMessage(id="t2", role="tool", tool_call_id="tooluse_2", content='{"ok":2}'),
+    ]
+
+    native = _build_strands_history(messages)
+
+    # Find the assistant(toolUse) message; its immediate successor must be the
+    # user(toolResult) message, with the wedged text pushed afterwards.
+    tooluse_idx = next(
+        i for i, m in enumerate(native)
+        if m["role"] == "assistant" and all("toolUse" in b for b in m["content"])
+    )
+    following = native[tooluse_idx + 1]
+    assert following["role"] == "user"
+    assert all("toolResult" in b for b in following["content"])
+    result_ids = [b["toolResult"]["toolUseId"] for b in following["content"]]
+    assert result_ids == ["tooluse_1", "tooluse_2"]
+
+
+# ---------------------------------------------------------------------------
+# _normalize_tool_turns regression tests (Fix #4)
+# ---------------------------------------------------------------------------
+
+
+def _assistant_tooluse(*ids):
+    return {"role": "assistant", "content": [{"toolUse": {"toolUseId": i}} for i in ids]}
+
+
+def _user_toolresult(*ids):
+    return {
+        "role": "user",
+        "content": [{"toolResult": {"toolUseId": i, "content": []}} for i in ids],
+    }
+
+
+def test_normalize_tool_turns_handles_many_turns_without_recursion():
+    """A previously-recursive implementation raised RecursionError past ~1000
+    tool turns; the iterative version must handle far more."""
+    msgs = []
+    for t in range(5000):
+        tid = f"t{t}"
+        msgs.append(_assistant_tooluse(tid))
+        msgs.append(_user_toolresult(tid))
+
+    out = _normalize_tool_turns(msgs)
+
+    # Each turn collapses to an assistant(toolUse)+user(toolResult) pair.
+    assert len(out) == 10000
+    assert out[0]["role"] == "assistant"
+    assert out[1]["role"] == "user"
+
+
+def test_normalize_tool_turns_deduplicates_repeated_tooluse_id():
+    """A repeated toolUseId must not emit a duplicate toolResult — Bedrock
+    rejects two result blocks with the same id."""
+    msgs = [
+        _assistant_tooluse("a", "b", "a"),  # "a" appears twice
+        _user_toolresult("a", "b"),
+    ]
+
+    out = _normalize_tool_turns(msgs)
+
+    user_msg = next(m for m in out if m["role"] == "user")
+    result_ids = [b["toolResult"]["toolUseId"] for b in user_msg["content"]]
+    assert result_ids == ["a", "b"]  # no duplicate "a"
+
+
+def test_normalize_tool_turns_preserves_messages_that_follow_results():
+    """Messages that legitimately follow a completed toolUse/toolResult pair
+    must be preserved in place, not reordered or dropped."""
+    trailing = {"role": "assistant", "content": [{"text": "all done"}]}
+    msgs = [
+        _assistant_tooluse("a"),
+        _user_toolresult("a"),
+        trailing,
+    ]
+
+    out = _normalize_tool_turns(msgs)
+
+    assert out[0]["role"] == "assistant" and "toolUse" in out[0]["content"][0]
+    assert out[1]["role"] == "user"
+    assert out[2] == trailing  # preserved, in place
+
+
 # ---------------------------------------------------------------------------
 # Scenario A – All parallel frontend tool calls must be emitted
 # ---------------------------------------------------------------------------
@@ -341,3 +483,58 @@ async def test_only_halting_result_emitted(self):
         assert len(result_events) == 1, (
             f"Expected exactly 1 result event, got {len(result_events)}: {result_ids}"
         )
+
+
+# ---------------------------------------------------------------------------
+# Fix #3 – Deferred hand-off flush order
+# ---------------------------------------------------------------------------
+
+class TestDeferredFrontendEndFlushOrder:
+    """When a turn mixes a frontend tool call with a backend tool result, the
+    frontend TOOL_CALL_END (which hands control to the client) must be emitted
+    *after* the backend TOOL_CALL_RESULT, so the client only starts executing
+    the frontend tool once backend work has reached it.
+    """
+
+    THREAD = "flush-order-thread"
+    TOOLS = [Tool(name="frontend_a", description="a", parameters={})]
+    STREAM = [
+        # Frontend tool call + its completion: buffers the END, sets pending_halt.
+        {"current_tool_use": {"name": "frontend_a", "toolUseId": "fe-1", "input": {}}},
+        {"event": {"contentBlockStop": {}}},
+        # Backend result for this turn arrives in a user message.
+        {
+            "message": {
+                "role": "user",
+                "content": [
+                    {
+                        "toolResult": {
+                            "toolUseId": "be-1",
+                            "content": [{"text": '{"ok": 1}'}],
+                        }
+                    }
+                ],
+            }
+        },
+    ]
+
+    async def test_backend_result_precedes_frontend_end(self):
+        agent = _build_agent(self.THREAD, self.STREAM)
+        events = await _collect(agent, _run_input(self.THREAD, tools=self.TOOLS))
+
+        types = [e.type for e in events]
+
+        # The frontend tool's ToolCallEnd is buffered and flushed last; the wire
+        # id is a freshly generated one (not the native toolUseId), so match by
+        # event type rather than id. The backend TOOL_CALL_RESULT must precede it.
+        assert EventType.TOOL_CALL_RESULT in types, "backend TOOL_CALL_RESULT not emitted"
+        assert EventType.TOOL_CALL_END in types, "frontend TOOL_CALL_END not emitted"
+
+        result_idx = types.index(EventType.TOOL_CALL_RESULT)
+        # The deferred frontend end is the last TOOL_CALL_END in the stream.
+        end_idx = max(i for i, t in enumerate(types) if t == EventType.TOOL_CALL_END)
+
+        assert result_idx < end_idx, (
+            "backend TOOL_CALL_RESULT must be emitted before the frontend "
+            f"TOOL_CALL_END (got result at {result_idx}, end at {end_idx})"
+        )