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 53eb22d847..221ed2f371 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 @@ -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, 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})" + )