Fix parallel toolcall in agent

[YiyuanMiao]

fix(aws-strands): correct frontend/backend parallel tool calls

Summary

Fixes several defects in ag_ui_strands that break parallel frontend + backend tool calls, and removes the temporary build-time patch script that previously carried these fixes out-of-tree.

The fixes now live directly in the package source (agent.py), so downstream consumers (e.g. opensearch-agent-server) get correct behavior simply by installing ag_ui_strands — no post-install patching step required.

Motivation

When the LLM issues frontend and backend tool calls in the same turn, the stock code produced hung runs, dropped results, and intermittent Bedrock errors. These were previously worked around by a build-time script that rewrote the installed agent.py in place — fragile (lost on every reinstall / venv rebuild) and easy to forget. This PR moves the fixes into source and deletes the script.

Changes

Added / modified — src/ag_ui_strands/agent.py

Fix	Problem	Change
#1 halt-break	On halt the loop did continue (mute only), leaving the Strands loop alive → extra Bedrock call + hung run	break so the existing aclose() cleanup terminates Strands
#2 keep-message	Halt discarded the entire tool-result message, dropping parallel backend tool results too	Set halt without dropping the message; the inner loop skips frontend placeholders and emits backend results
#3 defer-handoff	Frontend ToolCallEnd ("execute now" signal) was emitted before backend tools finished → client dispatched its follow-up run early → ConcurrencyException	Buffer frontend ToolCallEnd and flush only after backend results arrive
#3 safety-flush (new)	Turns with only parallel frontend tools never hit the backend-result flush path → buffered ToolCallEnds were lost, leaving TOOL_CALL_START with no matching END	Flush any remaining buffered ends when the stream ends, before cleanup
#4 normalize-history	Client replays prior turns as separate messages; Bedrock requires all toolUse of a turn in one assistant message and matching toolResult in the immediately-following user message → ValidationException on the 2nd call	Add _normalize_tool_turns() (+ _is_tooluse_only_assistant, _is_toolresult_only_user helpers) to normalize rebuilt history

Removed

• src/ag_ui_strands/patch_ag_ui_strands.py — the build-time patch script is no longer needed now that the fixes are in source. Confirmed no other code references it.

[YiyuanMiao]

Merged main to resolve the conflict in agent.py.

The conflict was with #1848 (fix(strands): bundle replayed tool results), which also touched _build_strands_history. Both changes address Bedrock's toolUse/toolResult pairing for parallel tool calls, so I kept upstream's approach as the base and layered this PR's normalization on top rather than replacing it:

• Base (from fix(strands): bundle replayed tool results #1848): pending_tool_results / flush_tool_results bundles parallel tool results into a single following user message.
• This PR adds _normalize_tool_turns as a post-pass, covering three cases fix(strands): bundle replayed tool results #1848 doesn't handle:
  1. Out-of-order results — reordered to match the toolUse order (Bedrock pairs positionally).
  2. Wedged messages — a non-tool message between a toolUse turn and its results is moved after them so the pair stays adjacent.
  3. Consecutive tool-call turns — merged into a single assistant/user pair.

Tests: kept #1848's bundling test and this PR's existing scenarios, and added two tests for the out-of-order and wedged-message cases. Full suite passes (156 passed, 2 skipped).

[YiyuanMiao]

Rebased onto latest main and resolved the conflict in agent.py.

The only conflict was the frontend tool-id handling. My original "reuse Strands' native id" fix is superseded by #2231's durable {wire_id: native_id} session-state mapping, so I dropped my version and kept that one — it's the stronger approach (survives across processes / multi-worker).

The remaining 4 fixes in this PR are independent of #2231 and still needed:

• halt-break: break instead of continue on halt so Strands' aclose() cleanup actually terminates the loop (avoids an extra Bedrock call + hung run).
• keep-message: halt without dropping the whole tool-result message, so parallel backend results aren't lost.
• defer-handoff + safety-flush: buffer frontend ToolCallEnd until backend results arrive (avoids ConcurrencyException), and flush any remaining buffered ends when the stream ends.
• normalize-history (_normalize_tool_turns): coalesce rebuilt multi-message turns so Bedrock doesn't throw ValidationException on the 2nd call.

Full suite green after the merge: 178 passed, 2 skipped.

@contextablemark @NathanTarbert since you both work in this area — could you take a look when you get a chance? Thanks!

[contextablemark]

Thanks for this, and sorry it's sat unreviewed — that's on us, not you.

Some context, because this file has become crowded: #2218 (@FriedhelmWS) touches the inbound prompt-building region, #2286 (@ciolo) adds the interrupt round-trip, and we have a halt-path fix on mark/strands-py-halt-stops-loop that we deliberately have not opened as a PR because two of its hunks overlap yours and you were here first.

Your diagnosis is correct on every count. We found fixes #1 and #2 independently, measured them the same way, and reached the same conclusions — so those two hunks are duplicate work, and the cause is that nobody reviewed this for 26 days.

Fix #1 — right change, wrong reason, and it needs one more line

The description says "break so the existing aclose() cleanup terminates Strands." That mechanism doesn't hold. A break leaves the async generator suspended at a yield, where ag_running is False — so the pre-existing guard

if not agent_stream.ag_running:
    pass
else:
    await agent_stream.aclose()

takes the pass branch and cleanup is skipped entirely. Teardown falls to GC, which in a long-lived server means the Strands cycle and its model stream stay open.

Measured on both builds: yours reports the stream not finalized during run() — still not, after an explicit gc.collect() — and it emits the OpenTelemetry Failed to detach context / ValueError: <Token …> was created in a different Context that the cleanup block exists to prevent. Ours reports finalized. The fix is to close explicitly on the halt path rather than rely on the guard:

if halt_event_stream:
    await agent_stream.aclose()
elif not agent_stream.ag_running:
    pass
else:
    await agent_stream.aclose()

Fix #3 — the placement doesn't deliver the stated guarantee

The flush sits at the top of the if pending_halt: branch, before the per-item loop that emits backend results. So the wire order is still frontend TOOL_CALL_END → then backend TOOL_CALL_RESULT. The goal — "the client only starts executing the frontend tool after backend work is done" — is narrowed rather than achieved: the END now follows backend execution, but still precedes the backend results reaching the client. Moving the flush to after the per-item loop delivers it.

Two more on the same fix:

• It patches only the use_streaming branch, so the args_streamer path still emits the frontend ToolCallEnd immediately.
• The safety flush sits inside the try, after the async for — so any exception from the stream skips it, losing the buffered ends and stranding TOOL_CALL_START with no END. That's a new failure mode the deferral introduces; moving it into the finally (or a try/finally of its own) would close it.

Fix #4 — three problems, all measured

Recursion. _normalize_tool_turns recurses once per tool turn (out.extend(_normalize_tool_turns(leftover)); return out). At ~1000 tool turns that's a RecursionError — fine at 800 (7.2 ms), fails at 1200. An iterative loop over the same structure avoids it.

Duplicate toolUseId emits a duplicate toolResult. ordered = [results_by_id[tid] for tid in tooluse_ids if tid in results_by_id] iterates tooluse_ids, which can contain the same id twice, so the same result block is emitted twice. Bedrock rejects that.

Reordering. Non-tool messages between a toolUse turn and its results get moved rather than preserved in place — the leftover list re-emits them after the merged pair. Deliberate for messages wedged between, per the docstring, but it also reorders messages that legitimately follow.

Your description lists a fix that's no longer in the diff

The "frontend-id" row — reuse Strands' stable id instead of a per-run UUID — has no corresponding uuid4 change in the diff any more. Your 2026-07-29 merge from upstream picked up main's approach to the same problem: frontend tools keep a generated wire id, and a durable {wire_id: native_id} map on agent state lets the continuation run reconcile against it (session_reconcile.py).

Worth removing that row, or saying explicitly if you think the map approach is worse — that's a reasonable thing to argue, but it should be an argument rather than a stale table entry.

Tests

test_parallel_tool_call_handling.py covers fixes #3 and #4 but adds nothing for #1 or #2 — the two that overlap our branch. Ours carries 12 tests for that path, including a two-run round-trip proving the emitted event stream replays into a transcript the model provider accepts.

How I'd suggest we land this

Easy either way — you were here first:

1. You drop hunks @@ -1044 and @@ -1229; we open ours for those two with its tests; yours keeps fixes trying to find websockets transport for ag-ui #3 and Fix the README file to improve visual appeal. #4, where you're the only one who's touched them.
2. You keep all six and we hand you our aclose() hunk, the flush-placement move, and the 12 tests to fold in.

Option 1 is less work for you and gets both landed sooner. Happy to open ours immediately on your word, or to pair on the #4 items.

[YiyuanMiao]

@contextablemark
Thanks for the thorough review — this is really helpful. I reproduced all of the findings locally and they hold:

• Fix#2: break leaves the generator suspended at a yield with ag_running == False, so the guard takes the pass branch and aclose() never runs — the resource stays
  open until GC. An explicit await agent_stream.aclose() on the halt path fixes it.
• Fix#3 placement: the emitted order is indeed frontend TOOL_CALL_END → backend TOOL_CALL_RESULT; moving the flush after the per-item loop corrects it.
• Fix#3 safety flush: an exception from the stream skips it inside the try; it needs to be in finally.
• Fix#4: RecursionError at ~1200 turns (fine at 800), and a duplicate toolUseId emits a duplicate toolResult.

Let's go with Option 1. You were here first on #1/#2, and you already have the correct implementation plus tests, so I'll drop those hunks and you open yours. I'll keep and fix #3 and #4, which only I've touched.

One thing I want to nail down before I start cutting — the hunk boundary, since two of your #3 sub-points sit right on the edge of the overlap:

1. args_streamer path still emitting the frontend ToolCallEnd immediately — does that live in the @@ -1044 region you're taking, or should I handle it as part of the Fix#3 logic I'm retaining?
2. Safety flush → finally — same question: your PR or mine?

My read is that both are entangled with the halt/handoff cleanup in #1/#2, so they'd naturally go with your hunks, but I'd rather confirm than double-fix or leave a
gap. Let me know how you'd like to split those two.

On my side, once the boundary is set, I'll:

• rewrite _normalize_tool_turns iteratively;
• dedupe toolUseId when building ordered;
• preserve messages that legitimately follow the results (keep dropping only the ones wedged between, per the docstring);
• move the flush after the per-item loop (if Fix#3 stays with me);
• add regression tests for each (1200-turn no-crash, duplicate-id, follow-message ordering, flush order).

I'll also drop the stale "frontend-id" row from the description — you're right that the session_reconcile.py {wire_id: native_id} map from the 07-29 upstream merge supersedes it, and I don't have a case that the map approach is worse.

Happy to pair on the #4 items whenever works for you. Go ahead and open yours for #1/#2 — I'll rebase around it.

[YiyuanMiao]

@contextablemark Pushed my half of Option 1.

Done in this branch (fixes #3 and #4 — the parts only I've touched):

• Fix#4 _normalize_tool_turns** — rewrote the recursion as an iterative loop (no more RecursionError past ~1000 turns); de-duplicate toolUseId so a repeated id no longer emits a duplicate toolResult; and preserve messages that legitimately follow a completed pair in place (still dropping only the ones wedged between).
• Fix#3 flush placement** — moved the buffered frontend ToolCallEnd flush to after the per-item backend loop, so the wire order is now backend TOOL_CALL_RESULT → frontend TOOL_CALL_END.
• Added regression tests for each (1200+ turn no-crash, duplicate-id dedup, follow-message ordering, and the flush order). All green.

On the two boundary sub-points from my last comment — the args_streamer path still emitting the frontend end immediately, and moving the safety flush into finally — I've left both untouched, since they're entangled with the halt-path cleanup in #1/#2 that's yours. They'll come with your branch.

Important: I have not dropped the two overlapping hunks (@@ -1044 and @@ -1229 — the Fix#1 break and the Fix#2 cleanup) yet. They're still live in my branch so #1/#2 aren't left unfixed in the meantime. Please open your PR for those two when you can — as soon as it's up, I'll drop my hunks and rebase around yours so we don't double-land them.

I'll also drop the stale "frontend-id" row from the description.

[contextablemark]

@YiyuanMiao Opened as #2334 — thanks for the unusually careful review cycle on this, and for reproducing the findings independently before acting on them.

It carries the two overlapping hunks (Fix #1 halt-break, Fix #2 keep-message) with 12 regression tests, plus the explicit aclose() on the halt path that the break alone doesn't cover.

Drop your @@ -1044 and @@ -1229 hunks whenever suits and rebase around it. On the two sub-points you left me — the args_streamer path emitting the frontend end immediately, and moving the safety flush into finally — I'm happy either way: I can take them as a follow-up commit on #2334 since it already rewrites that teardown region, or leave them with #2172 once mine lands. Say which you'd prefer.

One note on CI: #2172 is showing only a single check, so it looks fork-gated and needs a maintainer to approve the run. I'll get that approved so your new tests actually report.

[YiyuanMiao]

@contextablemark Thanks for opening #2334 so quickly.

Done on my side:

• Dropped the two overlapping hunks. Reverted Fix#1 (halt break → back to the original continue) so the halt-path fix — including the explicit aclose() — lands via fix(aws-strands): halt the Strands loop on frontend tools instead of muting the wire #2334. I kept the pending_halt message-retention (no longer discarding the tool-result message), since the deferred frontend-end flush depends on it and fix(aws-strands): halt the Strands loop on frontend tools instead of muting the wire #2334 does the same; updated its comment to stop referencing the reverted break.
• Removed the stale frontend-id row from the PR description.

So #2172 is now just Fix#3 (flush placement) + Fix#4 (_normalize_tool_turns: iterative rewrite, dup-id dedupe, follow-message ordering) with their regression tests.

On the two sub-points: please take them on #2334 — the args_streamer path and moving the safety flush into finally are both entangled with the teardown region you're already rewriting, so they're cleaner as one change there than split across our two PRs.

One heads-up: we'll likely hit a small merge conflict in the if pending_halt:block once #2334 lands, since we both removed that continue — trivial to resolve (same intent on both sides). Happy to rebase onto #2334 after it merges.

Also — CI on #2172 is still sitting at action_required (fork-gated), so my new tests haven't reported yet. Could you approve the run when you get a chance?