Fix parallel toolcall in agent

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
 
 
@@ -866,6 +948,12 @@ async def run(self, input_data: RunAgentInput) -> AsyncIterator[Any]:
             stop_text_streaming = False
             halt_event_stream = False
             pending_halt = False
+            # PATCHED (defer hand-off): frontend-tool ToolCallEnd ids buffered here
+            # so the client's 'execute this frontend tool' signal is delayed until
+            # AFTER this turn's backend tool results have been emitted. Prevents the
+            # client dispatching its follow-up run (tool result) before the current
+            # run finishes -> reduces the ConcurrencyException race window.
+            deferred_frontend_tool_ends = []
 
             # Reasoning/thinking state tracking
             reasoning_started = False
@@ -1229,8 +1317,14 @@ async def run(self, input_data: RunAgentInput) -> AsyncIterator[Any]:
                     # Handle tool results from Strands for backend tool rendering
                     elif "message" in event and event["message"].get("role") == "user":
                         if pending_halt:
+                            # Set halt but DO NOT discard this message: it packs
+                            # results for ALL tools in the turn, and the loop below
+                            # emits the backend-tool results. The buffered
+                            # frontend-tool ToolCallEnd(s) are flushed *after* that
+                            # per-item loop (not here), so the client receives the
+                            # backend TOOL_CALL_RESULT(s) before the frontend
+                            # TOOL_CALL_END that hands off control.
                             halt_event_stream = True
-                            continue
                         message_content = event["message"].get("content", [])
                         if not message_content or not isinstance(message_content, list):
                             continue
@@ -1405,6 +1499,21 @@ async def run(self, input_data: RunAgentInput) -> AsyncIterator[Any]:
                                 # Break inner loop — no further results should be emitted
                                 break
 
+                        # PATCHED (defer hand-off): now that this turn's backend
+                        # TOOL_CALL_RESULT(s) have been emitted above, flush the
+                        # buffered frontend-tool ToolCallEnd(s). Placing the flush
+                        # after the per-item loop guarantees the wire order
+                        # backend TOOL_CALL_RESULT -> frontend TOOL_CALL_END, so the
+                        # client only starts executing the frontend tool (and
+                        # dispatching its follow-up run) after backend work is done.
+                        if pending_halt and 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 = []
+
                     # Handle tool calls
                     elif "current_tool_use" in event and event["current_tool_use"]:
                         tool_use = event["current_tool_use"]
@@ -1709,10 +1818,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,
-                                    )
+                                    # PATCHED (defer hand-off): for frontend tools,
+                                    # buffer the ToolCallEnd instead of emitting now.
+                                    # It is flushed after this turn's backend results
+                                    # (see Fix #3c). Backend tools emit immediately.
+                                    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(
@@ -1902,6 +2021,20 @@ 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
+
+                # PATCHED (defer hand-off, safety flush): if the stream ended
+                # without a backend tool-result message (e.g. a turn with ONLY
+                # parallel frontend tool calls), the pending_halt flush path
+                # above never ran and the buffered frontend ToolCallEnd(s) would
+                # be lost — leaving TOOL_CALL_START events with no matching END.
+                # Flush any remaining buffered ends here before cleanup. This is
+                # a no-op when the mid-stream flush already drained the buffer.
+                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,