[eatmoreduck 多吃鸭 (eatmoreduck)]

Motivation & Context

When a model calls two declaration-only (client-side) tools in one response, the AG-UI endpoint starts each tool call twice: the client receives two TOOL_CALL_START events per call id and the thread snapshot records each call twice, so the next run fails when the provider rejects the duplicate call ids. A single declaration-only call per response is not affected. Reported in #9638, which includes a minimal reproduction against FlowState/_emit_content.

Description & Review Guide

• What are the major changes? _emit_tool_call now guards TOOL_CALL_START with the started-call set (tool_call_id not in flow.tool_calls_by_id) instead of the last-id cursor (tool_call_id != flow.tool_call_id), and adds a comment explaining why.
• What is the impact of these changes? With two declaration-only calls, the function loop pauses and replays both calls after streaming; the cursor sits on the second call, so the replay of the first (and then the second) re-emitted TOOL_CALL_START and re-appended the snapshot entry. Guarding on the set makes each call start exactly once. A single call was unaffected before only because the cursor happened to equal its own id; the set makes that correctness explicit rather than incidental. Resume semantics are unchanged: a fresh FlowState starts nothing, so replayed ids on a new run still START normally.
• What do you want reviewers to focus on? That tool_calls_by_id is only populated on the START path here (the approval confirm_changes and MCP paths are separate), so membership means "already STARTed by _emit_tool_call in this run".

The new regression test mirrors the issue's reproduction: stream two fully-argumented function calls, then replay both as declaration-only pauses, and assert exactly one TOOL_CALL_START per id and no duplicate snapshot entries (fails on main, passes with the fix).

Related Issue

Fixes #9638

[Copilot]

🟡 Changes recommended

Run-wide ID deduplication suppresses distinct calls that reuse completed IDs and can corrupt their snapshot arguments.

1 open issue

• Handle reused provider call IDs as distinct call occurrences · New

What changed in this PR

This PR addresses #9638 in the Python AG-UI integration by preventing duplicate tool-call starts and snapshot entries when multiple declaration-only calls are replayed.

Changes:

• Replace the last-call-ID guard with a lookup of already registered call IDs.
• Explain the replay scenario in a code comment.
• Add a regression test checking that two replayed calls each start and appear in the snapshot once.

File	Description
python/​packages/​ag-ui/​tests/​ag_ui/​test_run_common.py	Add coverage for replaying two declaration-only calls.
python/​packages/​ag-ui/​agent_framework_ag_ui/​_run_common.py	Change and document the tool-call START deduplication guard.

_{Review effort: Balanced}

[eatmoreduck 多吃鸭 (eatmoreduck)]

Good catch — addressed in 803db77. The guard is now "is this call still open" (tool_call_id in flow.tool_calls_by_id and tool_call_id not in flow.tool_calls_ended) rather than run-wide id membership:

• a declaration-only replay of an open call stays deduplicated (the Python: [Bug]: AG-UI starts each tool call twice when an agent calls two declaration-only tools in one response #9638 fix holds);
• a completed id reused by a later distinct call STARTs again, replaces the snapshot entry instead of appending arguments into the old one, and discards the ended marker so the second occurrence can emit TOOL_CALL_END again.

I first tried deduplicating by Content.id per the spec, but _emit_content can receive id-less function-call content on both the streamed and replayed paths, where the occurrence identity degrades back to the call id — the open/ended lifecycle is the signal that separates "replay of this occurrence" from "a new occurrence reusing the id" in both cases. The regression you suggested is in: stream → result → a distinct call reusing the same provider id asserts a fresh START, arguments not concatenated, and the call endable again. Full ag-ui suite: 1,695 passed.

[eatmoreduck 多吃鸭 (eatmoreduck)]

@microsoft-github-policy-service agree [company=""]

[eatmoreduck 多吃鸭 (eatmoreduck)]

@microsoft-github-policy-service agree

[Gian Paolo Santopaolo (gian-alloy)]

Thanks for picking this up, eatmoreduck 多吃鸭 (@eatmoreduck). I reported #9638, and I checked this PR locally at 803db775.

The fix works for #9638. The reproduction from the issue now gives one TOOL_CALL_START and one snapshot entry per call id. The ag-ui suite passes locally (1,693 passed). Two span-correlation tests fail on my machine, but they fail the same way with main's _run_common.py, so they are not caused by this PR.

Two CI checks fail because of this PR:

1. Package Checks: ruff I001. In tests/ag_ui/test_run_common.py, the new _emit_content import sits above _close_reasoning_block.
2. Test Typing Checks: mypy, ty and zuban all report "BaseEvent" has no attribute "tool_call_id" at test_run_common.py:1130 and :1152. Filtering with isinstance(e, ToolCallStartEvent) instead of e.type == EventType.TOOL_CALL_START narrows the type and should clear both (I did not run the type checkers myself).

One problem in the reused-id path. tool_calls_ended is keyed by call id. So tool_calls_ended.discard(tool_call_id) (_run_common.py:764) reopens every snapshot entry with that id, including the first occurrence, which was already ended. get_pending_without_end() (_run_common.py:564-566) then returns both entries. The end-of-run loop in _agent_run.py:3995-4004 sends TOOL_CALL_END once for each entry it returns, so one open call gets two ENDs:

flow = FlowState()
_emit_content(Content.from_function_call(call_id="call_x", name="get_weather", arguments='{"city": "NYC"}'), flow)
_emit_content(Content.from_function_result(call_id="call_x", result="72F"), flow)
_emit_content(Content.from_function_call(call_id="call_x", name="get_weather", arguments='{"city": "LA"}'), flow)
print([tc["id"] for tc in flow.get_pending_without_end()])

	END ids from the end-of-run loop
main	[] (the second occurrence is never ended)
this PR	['call_x', 'call_x'] (the first occurrence is ended a second time)
expected	['call_x']

The END state has to be tracked per occurrence, not per id. One way I tried: make get_pending_without_end() read from tool_calls_by_id, which holds only the latest entry for each id:

return [tc for call_id, tc in self.tool_calls_by_id.items() if call_id not in self.tool_calls_ended]

That gives exactly one END here, and the only new test failure is test_run.py::TestFlowState::test_get_pending_without_end. That test fills pending_tool_calls by hand and leaves tool_calls_by_id empty, so it would need to fill both. Any approach is fine with me. Could you also assert the END count in test_reused_provider_call_id_starts_a_new_occurrence, so this case is covered?

Once CI is green and the END count is fixed, this looks good to me.

I used AI assistance for this review; I ran the checks and the reproduction above myself.

[Gian Paolo Santopaolo (gian-alloy)]

Thanks for picking this up, eatmoreduck 多吃鸭 (@eatmoreduck). I reported #9638, and I checked this PR locally at 803db775.

The fix works for #9638. The reproduction from the issue now gives one TOOL_CALL_START and one snapshot entry per call id. The ag-ui suite passes locally (1,693 passed). Two span-correlation tests fail on my machine, but they fail the same way with main's _run_common.py, so they are not caused by this PR.

Two CI checks fail because of this PR:

1. Package Checks: ruff I001. In tests/ag_ui/test_run_common.py, the new _emit_content import sits above _close_reasoning_block.
2. Test Typing Checks: mypy, ty and zuban all report "BaseEvent" has no attribute "tool_call_id" at test_run_common.py:1130 and :1152. Filtering with isinstance(e, ToolCallStartEvent) instead of e.type == EventType.TOOL_CALL_START narrows the type and should clear both (I did not run the type checkers myself).

One problem in the reused-id path. tool_calls_ended is keyed by call id. So tool_calls_ended.discard(tool_call_id) (_run_common.py:764) reopens every snapshot entry with that id, including the first occurrence, which was already ended. get_pending_without_end() (_run_common.py:564-566) then returns both entries. The end-of-run loop in _agent_run.py:3995-4004 sends TOOL_CALL_END once for each entry it returns, so one open call gets two ENDs:

flow = FlowState()
_emit_content(Content.from_function_call(call_id="call_x", name="get_weather", arguments='{"city": "NYC"}'), flow)
_emit_content(Content.from_function_result(call_id="call_x", result="72F"), flow)
_emit_content(Content.from_function_call(call_id="call_x", name="get_weather", arguments='{"city": "LA"}'), flow)
print([tc["id"] for tc in flow.get_pending_without_end()])

	END ids from the end-of-run loop
main	[] (the second occurrence is never ended)
this PR	['call_x', 'call_x'] (the first occurrence is ended a second time)
expected	['call_x']

The END state has to be tracked per occurrence, not per id. One way I tried: make get_pending_without_end() read from tool_calls_by_id, which holds only the latest entry for each id:

return [tc for call_id, tc in self.tool_calls_by_id.items() if call_id not in self.tool_calls_ended]

That gives exactly one END here, and the only new test failure is test_run.py::TestFlowState::test_get_pending_without_end. That test fills pending_tool_calls by hand and leaves tool_calls_by_id empty, so it would need to fill both. Any approach is fine with me. Could you also assert the END count in test_reused_provider_call_id_starts_a_new_occurrence, so this case is covered?

Once CI is green and the END count is fixed, this looks good to me.

I used AI assistance for this review; I ran the checks and the reproduction above myself.

[eatmoreduck 多吃鸭 (eatmoreduck)]

Thank you for the exceptionally thorough local review — everything is addressed in 17d6a6e:

1. Double END on the reused-id path — exactly as you diagnosed: get_pending_without_end() now reads from tool_calls_by_id (the latest occurrence per id), so an already-ended first occurrence can never be re-endable, and your three-step reproduction now yields ['call_x']. Your suggested approach, taken as-is.
2. test_get_pending_without_end — the hand-filled fixture now populates tool_calls_by_id alongside pending_tool_calls, and a new case pins the reuse semantics directly (no pending entries while the id is marked ended; exactly one after the discard).
3. END count in the reused-id scenario — test_reused_provider_call_id_starts_a_new_occurrence now asserts get_pending_without_end() holds exactly one call_x entry after the second occurrence opens.
4. CI — the _emit_content import moved below _close_reasoning_block (ruff I001 clean), and both TOOL_CALL_START filters use isinstance(e, ToolCallStartEvent).

Verified locally: test_run_common.py 81/81, TestFlowState 5/5; the wider test_run.py failures reproduce identically on main in this environment (unrelated fixtures). The two span-correlation failures you saw match what CI reported on main's _run_common.py as well, as you noted.