Python: Keep every queued approval result in the AG-UI thread snapshot

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)
@@ -2221,11 +2225,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, 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[tuple[str, int], list[dict[str, Any]]]:
+    """Build replayable AG-UI tool messages from resolved approval results.
+
+    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.
+    """
+    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
@@ -2248,48 +2270,75 @@ 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
+            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(
     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.
+
+    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_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":
-            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:
+        occurrence = tool_message_occurrences.get(id(message))
+        if occurrence is None 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_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))
 
+    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":
-            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:
-                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):
-            result_message = result_by_call_id.pop(call_id, None)
-            if result_message is not None:
-                merged_messages.append(result_message)
+        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]), []))
 
-    merged_messages.extend(result_by_call_id.values())
+    for queued_results in results_by_occurrence.values():
+        merged_messages.extend(queued_results)
     snapshot_messages[:] = merged_messages
 
 
@@ -3588,6 +3637,7 @@ async def save_safe_point_snapshot() -> None:
     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():
@@ -3682,6 +3732,7 @@ async def save_safe_point_snapshot() -> None:
                         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]]
@@ -3949,7 +4000,23 @@ async def save_safe_point_snapshot() -> None:
     # 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.
+        # 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"
+        }
+        carried_results = [
+            Message(role="tool", contents=[result])
+            for result in native_approval_results
+            if id(result) not in folded_result_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
         ]

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,26 +74,29 @@ 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 _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.
 
-def _function_result_call_ids(messages: list[Message]) -> set[str]:
+    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
 
 
@@ -104,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":
@@ -119,12 +128,12 @@ 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
 
-            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":
@@ -246,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
@@ -264,12 +273,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
@@ -287,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)
@@ -303,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
@@ -313,45 +321,57 @@ 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, 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):
         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 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
-
-                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)
+            # 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}")
                 continue
 
             seen_keys[key] = len(unique_messages)
             unique_messages.append(msg)
+            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