From 1b7ef23054a5260d0cec15ded2ebf7fccb982924 Mon Sep 17 00:00:00 2001 From: Shivam Sharma <192242476+ShivamSharma43@users.noreply.github.com> Date: Tue, 29 Sep 2026 20:48:27 +0530 Subject: [PATCH 1/7] Python: Keep every queued approval result in the AG-UI thread snapshot With ToolApprovalMiddleware, decisions collected across separate resumes are released as one batch. Both calls ran and streamed TOOL_CALL_RESULT, but the snapshot kept only the result whose approval arrived in the final request, so the next turn replayed both calls with one result. Native approval results are dropped from the flow on the assumption that the snapshot merge restores them from the request's messages. A decision collected in an earlier request has no approval response there, so its result was lost. Carry those results into the merge so each is persisted next to its call. Fixes #8789 Co-Authored-By: Claude Opus 5.5 --- .../ag-ui/agent_framework_ag_ui/_agent_run.py | 19 +++++- .../ag-ui/tests/ag_ui/test_endpoint.py | 60 +++++++++++++++++++ 2 files changed, 78 insertions(+), 1 deletion(-) diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py b/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py index cece3ed33dd..a931fb77456 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py @@ -3406,6 +3406,7 @@ async def _run_agent_stream( telemetry_context = partial(_use_telemetry_conversation_id, telemetry_conversation_id) stream_completed = False native_approval_flow_result_ids: set[int] = set() + native_approval_results: list[Content] = [] try: with telemetry_context(): for queued_executions in forwarded_executions.values(): @@ -3492,6 +3493,7 @@ async def _run_agent_stream( and intent.owner is not ApprovalExecutionOwner.HOSTED ): native_approval_result = True + native_approval_results.append(content) suppress_approval_call_end = intent.owner is ApprovalExecutionOwner.LOCAL _replace_approval_contents_with_results( messages, _collect_approval_responses(messages), [[content]] @@ -3748,7 +3750,22 @@ async def _run_agent_stream( # Feature #5: Suppress intermediate snapshots for predictive tools without confirmation _clean_resolved_approvals_from_snapshot(snapshot_messages, messages) if native_approval_flow_result_ids: - _merge_resolved_approval_results_into_snapshot(snapshot_messages, messages) + # A decision collected in an earlier request has no approval response in this + # request's messages, so its result was never folded into them. Carry it into + # the merge so it is persisted next to its call like the rest of the batch. + resolved_call_ids = { + str(content.call_id) + for message in messages + if get_role_value(message) == "tool" + for content in message.contents or [] + if content.type == "function_result" and content.call_id + } + carried_results = [ + Message(role="tool", contents=[result]) + for result in native_approval_results + if str(result.call_id) not in resolved_call_ids + ] + _merge_resolved_approval_results_into_snapshot(snapshot_messages, [*messages, *carried_results]) flow.tool_results[:] = [ result for result in flow.tool_results if id(result) not in native_approval_flow_result_ids ] diff --git a/python/packages/ag-ui/tests/ag_ui/test_endpoint.py b/python/packages/ag-ui/tests/ag_ui/test_endpoint.py index c10e20fcfc3..7fee1002195 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_endpoint.py +++ b/python/packages/ag-ui/tests/ag_ui/test_endpoint.py @@ -3775,6 +3775,7 @@ def _build_tool_approval_queue_endpoint( streaming_chat_client_stub: Any, *, lifecycle: ApprovalLifecycle | None = None, + snapshot_store: InMemoryAGUIThreadSnapshotStore | None = None, ) -> tuple[TestClient, list[str], list[Message], dict[str, str], AgentFrameworkAgent]: executed: list[str] = [] messages_received: list[Message] = [] @@ -3826,6 +3827,8 @@ async def stream_fn( app, wrapped_agent, path="/approval", + snapshot_store=snapshot_store, + snapshot_scope_resolver=(lambda _request: "test") if snapshot_store is not None else None, ) return TestClient(app), executed, messages_received, state, wrapped_agent @@ -5916,6 +5919,63 @@ async def test_endpoint_agent_approval_resume_surfaces_queued_tool_approval(stre assert [content.call_id for content in replayed_results] == ["call_first", "call_second"] +async def test_endpoint_queued_approval_snapshot_retains_every_released_result(streaming_chat_client_stub): + """Results for decisions collected in earlier requests survive the snapshot round trip (#8789).""" + store = InMemoryAGUIThreadSnapshotStore() + client, executed, messages_received, state, _ = _build_tool_approval_queue_endpoint( + streaming_chat_client_stub, snapshot_store=store + ) + thread_id = "thread-queued-snapshot" + + def post(run_id: str, *, text: str | None = None, approval_id: str | None = None) -> list[dict[str, Any]]: + payload: dict[str, Any] = { + "runId": run_id, + "threadId": thread_id, + "messages": [{"role": "user", "content": text}] if text else [], + } + if approval_id: + payload["resume"] = [{"interruptId": approval_id, "status": "resolved", "payload": {"accepted": True}}] + response = client.post("/approval", json=payload) + assert response.status_code == 200 + events = _decode_sse_events(response) + assert not [event for event in events if event["type"] == "RUN_ERROR"] + return events + + def next_approval_id(events: list[dict[str, Any]]) -> str: + finished = [event for event in events if event["type"] == "RUN_FINISHED"] + return _run_finished_interrupts(finished[-1])[0]["id"] + + def tool_result_call_ids(messages: list[dict[str, Any]]) -> list[str]: + return [ + str(message.get("toolCallId") or message.get("tool_call_id")) + for message in messages + if message.get("role") == "tool" + ] + + pause_events = post("run-pause", text="Run both tools") + state["phase"] = "resume" + first_resume_events = post("run-resume-first", approval_id=next_approval_id(pause_events)) + final_events = post("run-resume-second", approval_id=next_approval_id(first_resume_events)) + assert executed == ["first", "second"] + + snapshot_events = [event for event in final_events if event["type"] == "MESSAGES_SNAPSHOT"] + assert snapshot_events + event_messages = snapshot_events[-1]["messages"] + assert tool_result_call_ids(event_messages) == ["call_first", "call_second"] + # Both results follow the call message and precede the model's reply. + assert [message["role"] for message in event_messages] == ["user", "assistant", "tool", "tool", "assistant"] + + saved = await store.get(scope="test", thread_id=thread_id) + assert saved is not None + assert tool_result_call_ids(saved.messages) == ["call_first", "call_second"] + + post("run-next-turn", text="What happened?") + replayed_results = [ + content for message in messages_received for content in message.contents if content.type == "function_result" + ] + assert [content.call_id for content in replayed_results] == ["call_first", "call_second"] + + @pytest.mark.parametrize("retry_time", [25.0, 33.0]) async def test_endpoint_collected_approval_preserves_original_expiration( streaming_chat_client_stub, retry_time: float From 439f3669f81ca19c4e2aeeba8ff8482f4e3b66a3 Mon Sep 17 00:00:00 2001 From: Shivam Sharma <192242476+ShivamSharma43@users.noreply.github.com> Date: Fri, 2 Oct 2026 00:23:23 +0530 Subject: [PATCH 2/7] Python: Keep every result for a reused call id in the AG-UI snapshot merge Distinct function-call occurrences can share one provider call id, for example hidden siblings released with a guarded call. The snapshot merge keyed resolved results by call id, so only the last one survived and the emitted and persisted snapshots dropped the rest. _resolved_tool_result_snapshot_messages now groups results by call id in resolution order, and the merge places every result for a call id after the assistant message that declares it. Results carried from earlier requests are matched by result object instead of call id, so a reused id cannot hide a result that was never folded into this request's messages. Co-Authored-By: Claude Opus 5.5 --- .../ag-ui/agent_framework_ag_ui/_agent_run.py | 56 ++++++++------ .../ag-ui/tests/ag_ui/test_endpoint.py | 73 +++++++++++++++++-- .../ag-ui/tests/ag_ui/test_run_common.py | 6 +- 3 files changed, 105 insertions(+), 30 deletions(-) diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py b/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py index a931fb77456..ff54a35820d 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py @@ -2220,9 +2220,13 @@ def _snapshot_tool_call_ids(message: Mapping[str, Any]) -> list[str]: return call_ids -def _resolved_tool_result_snapshot_messages(resolved_messages: list[Message]) -> dict[str, dict[str, Any]]: - """Build replayable AG-UI tool messages from resolved approval results.""" - result_by_call_id: dict[str, dict[str, Any]] = {} +def _resolved_tool_result_snapshot_messages(resolved_messages: list[Message]) -> dict[str, list[dict[str, Any]]]: + """Build replayable AG-UI tool messages from resolved approval results. + + Distinct function-call occurrences can reuse one provider call id, so results are + grouped by call id in resolution order rather than collapsed to one per id. + """ + results_by_call_id: dict[str, list[dict[str, Any]]] = {} for msg in resolved_messages: if get_role_value(msg) != "tool": continue @@ -2247,48 +2251,57 @@ def _resolved_tool_result_snapshot_messages(resolved_messages: list[Message]) -> _model_items_for_agui_replay(content, llm_result), ) ) - result_by_call_id[call_id] = snapshot_message - return result_by_call_id + results_by_call_id.setdefault(call_id, []).append(snapshot_message) + return results_by_call_id def _merge_resolved_approval_results_into_snapshot( snapshot_messages: list[dict[str, Any]], resolved_messages: list[Message], ) -> None: - """Persist approval-resolved tool results under their original tool call ids.""" - result_by_call_id = _resolved_tool_result_snapshot_messages(resolved_messages) - if not result_by_call_id: + """Persist approval-resolved tool results under their original tool call ids. + + Every result for a call id follows the assistant message that declares it, in resolution + order. Sibling occurrences hidden from the client share the visible call's id, so one + declared call can be followed by several results. + """ + results_by_call_id = _resolved_tool_result_snapshot_messages(resolved_messages) + if not results_by_call_id: snapshot_messages[:] = [message for message in snapshot_messages if not message.get("function_approvals")] return + retained_host_messages: set[int] = set() for message in snapshot_messages: if normalize_agui_role(message.get("role", "")) != "tool": continue tool_call_id = message.get("toolCallId") or message.get("tool_call_id") if not tool_call_id or message.get(_AGUI_MCP_TOOL_RESULT_KEY) is not True: continue - replacement = result_by_call_id.get(str(tool_call_id)) - if replacement is not None and replacement.get(_AGUI_MCP_TOOL_RESULT_KEY) is not True: - result_by_call_id.pop(str(tool_call_id)) + queued_results = results_by_call_id.get(str(tool_call_id)) + if queued_results and queued_results[0].get(_AGUI_MCP_TOOL_RESULT_KEY) is not True: + # Keep the host-rendered result already in the snapshot for this occurrence. + queued_results.pop(0) + retained_host_messages.add(id(message)) + if not queued_results: + results_by_call_id.pop(str(tool_call_id)) merged_messages: list[dict[str, Any]] = [] for message in snapshot_messages: if message.get("function_approvals"): continue role = normalize_agui_role(message.get("role", "")) - if role == "tool": + if role == "tool" and id(message) not in retained_host_messages: tool_call_id = message.get("toolCallId") or message.get("tool_call_id") - if tool_call_id and str(tool_call_id) in result_by_call_id: + if tool_call_id and str(tool_call_id) in results_by_call_id: continue merged_messages.append(message) if role != "assistant": continue for call_id in _snapshot_tool_call_ids(message): - result_message = result_by_call_id.pop(call_id, None) - if result_message is not None: - merged_messages.append(result_message) + merged_messages.extend(results_by_call_id.pop(call_id, [])) - merged_messages.extend(result_by_call_id.values()) + for queued_results in results_by_call_id.values(): + merged_messages.extend(queued_results) snapshot_messages[:] = merged_messages @@ -3753,17 +3766,18 @@ async def _run_agent_stream( # A decision collected in an earlier request has no approval response in this # request's messages, so its result was never folded into them. Carry it into # the merge so it is persisted next to its call like the rest of the batch. - resolved_call_ids = { - str(content.call_id) + # Match by result object, not call id: distinct occurrences can share a call id. + folded_result_ids = { + id(content) for message in messages if get_role_value(message) == "tool" for content in message.contents or [] - if content.type == "function_result" and content.call_id + if content.type == "function_result" } carried_results = [ Message(role="tool", contents=[result]) for result in native_approval_results - if str(result.call_id) not in resolved_call_ids + if id(result) not in folded_result_ids ] _merge_resolved_approval_results_into_snapshot(snapshot_messages, [*messages, *carried_results]) flow.tool_results[:] = [ diff --git a/python/packages/ag-ui/tests/ag_ui/test_endpoint.py b/python/packages/ag-ui/tests/ag_ui/test_endpoint.py index 7fee1002195..18be1af0dab 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_endpoint.py +++ b/python/packages/ag-ui/tests/ag_ui/test_endpoint.py @@ -5244,10 +5244,12 @@ async def test_endpoint_agent_approval_resume_releases_already_approved_sibling( assert sorted(replayed_call_ids) == ["call_sensitive", "call_weather"] -async def test_endpoint_agent_approval_resume_distinguishes_hidden_siblings_with_reused_call_id( - streaming_chat_client_stub, -) -> None: - """Distinct hidden occurrences sharing a provider call ID resume and execute once.""" +def _build_reused_provider_call_id_endpoint( + streaming_chat_client_stub: Any, + *, + snapshot_store: InMemoryAGUIThreadSnapshotStore | None = None, +) -> tuple[TestClient, list[str], dict[str, str], AgentFrameworkAgent]: + """Pause on one guarded call whose two safe siblings reuse its provider call ID.""" executed: list[str] = [] state = {"phase": "pause"} @@ -5313,8 +5315,21 @@ async def stream_fn( ) wrapped_agent = AgentFrameworkAgent(agent=agent, require_confirmation=False) app = FastAPI() - add_agent_framework_fastapi_endpoint(app, wrapped_agent, path="/approval") - client = TestClient(app) + add_agent_framework_fastapi_endpoint( + app, + wrapped_agent, + path="/approval", + snapshot_store=snapshot_store, + snapshot_scope_resolver=(lambda _request: "test") if snapshot_store is not None else None, + ) + return TestClient(app), executed, state, wrapped_agent + + +async def test_endpoint_agent_approval_resume_distinguishes_hidden_siblings_with_reused_call_id( + streaming_chat_client_stub, +) -> None: + """Distinct hidden occurrences sharing a provider call ID resume and execute once.""" + client, executed, state, wrapped_agent = _build_reused_provider_call_id_endpoint(streaming_chat_client_stub) thread_id = "thread-shared-provider-call" pause_response = client.post( @@ -5375,6 +5390,52 @@ async def stream_fn( assert {occurrence.status for occurrence in occurrences} == {ApprovalStatus.SETTLED} +async def test_endpoint_approval_snapshot_keeps_every_result_for_reused_call_id(streaming_chat_client_stub) -> None: + """Results of distinct occurrences sharing a provider call ID all survive the snapshot round trip.""" + store = InMemoryAGUIThreadSnapshotStore() + client, executed, state, _ = _build_reused_provider_call_id_endpoint( + streaming_chat_client_stub, snapshot_store=store + ) + thread_id = "thread-shared-provider-call-snapshot" + expected_results = ["guarded result", "first safe result", "second safe result"] + + def post(run_id: str, *, text: str | None = None, interrupt_id: str | None = None) -> list[dict[str, Any]]: + payload: dict[str, Any] = { + "runId": run_id, + "threadId": thread_id, + "messages": [{"role": "user", "content": text}] if text else [], + } + if interrupt_id: + payload["resume"] = [{"interruptId": interrupt_id, "status": "resolved", "payload": {"accepted": True}}] + response = client.post("/approval", json=payload) + assert response.status_code == 200 + events = _decode_sse_events(response) + assert not [event for event in events if event["type"] == "RUN_ERROR"] + return events + + def tool_results(messages: list[dict[str, Any]]) -> list[tuple[str, str]]: + return [ + (str(message.get("toolCallId") or message.get("tool_call_id")), message["content"]) + for message in messages + if message.get("role") == "tool" + ] + + post("run-pause", text="Run all three tools") + state["phase"] = "resume" + resume_events = post("run-resume", interrupt_id="guarded-occurrence") + assert Counter(executed) == {"guarded": 1, "first-safe": 1, "second-safe": 1} + + snapshot_events = [event for event in resume_events if event["type"] == "MESSAGES_SNAPSHOT"] + assert snapshot_events + event_messages = snapshot_events[-1]["messages"] + assert tool_results(event_messages) == [("provider-shared", result) for result in expected_results] + assert [message["role"] for message in event_messages] == ["user", "assistant", "tool", "tool", "tool", "assistant"] + + saved = await store.get(scope="test", thread_id=thread_id) + assert saved is not None + assert tool_results(saved.messages) == [("provider-shared", result) for result in expected_results] + + def _build_fides_policy_approval_endpoint( streaming_chat_client_stub: Any, *, diff --git a/python/packages/ag-ui/tests/ag_ui/test_run_common.py b/python/packages/ag-ui/tests/ag_ui/test_run_common.py index 1841a0d37cc..ed6093577ce 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_run_common.py +++ b/python/packages/ag-ui/tests/ag_ui/test_run_common.py @@ -583,7 +583,7 @@ def test_mcp_host_payload_has_live_snapshot_and_approval_parity(self): approval_event = _make_approval_tool_result_events([content])[0] approval_snapshot = _resolved_tool_result_snapshot_messages( [Message(role="tool", contents=[content], message_id="approval-result")] - )["mcp-1"] + )["mcp-1"][0] assert json.loads(approval_event.content) == host_payload assert approval_snapshot["content"] == "Summary" assert json.loads(approval_snapshot[_AGUI_TOOL_RESULT_HOST_PAYLOAD_KEY]) == host_payload @@ -610,7 +610,7 @@ def test_empty_mcp_model_projection_remains_empty_across_live_snapshot_and_appro approval_event = _make_approval_tool_result_events([content])[0] approval_snapshot = _resolved_tool_result_snapshot_messages([Message(role="tool", contents=[content])])[ "mcp-empty" - ] + ][0] assert getattr(result_event, _AGUI_TOOL_RESULT_MODEL_CONTENT_KEY) == [] assert flow.tool_results[-1]["content"] == "" @@ -647,7 +647,7 @@ def test_explicit_display_payload_wins_over_mcp_host_projection(self): approval_event = _make_approval_tool_result_events([content])[0] approval_snapshot = _resolved_tool_result_snapshot_messages([Message(role="tool", contents=[content])])[ "mcp-display" - ] + ][0] assert json.loads(result_event.content) == display_payload # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] assert flow.tool_results[-1]["content"] == "Summary" From 07c6c69d63b394711d6268cb51be186a64643bb0 Mon Sep 17 00:00:00 2001 From: Shivam Sharma <192242476+ShivamSharma43@users.noreply.github.com> Date: Wed, 7 Oct 2026 00:40:20 +0530 Subject: [PATCH 3/7] Python: Keep one replayed tool result per reused call id occurrence The provider history sanitizer collapsed function calls that share a call id into a single pending id, so the first result consumed it and the results of later occurrences were dropped from replay. _deduplicate_messages did the same by keying tool results on call id alone. The sanitizer now tracks one pending entry per function call and clears one per result, while still injecting a single synthetic result per missing call id. Deduplication keeps as many results for a call id as the kept assistant messages declare occurrences of it, and still drops or fills in replays beyond that. Co-Authored-By: Claude Opus 5.5 --- .../_message_adapters.py | 62 +++++++-------- .../tests/ag_ui/test_message_adapters.py | 77 +++++++++++++++++++ 2 files changed, 106 insertions(+), 33 deletions(-) diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_message_adapters.py b/python/packages/ag-ui/agent_framework_ag_ui/_message_adapters.py index a03a142ea97..895ac9b4267 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_message_adapters.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_message_adapters.py @@ -8,6 +8,7 @@ import binascii import json import logging +from collections import Counter from typing import Any, cast, get_args from agent_framework import ( @@ -54,7 +55,8 @@ def _append_synthetic_tool_results( excluded_tool_call_ids: set[str] | None = None, ) -> None: excluded_tool_call_ids = excluded_tool_call_ids or set() - for pending_call_id in pending_tool_call_ids: + # A call ID awaiting several occurrences still gets a single synthetic result. + for pending_call_id in dict.fromkeys(pending_tool_call_ids): if pending_call_id in excluded_tool_call_ids: logger.info("Not injecting synthetic tool result for non-abandoned call_id=%s", pending_call_id) continue @@ -72,18 +74,9 @@ def _append_synthetic_tool_results( ) -def _ordered_unique_tool_call_ids(contents: list[Content]) -> list[str]: - tool_ids: list[str] = [] - seen: set[str] = set() - for content in contents: - if content.type != "function_call" or not content.call_id: - continue - tool_id = str(content.call_id) - if tool_id in seen: - continue - tool_ids.append(tool_id) - seen.add(tool_id) - return tool_ids +def _tool_call_ids_per_occurrence(contents: list[Content]) -> list[str]: + """Return one call ID per function call, so a call ID reused by several occurrences awaits each result.""" + return [str(content.call_id) for content in contents if content.type == "function_call" and content.call_id] def _function_result_call_ids(messages: list[Message]) -> set[str]: @@ -124,7 +117,7 @@ def _sanitize_tool_history( pending_tool_call_ids = None pending_confirm_changes_id = None - tool_ids = _ordered_unique_tool_call_ids(msg.contents or []) + tool_ids = _tool_call_ids_per_occurrence(msg.contents or []) confirm_changes_call = None for content in msg.contents or []: if content.type == "function_call" and content.name == "confirm_changes": @@ -264,12 +257,11 @@ def _sanitize_tool_history( call_id = str(content.call_id) if call_id in pending_tool_call_ids: keep = True - # Remove the call_id from pending since we now have its result. + # Remove one pending occurrence of the call_id since we now have its result. # This prevents duplicate synthetic "skipped" results from being - # injected when a user message arrives later. - pending_tool_call_ids = [ - pending_id for pending_id in pending_tool_call_ids if pending_id != call_id - ] + # injected when a user message arrives later, while later occurrences + # that reuse the call_id still wait for their own results. + pending_tool_call_ids.remove(call_id) if call_id == pending_confirm_changes_id: pending_confirm_changes_id = None break @@ -313,38 +305,41 @@ def _deduplicate_messages(messages: list[Message]) -> list[Message]: """Remove duplicate messages while preserving order.""" seen_keys: dict[Any, int] = {} unique_messages: list[Message] = [] + # A provider call ID can be reused by several function calls; keep one result per declared occurrence. + declared_call_occurrences: Counter[str] = Counter() + tool_result_indexes: dict[str, list[int]] = {} for idx, msg in enumerate(messages): role_value = get_role_value(msg) if role_value == "tool" and msg.contents and msg.contents[0].type == "function_result": call_id = str(msg.contents[0].call_id) - key: Any = (role_value, call_id) - - if key in seen_keys: - existing_idx = seen_keys[key] - existing_msg = unique_messages[existing_idx] + kept_indexes = tool_result_indexes.setdefault(call_id, []) - existing_result = None - if existing_msg.contents and existing_msg.contents[0].type == "function_result": - existing_result = existing_msg.contents[0].result + if len(kept_indexes) >= max(1, declared_call_occurrences[call_id]): new_result = msg.contents[0].result - - if (not existing_result or existing_result == "") and new_result: - logger.info(f"Replacing empty tool result at index {existing_idx} with data from index {idx}") - unique_messages[existing_idx] = msg + for existing_idx in kept_indexes: + existing_msg = unique_messages[existing_idx] + existing_result = None + if existing_msg.contents and existing_msg.contents[0].type == "function_result": + existing_result = existing_msg.contents[0].result + + if (not existing_result or existing_result == "") and new_result: + logger.info(f"Replacing empty tool result at index {existing_idx} with data from index {idx}") + unique_messages[existing_idx] = msg + break else: logger.info(f"Skipping duplicate tool result at index {idx}: call_id={call_id}") continue - seen_keys[key] = len(unique_messages) + kept_indexes.append(len(unique_messages)) unique_messages.append(msg) elif role_value == "assistant" and msg.contents and any(c.type == "function_call" for c in msg.contents): tool_call_ids = tuple( sorted(str(c.call_id) for c in msg.contents if c.type == "function_call" and c.call_id) ) - key = (role_value, tool_call_ids) + key: Any = (role_value, tool_call_ids) if key in seen_keys: logger.info(f"Skipping duplicate assistant tool call at index {idx}") @@ -352,6 +347,7 @@ def _deduplicate_messages(messages: list[Message]) -> list[Message]: seen_keys[key] = len(unique_messages) unique_messages.append(msg) + declared_call_occurrences.update(tool_call_ids) else: # Use message_id for deduplication when available — two messages with the diff --git a/python/packages/ag-ui/tests/ag_ui/test_message_adapters.py b/python/packages/ag-ui/tests/ag_ui/test_message_adapters.py index a08ab19af36..b3d689efc54 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_message_adapters.py +++ b/python/packages/ag-ui/tests/ag_ui/test_message_adapters.py @@ -2360,6 +2360,60 @@ def test_sanitize_history_ending_with_tool_call_closes_pending_call(): assert all("skipped" in str(msg.contents[0].result).lower() for msg in result[1:]) +def test_sanitize_history_keeps_every_result_for_reused_call_id(): + """Each function call occurrence sharing a provider call ID keeps its own result.""" + from agent_framework_ag_ui._message_adapters import _sanitize_tool_history + + for call_ids in [("shared",) * 3, ("call-a", "call-b", "call-c")]: + messages = [ + Message( + role="assistant", + contents=[ + Content.from_function_call(id=f"occurrence-{n}", call_id=call_id, name=f"tool_{n}", arguments="{}") + for n, call_id in enumerate(call_ids) + ], + ) + ] + messages += [ + Message(role="tool", contents=[Content.from_function_result(call_id=call_id, result=f"result-{n}")]) + for n, call_id in enumerate(call_ids) + ] + messages.append(Message(role="user", contents=[Content.from_text(text="Thanks")])) + + result = _sanitize_tool_history(messages) + + assert [msg.role for msg in result] == ["assistant", "tool", "tool", "tool", "user"] + assert [content.result for msg in result for content in msg.contents if content.type == "function_result"] == [ + "result-0", + "result-1", + "result-2", + ] + + +def test_sanitize_history_drops_results_beyond_reused_call_id_occurrences(): + """Results for a reused call ID beyond the declared occurrences are still dropped.""" + from agent_framework_ag_ui._message_adapters import _sanitize_tool_history + + assistant_msg = Message( + role="assistant", + contents=[ + Content.from_function_call(call_id="shared", name="tool_a", arguments="{}"), + Content.from_function_call(call_id="shared", name="tool_b", arguments="{}"), + ], + ) + tool_msgs = [ + Message(role="tool", contents=[Content.from_function_result(call_id="shared", result=f"result-{n}")]) + for n in range(3) + ] + + result = _sanitize_tool_history([assistant_msg, *tool_msgs]) + + assert [content.result for msg in result for content in msg.contents if content.type == "function_result"] == [ + "result-0", + "result-1", + ] + + def test_sanitize_tool_result_clears_pending_confirm(): """Tool result for pending confirm_changes call_id clears pending state.""" from agent_framework_ag_ui._message_adapters import _sanitize_tool_history @@ -2435,6 +2489,29 @@ def test_deduplicate_tool_results(): assert len(result) == 1 +def test_deduplicate_keeps_one_tool_result_per_reused_call_id_occurrence(): + """Results for a call ID reused by several function calls are kept once per occurrence.""" + from agent_framework_ag_ui._message_adapters import _deduplicate_messages + + assistant_msg = Message( + role="assistant", + contents=[Content.from_function_call(call_id="shared", name=f"tool_{n}", arguments="{}") for n in range(3)], + ) + tool_msgs = [ + Message(role="tool", contents=[Content.from_function_result(call_id="shared", result=f"result-{n}")]) + for n in range(3) + ] + replayed = Message(role="tool", contents=[Content.from_function_result(call_id="shared", result="result-0")]) + + result = _deduplicate_messages([assistant_msg, *tool_msgs, replayed]) + + assert [content.result for msg in result for content in msg.contents if content.type == "function_result"] == [ + "result-0", + "result-1", + "result-2", + ] + + def test_deduplicate_assistant_tool_calls(): """Duplicate assistant messages with same tool_calls are deduplicated.""" from agent_framework_ag_ui._message_adapters import _deduplicate_messages From d1d99cd25affa9dd7c7f8ad17d718258bf82ff4e Mon Sep 17 00:00:00 2001 From: Shivam Sharma <192242476+ShivamSharma43@users.noreply.github.com> Date: Fri, 9 Oct 2026 18:07:04 +0530 Subject: [PATCH 4/7] Python: Pair reused call id results with their own turn in AG-UI history A later turn that reuses a call id already completed earlier lost its pairing in two places: - _deduplicate_messages keyed assistant tool-call messages by call ids only, so the later turn's call was dropped as a duplicate and its result went with it. Assistant messages with a message_id are now deduplicated by id like other messages, and each declared occurrence opens one result slot, so a replayed result is still dropped. - _merge_resolved_approval_results_into_snapshot bucketed results by call id across the whole snapshot and drained the bucket at the first declaring assistant message, attaching the new result to the earlier call and keeping the old one twice. Results are now keyed by call id and declaration occurrence and placed after the matching assistant message. Co-Authored-By: Claude Opus 5.5 --- .../ag-ui/agent_framework_ag_ui/_agent_run.py | 88 ++++++++----- .../_message_adapters.py | 18 ++- .../ag-ui/tests/ag_ui/test_endpoint.py | 116 ++++++++++++++++++ .../tests/ag_ui/test_message_adapters.py | 34 +++++ .../ag-ui/tests/ag_ui/test_run_common.py | 6 +- 5 files changed, 226 insertions(+), 36 deletions(-) diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py b/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py index 13eb9a55e74..65faedef63d 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py @@ -2221,15 +2221,29 @@ def _snapshot_tool_call_ids(message: Mapping[str, Any]) -> list[str]: return call_ids -def _resolved_tool_result_snapshot_messages(resolved_messages: list[Message]) -> dict[str, list[dict[str, Any]]]: +def _resolved_tool_result_snapshot_messages( + resolved_messages: list[Message], +) -> dict[tuple[str, int], list[dict[str, Any]]]: """Build replayable AG-UI tool messages from resolved approval results. - Distinct function-call occurrences can reuse one provider call id, so results are - grouped by call id in resolution order rather than collapsed to one per id. + Results are keyed by call id and declaration occurrence: the number of assistant messages + declaring that call id up to the result (at least one). A later turn that reuses a call id + therefore keeps its result apart from the earlier turn's. Distinct function calls inside one + assistant message can also share a call id, so each key holds every result in resolution order. """ - results_by_call_id: dict[str, list[dict[str, Any]]] = {} + declarations: dict[str, int] = {} + results_by_occurrence: dict[tuple[str, int], list[dict[str, Any]]] = {} for msg in resolved_messages: - if get_role_value(msg) != "tool": + role = get_role_value(msg) + if role == "assistant": + for call_id in { + str(content.call_id) + for content in msg.contents or [] + if content.type == "function_call" and content.call_id + }: + declarations[call_id] = declarations.get(call_id, 0) + 1 + continue + if role != "tool": continue function_results = [ content for content in msg.contents or [] if content.type == "function_result" and content.call_id @@ -2252,8 +2266,26 @@ def _resolved_tool_result_snapshot_messages(resolved_messages: list[Message]) -> _model_items_for_agui_replay(content, llm_result), ) ) - results_by_call_id.setdefault(call_id, []).append(snapshot_message) - return results_by_call_id + occurrence = (call_id, max(1, declarations.get(call_id, 0))) + results_by_occurrence.setdefault(occurrence, []).append(snapshot_message) + return results_by_occurrence + + +def _snapshot_tool_message_occurrences(snapshot_messages: list[dict[str, Any]]) -> dict[int, tuple[str, int]]: + """Map each snapshot tool message to its call id and declaration occurrence.""" + declarations: dict[str, int] = {} + occurrences: dict[int, tuple[str, int]] = {} + for message in snapshot_messages: + role = normalize_agui_role(message.get("role", "")) + if role == "assistant": + for call_id in set(_snapshot_tool_call_ids(message)): + declarations[call_id] = declarations.get(call_id, 0) + 1 + elif role == "tool": + tool_call_id = message.get("toolCallId") or message.get("tool_call_id") + if tool_call_id: + call_id = str(tool_call_id) + occurrences[id(message)] = (call_id, max(1, declarations.get(call_id, 0))) + return occurrences def _merge_resolved_approval_results_into_snapshot( @@ -2262,46 +2294,46 @@ def _merge_resolved_approval_results_into_snapshot( ) -> None: """Persist approval-resolved tool results under their original tool call ids. - Every result for a call id follows the assistant message that declares it, in resolution - order. Sibling occurrences hidden from the client share the visible call's id, so one - declared call can be followed by several results. + The results of the n-th assistant message declaring a call id follow the n-th snapshot + assistant message declaring it, in resolution order, so a call id reused by a later turn + stays paired with its own call. Sibling occurrences hidden from the client share the + visible call's id, so one declared call can be followed by several results. """ - results_by_call_id = _resolved_tool_result_snapshot_messages(resolved_messages) - if not results_by_call_id: + results_by_occurrence = _resolved_tool_result_snapshot_messages(resolved_messages) + if not results_by_occurrence: snapshot_messages[:] = [message for message in snapshot_messages if not message.get("function_approvals")] return + tool_message_occurrences = _snapshot_tool_message_occurrences(snapshot_messages) + replaced_occurrences = set(results_by_occurrence) retained_host_messages: set[int] = set() for message in snapshot_messages: - if normalize_agui_role(message.get("role", "")) != "tool": + occurrence = tool_message_occurrences.get(id(message)) + if occurrence is None or message.get(_AGUI_MCP_TOOL_RESULT_KEY) is not True: continue - tool_call_id = message.get("toolCallId") or message.get("tool_call_id") - if not tool_call_id or message.get(_AGUI_MCP_TOOL_RESULT_KEY) is not True: - continue - queued_results = results_by_call_id.get(str(tool_call_id)) + queued_results = results_by_occurrence.get(occurrence) if queued_results and queued_results[0].get(_AGUI_MCP_TOOL_RESULT_KEY) is not True: # Keep the host-rendered result already in the snapshot for this occurrence. queued_results.pop(0) retained_host_messages.add(id(message)) - if not queued_results: - results_by_call_id.pop(str(tool_call_id)) + declarations: dict[str, int] = {} merged_messages: list[dict[str, Any]] = [] for message in snapshot_messages: if message.get("function_approvals"): continue - role = normalize_agui_role(message.get("role", "")) - if role == "tool" and id(message) not in retained_host_messages: - tool_call_id = message.get("toolCallId") or message.get("tool_call_id") - if tool_call_id and str(tool_call_id) in results_by_call_id: - continue + if id(message) not in retained_host_messages and tool_message_occurrences.get(id(message)) in ( + replaced_occurrences + ): + continue merged_messages.append(message) - if role != "assistant": + if normalize_agui_role(message.get("role", "")) != "assistant": continue - for call_id in _snapshot_tool_call_ids(message): - merged_messages.extend(results_by_call_id.pop(call_id, [])) + for call_id in dict.fromkeys(_snapshot_tool_call_ids(message)): + declarations[call_id] = declarations.get(call_id, 0) + 1 + merged_messages.extend(results_by_occurrence.pop((call_id, declarations[call_id]), [])) - for queued_results in results_by_call_id.values(): + for queued_results in results_by_occurrence.values(): merged_messages.extend(queued_results) snapshot_messages[:] = merged_messages diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_message_adapters.py b/python/packages/ag-ui/agent_framework_ag_ui/_message_adapters.py index 895ac9b4267..f1e8eb36172 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_message_adapters.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_message_adapters.py @@ -305,8 +305,11 @@ def _deduplicate_messages(messages: list[Message]) -> list[Message]: """Remove duplicate messages while preserving order.""" seen_keys: dict[Any, int] = {} unique_messages: list[Message] = [] - # A provider call ID can be reused by several function calls; keep one result per declared occurrence. - declared_call_occurrences: Counter[str] = Counter() + # A provider call ID can be reused by several function calls, in one assistant message or in a + # later turn. Each declared occurrence opens one result slot, so every occurrence keeps its own + # result while replayed results beyond the open slots are dropped. + open_result_slots: Counter[str] = Counter() + declared_call_ids: set[str] = set() tool_result_indexes: dict[str, list[int]] = {} for idx, msg in enumerate(messages): @@ -316,7 +319,9 @@ def _deduplicate_messages(messages: list[Message]) -> list[Message]: call_id = str(msg.contents[0].call_id) kept_indexes = tool_result_indexes.setdefault(call_id, []) - if len(kept_indexes) >= max(1, declared_call_occurrences[call_id]): + if open_result_slots[call_id] > 0: + open_result_slots[call_id] -= 1 + elif call_id in declared_call_ids or kept_indexes: new_result = msg.contents[0].result for existing_idx in kept_indexes: existing_msg = unique_messages[existing_idx] @@ -339,7 +344,9 @@ def _deduplicate_messages(messages: list[Message]) -> list[Message]: tool_call_ids = tuple( sorted(str(c.call_id) for c in msg.contents if c.type == "function_call" and c.call_id) ) - key: Any = (role_value, tool_call_ids) + # As for other messages, a shared message_id marks a replay of the same message, while a + # later turn reusing the same call IDs arrives as a distinct message. + key: Any = ("id", msg.message_id) if msg.message_id else (role_value, tool_call_ids) if key in seen_keys: logger.info(f"Skipping duplicate assistant tool call at index {idx}") @@ -347,7 +354,8 @@ def _deduplicate_messages(messages: list[Message]) -> list[Message]: seen_keys[key] = len(unique_messages) unique_messages.append(msg) - declared_call_occurrences.update(tool_call_ids) + open_result_slots.update(tool_call_ids) + declared_call_ids.update(tool_call_ids) else: # Use message_id for deduplication when available — two messages with the diff --git a/python/packages/ag-ui/tests/ag_ui/test_endpoint.py b/python/packages/ag-ui/tests/ag_ui/test_endpoint.py index f164936c20a..f88480e970f 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_endpoint.py +++ b/python/packages/ag-ui/tests/ag_ui/test_endpoint.py @@ -6983,6 +6983,122 @@ def tool_result_call_ids(messages: list[dict[str, Any]]) -> list[str]: assert [content.call_id for content in replayed_results] == ["call_first", "call_second"] +async def test_endpoint_approval_snapshot_pairs_result_with_later_turn_reusing_call_id(streaming_chat_client_stub): + """A later approved call reusing an earlier completed call ID keeps both pairs in place.""" + store = InMemoryAGUIThreadSnapshotStore() + executed: list[str] = [] + messages_received: list[Message] = [] + turn = {"number": 0} + + def lookup() -> str: + executed.append(f"turn-{turn['number']}") + return f"result {turn['number']}" + + async def stream_fn( + messages: list[Message], + options: dict[str, Any], + **kwargs: Any, + ) -> AsyncIterator[ChatResponseUpdate]: + del options, kwargs + messages_received[:] = list(messages) + if messages[-1].role == "user": + turn["number"] += 1 + yield ChatResponseUpdate( + contents=[Content.from_function_call(call_id="call_reused", name="lookup", arguments="{}")], + role="assistant", + ) + return + yield ChatResponseUpdate(contents=[Content.from_text(text=f"Done {turn['number']}.")], role="assistant") + + agent = Agent( + name="test_agent", + instructions="Test", + client=streaming_chat_client_stub(stream_fn), + tools=[FunctionTool(name="lookup", description="Lookup", func=lookup, approval_mode="always_require")], + middleware=[ToolApprovalMiddleware()], + ) + app = FastAPI() + add_agent_framework_fastapi_endpoint( + app, + AgentFrameworkAgent(agent=agent, require_confirmation=False), + path="/approval", + snapshot_store=store, + snapshot_scope_resolver=lambda _request: "test", + ) + client = TestClient(app) + thread_id = "thread-reused-call-id-across-turns" + + def post(run_id: str, *, text: str | None = None, approval_id: str | None = None) -> list[dict[str, Any]]: + payload: dict[str, Any] = { + "runId": run_id, + "threadId": thread_id, + "messages": [{"role": "user", "content": text}] if text else [], + } + if approval_id: + payload["resume"] = [{"interruptId": approval_id, "status": "resolved", "payload": {"accepted": True}}] + response = client.post("/approval", json=payload) + assert response.status_code == 200 + events = _decode_sse_events(response) + assert not [event for event in events if event["type"] == "RUN_ERROR"] + return events + + def approval_id(events: list[dict[str, Any]]) -> str: + finished = [event for event in events if event["type"] == "RUN_FINISHED"] + return _run_finished_interrupts(finished[-1])[0]["id"] + + def transcript(messages: list[dict[str, Any]]) -> list[tuple[str, str]]: + entries: list[tuple[str, str]] = [] + for message in messages: + role = message["role"] + if role == "tool": + entries.append( + ("tool", f"{message.get('toolCallId') or message.get('tool_call_id')}={message['content']}") + ) + elif role == "assistant" and (message.get("toolCalls") or message.get("tool_calls")): + tool_calls = message.get("toolCalls") or message.get("tool_calls") + entries.append(("call", ",".join(tool_call["id"] for tool_call in tool_calls))) + else: + entries.append((role, message.get("content") or "")) + return entries + + expected = [ + ("user", "First"), + ("call", "call_reused"), + ("tool", "call_reused=result 1"), + ("assistant", "Done 1."), + ("user", "Second"), + ("call", "call_reused"), + ("tool", "call_reused=result 2"), + ("assistant", "Done 2."), + ] + + post("run-first-resume", approval_id=approval_id(post("run-first", text="First"))) + final_events = post("run-second-resume", approval_id=approval_id(post("run-second", text="Second"))) + assert executed == ["turn-1", "turn-2"] + + snapshot_events = [event for event in final_events if event["type"] == "MESSAGES_SNAPSHOT"] + assert snapshot_events + assert transcript(snapshot_events[-1]["messages"]) == expected + + saved = await store.get(scope="test", thread_id=thread_id) + assert saved is not None + assert transcript(saved.messages) == expected + + post("run-third", text="Third") + replayed = [ + (content.type, content.call_id, content.result if content.type == "function_result" else None) + for message in messages_received + for content in message.contents + if content.type in {"function_call", "function_result"} + ] + assert replayed == [ + ("function_call", "call_reused", None), + ("function_result", "call_reused", "result 1"), + ("function_call", "call_reused", None), + ("function_result", "call_reused", "result 2"), + ] + + @pytest.mark.parametrize("retry_time", [25.0, 33.0]) async def test_endpoint_collected_approval_preserves_original_expiration( streaming_chat_client_stub, retry_time: float diff --git a/python/packages/ag-ui/tests/ag_ui/test_message_adapters.py b/python/packages/ag-ui/tests/ag_ui/test_message_adapters.py index b3d689efc54..f67c71061af 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_message_adapters.py +++ b/python/packages/ag-ui/tests/ag_ui/test_message_adapters.py @@ -2512,6 +2512,40 @@ def test_deduplicate_keeps_one_tool_result_per_reused_call_id_occurrence(): ] +def test_deduplicate_keeps_later_turn_reusing_call_id(): + """A later turn reusing a completed call ID keeps its own call and result; a replay is still dropped.""" + from agent_framework_ag_ui._message_adapters import _deduplicate_messages + + def call(message_id: str) -> Message: + return Message( + role="assistant", + contents=[Content.from_function_call(call_id="reused", name="lookup", arguments="{}")], + message_id=message_id, + ) + + def result(message_id: str, value: str) -> Message: + return Message( + role="tool", + contents=[Content.from_function_result(call_id="reused", result=value)], + message_id=message_id, + ) + + messages = [ + Message(role="user", contents=[Content.from_text(text="First")], message_id="user-1"), + call("call-1"), + result("result-1", "first"), + Message(role="user", contents=[Content.from_text(text="Second")], message_id="user-2"), + call("call-2"), + result("result-2", "second"), + call("call-2"), + result("result-2-replayed", "second"), + ] + + deduplicated = _deduplicate_messages(messages) + + assert [msg.message_id for msg in deduplicated] == ["user-1", "call-1", "result-1", "user-2", "call-2", "result-2"] + + def test_deduplicate_assistant_tool_calls(): """Duplicate assistant messages with same tool_calls are deduplicated.""" from agent_framework_ag_ui._message_adapters import _deduplicate_messages diff --git a/python/packages/ag-ui/tests/ag_ui/test_run_common.py b/python/packages/ag-ui/tests/ag_ui/test_run_common.py index cb7266dce7c..69822014420 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_run_common.py +++ b/python/packages/ag-ui/tests/ag_ui/test_run_common.py @@ -620,7 +620,7 @@ def test_mcp_host_payload_has_live_snapshot_and_approval_parity(self): approval_event = _make_approval_tool_result_events([content])[0] approval_snapshot = _resolved_tool_result_snapshot_messages( [Message(role="tool", contents=[content], message_id="approval-result")] - )["mcp-1"][0] + )[("mcp-1", 1)][0] assert json.loads(approval_event.content) == host_payload assert approval_snapshot["content"] == "Summary" assert json.loads(approval_snapshot[_AGUI_TOOL_RESULT_HOST_PAYLOAD_KEY]) == host_payload @@ -646,7 +646,7 @@ def test_empty_mcp_model_projection_remains_empty_across_live_snapshot_and_appro snapshot = _build_messages_snapshot(flow, []).messages[-1].model_dump(by_alias=True, exclude_none=True) approval_event = _make_approval_tool_result_events([content])[0] approval_snapshot = _resolved_tool_result_snapshot_messages([Message(role="tool", contents=[content])])[ - "mcp-empty" + ("mcp-empty", 1) ][0] assert getattr(result_event, _AGUI_TOOL_RESULT_MODEL_CONTENT_KEY) == [] @@ -683,7 +683,7 @@ def test_explicit_display_payload_wins_over_mcp_host_projection(self): ) approval_event = _make_approval_tool_result_events([content])[0] approval_snapshot = _resolved_tool_result_snapshot_messages([Message(role="tool", contents=[content])])[ - "mcp-display" + ("mcp-display", 1) ][0] assert json.loads(result_event.content) == display_payload # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] From e016689982d0b0a6e2f7bdac2f1c22911d7c555d Mon Sep 17 00:00:00 2001 From: Shivam Sharma <192242476+ShivamSharma43@users.noreply.github.com> Date: Fri, 9 Oct 2026 18:11:18 +0530 Subject: [PATCH 5/7] Python: Skip abandoned AG-UI calls that reuse an earlier turn's call id _sanitize_tool_history treated a call id as answered when a result for it appeared anywhere in history, so a later turn reusing a completed call id and then abandoning it got no synthetic skipped result and replayed as an unpaired call. Only results from the flush point on, before the call id is declared again, now count as answering the pending occurrence. Co-Authored-By: Claude Opus 5.5 --- .../_message_adapters.py | 36 +++++++++++++------ .../tests/ag_ui/test_message_adapters.py | 26 ++++++++++++++ 2 files changed, 52 insertions(+), 10 deletions(-) diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_message_adapters.py b/python/packages/ag-ui/agent_framework_ag_ui/_message_adapters.py index f1e8eb36172..e00b4a389d0 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_message_adapters.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_message_adapters.py @@ -79,12 +79,24 @@ def _tool_call_ids_per_occurrence(contents: list[Content]) -> list[str]: return [str(content.call_id) for content in contents if content.type == "function_call" and content.call_id] -def _function_result_call_ids(messages: list[Message]) -> set[str]: +def _answered_call_ids_from(messages: list[Message], start: int) -> set[str]: + """Return call IDs that get a function result from ``start`` on, before the call ID is declared again. + + A result from an earlier turn belongs to that turn's occurrence, and a result after the call ID + is declared again belongs to the later occurrence, so neither answers the pending occurrence. + """ result_ids: set[str] = set() - for msg in messages: + redeclared_ids: set[str] = set() + for msg in messages[start:]: + is_assistant = get_role_value(msg) == "assistant" for content in msg.contents or []: - if content.type == "function_result" and content.call_id: - result_ids.add(str(content.call_id)) + if not content.call_id: + continue + call_id = str(content.call_id) + if content.type == "function_call" and is_assistant: + redeclared_ids.add(call_id) + elif content.type == "function_result" and call_id not in redeclared_ids: + result_ids.add(call_id) return result_ids @@ -97,9 +109,13 @@ def _sanitize_tool_history( sanitized: list[Message] = [] pending_tool_call_ids: list[str] | None = None pending_confirm_changes_id: str | None = None - non_abandoned_tool_call_ids = set(protected_tool_call_ids or set()) | _function_result_call_ids(messages) + protected_ids = set(protected_tool_call_ids or set()) - for msg in messages: + def non_abandoned_tool_call_ids(start: int) -> set[str]: + """Protected calls and calls answered from ``start`` on; earlier turns do not count.""" + return protected_ids | _answered_call_ids_from(messages, start) + + for idx, msg in enumerate(messages): role_value = get_role_value(msg) if role_value == "assistant": @@ -112,7 +128,7 @@ def _sanitize_tool_history( sanitized, pending_tool_call_ids, "Tool execution skipped - assistant continued before the tool result was available.", - excluded_tool_call_ids=non_abandoned_tool_call_ids, + excluded_tool_call_ids=non_abandoned_tool_call_ids(idx), ) pending_tool_call_ids = None pending_confirm_changes_id = None @@ -239,7 +255,7 @@ def _sanitize_tool_history( sanitized, pending_tool_call_ids, "Tool execution skipped - user provided follow-up message", - excluded_tool_call_ids=non_abandoned_tool_call_ids, + excluded_tool_call_ids=non_abandoned_tool_call_ids(idx + 1), ) pending_tool_call_ids = None pending_confirm_changes_id = None @@ -279,7 +295,7 @@ def _sanitize_tool_history( sanitized, pending_tool_call_ids, "Tool execution skipped - conversation continued before the tool result was available.", - excluded_tool_call_ids=non_abandoned_tool_call_ids, + excluded_tool_call_ids=non_abandoned_tool_call_ids(idx), ) sanitized.append(msg) @@ -295,7 +311,7 @@ def _sanitize_tool_history( sanitized, pending_tool_call_ids, "Tool execution skipped - conversation ended before the tool result was available.", - excluded_tool_call_ids=non_abandoned_tool_call_ids, + excluded_tool_call_ids=non_abandoned_tool_call_ids(len(messages)), ) return sanitized diff --git a/python/packages/ag-ui/tests/ag_ui/test_message_adapters.py b/python/packages/ag-ui/tests/ag_ui/test_message_adapters.py index f67c71061af..fccf40e236c 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_message_adapters.py +++ b/python/packages/ag-ui/tests/ag_ui/test_message_adapters.py @@ -2274,6 +2274,32 @@ def test_sanitize_pending_tool_skip_on_user_followup(): assert "skipped" in str(tool_results[0].contents[0].result).lower() +def test_sanitize_skips_abandoned_later_turn_reusing_completed_call_id(): + """A result from an earlier turn does not stop a later abandoned call reusing its ID from being skipped.""" + from agent_framework_ag_ui._message_adapters import _sanitize_tool_history + + def call() -> Message: + return Message( + role="assistant", + contents=[Content.from_function_call(call_id="reused", name="get_weather", arguments="{}")], + ) + + messages = [ + Message(role="user", contents=[Content.from_text(text="First")]), + call(), + Message(role="tool", contents=[Content.from_function_result(call_id="reused", result="sunny")]), + Message(role="user", contents=[Content.from_text(text="Second")]), + call(), + Message(role="user", contents=[Content.from_text(text="Actually, never mind")]), + ] + + result = _sanitize_tool_history(messages) + + assert [msg.role for msg in result] == ["user", "assistant", "tool", "user", "assistant", "tool", "user"] + assert result[2].contents[0].result == "sunny" + assert "skipped" in str(result[5].contents[0].result).lower() + + def test_sanitize_pending_tool_does_not_skip_server_owned_approval_call(): """Server-owned approval state means a pending tool call is not abandoned.""" from agent_framework_ag_ui._message_adapters import _sanitize_tool_history From 5647ee9b0565c1205be0b8492edd577cbad0f54e Mon Sep 17 00:00:00 2001 From: Shivam Sharma <192242476+ShivamSharma43@users.noreply.github.com> Date: Sun, 11 Oct 2026 20:19:42 +0530 Subject: [PATCH 6/7] Python: Narrow tool_calls type in reused-call-id endpoint test Fixes the reportOptionalIterable error from the test typing checks. Co-Authored-By: Claude Opus 5.5 --- python/packages/ag-ui/tests/ag_ui/test_endpoint.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/python/packages/ag-ui/tests/ag_ui/test_endpoint.py b/python/packages/ag-ui/tests/ag_ui/test_endpoint.py index f88480e970f..232535c0e5b 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_endpoint.py +++ b/python/packages/ag-ui/tests/ag_ui/test_endpoint.py @@ -7054,8 +7054,7 @@ def transcript(messages: list[dict[str, Any]]) -> list[tuple[str, str]]: entries.append( ("tool", f"{message.get('toolCallId') or message.get('tool_call_id')}={message['content']}") ) - elif role == "assistant" and (message.get("toolCalls") or message.get("tool_calls")): - tool_calls = message.get("toolCalls") or message.get("tool_calls") + elif role == "assistant" and (tool_calls := message.get("toolCalls") or message.get("tool_calls")): entries.append(("call", ",".join(tool_call["id"] for tool_call in tool_calls))) else: entries.append((role, message.get("content") or "")) From 100f7c83267c892f3c4a5836d8a65bde35f15032 Mon Sep 17 00:00:00 2001 From: Shivam Sharma <192242476+ShivamSharma43@users.noreply.github.com> Date: Mon, 12 Oct 2026 15:33:21 +0530 Subject: [PATCH 7/7] Python: Don't let settled approvals protect a later reused AG-UI call id _approval_state_tool_call_ids returned the call ids of every approval occurrence the lifecycle still holds, including terminal ones that are only retained for replay. A later turn that reuses such a call id and is then abandoned was treated as server-owned, so _sanitize_tool_history never added a skipped result and the next request replayed an unpaired call. Only non-terminal occurrences (and queued approval state) now protect their call ids. Co-Authored-By: Claude Opus 5.5 --- .../ag-ui/agent_framework_ag_ui/_agent_run.py | 4 + .../ag-ui/tests/ag_ui/test_endpoint.py | 100 ++++++++++++++++++ 2 files changed, 104 insertions(+) diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py b/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py index 65faedef63d..6b78337291b 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py @@ -1096,6 +1096,10 @@ def _approval_state_tool_call_ids( """Return server-owned approval call ids that are not abandoned.""" call_ids: set[str] = set() for occurrence in approval_state_store.lifecycle.occurrences_for_thread(thread_id=thread_id): + # Terminal occurrences are retained only for replay. Their call ID no longer belongs to an + # open call, so a later turn reusing it can still be abandoned and needs a skipped result. + if occurrence.status.is_terminal: + continue call_ids.add(occurrence.identity.call_id) call_ids.add(occurrence.function_call_id) call_ids.update(occurrence.aliases) diff --git a/python/packages/ag-ui/tests/ag_ui/test_endpoint.py b/python/packages/ag-ui/tests/ag_ui/test_endpoint.py index 232535c0e5b..a69ed8d3f12 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_endpoint.py +++ b/python/packages/ag-ui/tests/ag_ui/test_endpoint.py @@ -7098,6 +7098,106 @@ def transcript(messages: list[dict[str, Any]]) -> list[tuple[str, str]]: ] +async def test_endpoint_skips_abandoned_later_call_reusing_settled_approval_call_id(streaming_chat_client_stub): + """A settled approval retained for replay does not protect a later abandoned call that reuses its ID.""" + store = InMemoryAGUIThreadSnapshotStore() + executed: list[str] = [] + messages_received: list[Message] = [] + + def lookup() -> str: + executed.append("lookup") + return "result 1" + + async def stream_fn( + messages: list[Message], + options: dict[str, Any], + **kwargs: Any, + ) -> AsyncIterator[ChatResponseUpdate]: + del options, kwargs + messages_received[:] = list(messages) + last = messages[-1] + if last.role == "user" and last.text == "First": + yield ChatResponseUpdate( + contents=[Content.from_function_call(call_id="call_reused", name="lookup", arguments="{}")], + role="assistant", + ) + return + if last.role == "user" and last.text == "Second": + yield ChatResponseUpdate( + contents=[Content.from_function_call(call_id="call_reused", name="pick_color", arguments="{}")], + role="assistant", + ) + return + yield ChatResponseUpdate(contents=[Content.from_text(text="Done.")], role="assistant") + + agent = Agent( + name="test_agent", + instructions="Test", + client=streaming_chat_client_stub(stream_fn), + tools=[FunctionTool(name="lookup", description="Lookup", func=lookup, approval_mode="always_require")], + middleware=[ToolApprovalMiddleware()], + ) + wrapped_agent = AgentFrameworkAgent(agent=agent, require_confirmation=False) + app = FastAPI() + add_agent_framework_fastapi_endpoint( + app, + wrapped_agent, + path="/approval", + snapshot_store=store, + snapshot_scope_resolver=lambda _request: "test", + ) + client = TestClient(app) + thread_id = "thread-reused-call-id-abandoned" + client_tool = {"name": "pick_color", "description": "Pick a color", "parameters": {"type": "object"}} + + def post(run_id: str, *, text: str | None = None, approval_id: str | None = None) -> list[dict[str, Any]]: + payload: dict[str, Any] = { + "runId": run_id, + "threadId": thread_id, + "messages": [{"role": "user", "content": text}] if text else [], + "tools": [client_tool], + } + if approval_id: + payload["resume"] = [{"interruptId": approval_id, "status": "resolved", "payload": {"accepted": True}}] + response = client.post("/approval", json=payload) + assert response.status_code == 200 + events = _decode_sse_events(response) + assert not [event for event in events if event["type"] == "RUN_ERROR"] + return events + + first_events = post("run-first", text="First") + first_finished = [event for event in first_events if event["type"] == "RUN_FINISHED"] + post("run-first-resume", approval_id=_run_finished_interrupts(first_finished[-1])[0]["id"]) + assert executed == ["lookup"] + + # The settled approval stays in the lifecycle for replay while the next turn reuses its call ID. + occurrences = wrapped_agent._approval_state_store.lifecycle.occurrences_for_thread( + thread_id=approval_state_thread_id(scope="test", thread_id=thread_id) + ) + assert [(occurrence.identity.call_id, occurrence.status.is_terminal) for occurrence in occurrences] == [ + ("call_reused", True) + ] + + post("run-second", text="Second") + # The client never answers the client tool call and moves on. + post("run-third", text="Never mind") + + replayed = [ + (content.type, content.call_id, content.result if content.type == "function_result" else None) + for message in messages_received + for content in message.contents + if content.type in {"function_call", "function_result"} + ] + assert [entry[:2] for entry in replayed] == [ + ("function_call", "call_reused"), + ("function_result", "call_reused"), + ("function_call", "call_reused"), + ("function_result", "call_reused"), + ] + assert replayed[1][2] == "result 1" + assert "Tool execution skipped" in str(replayed[3][2]) + + @pytest.mark.parametrize("retry_time", [25.0, 33.0]) async def test_endpoint_collected_approval_preserves_original_expiration( streaming_chat_client_stub, retry_time: float