diff --git a/integrations/aws-strands/python/src/ag_ui_strands/agent.py b/integrations/aws-strands/python/src/ag_ui_strands/agent.py index 60896ac029..f94a15fdb4 100644 --- a/integrations/aws-strands/python/src/ag_ui_strands/agent.py +++ b/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, diff --git a/integrations/aws-strands/python/tests/test_parallel_tool_call_handling.py b/integrations/aws-strands/python/tests/test_parallel_tool_call_handling.py index 76e6eedc49..5f91bb6279 100644 --- a/integrations/aws-strands/python/tests/test_parallel_tool_call_handling.py +++ b/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})" + )