Skip to content

Commit 5ffe5ca

Browse files
ShivamSharma43claude
ShivamSharma43
and
claude
authored
Python: Keep every queued approval result in the AG-UI thread snapshot (#8851)
* 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 <noreply@anthropic.com> * 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 <noreply@anthropic.com> * 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 <noreply@anthropic.com> * 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 <noreply@anthropic.com> * 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 <noreply@anthropic.com> * 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 <noreply@anthropic.com> * 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 <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
1 parent c9b237c commit 5ffe5ca

5 files changed

Lines changed: 642 additions & 82 deletions

File tree

‎python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py‎

Lines changed: 95 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -1109,6 +1109,10 @@ def _approval_state_tool_call_ids(
11091109
"""Return server-owned approval call ids that are not abandoned."""
11101110
call_ids: set[str] = set()
11111111
for occurrence in approval_state_store.lifecycle.occurrences_for_thread(thread_id=thread_id):
1112+
# Terminal occurrences are retained only for replay. Their call ID no longer belongs to an
1113+
# open call, so a later turn reusing it can still be abandoned and needs a skipped result.
1114+
if occurrence.status.is_terminal:
1115+
continue
11121116
call_ids.add(occurrence.identity.call_id)
11131117
call_ids.add(occurrence.function_call_id)
11141118
call_ids.update(occurrence.aliases)
@@ -2238,11 +2242,29 @@ def _snapshot_tool_call_ids(message: Mapping[str, Any]) -> list[str]:
22382242
return call_ids
22392243

22402244

2241-
def _resolved_tool_result_snapshot_messages(resolved_messages: list[Message]) -> dict[str, dict[str, Any]]:
2242-
"""Build replayable AG-UI tool messages from resolved approval results."""
2243-
result_by_call_id: dict[str, dict[str, Any]] = {}
2245+
def _resolved_tool_result_snapshot_messages(
2246+
resolved_messages: list[Message],
2247+
) -> dict[tuple[str, int], list[dict[str, Any]]]:
2248+
"""Build replayable AG-UI tool messages from resolved approval results.
2249+
2250+
Results are keyed by call id and declaration occurrence: the number of assistant messages
2251+
declaring that call id up to the result (at least one). A later turn that reuses a call id
2252+
therefore keeps its result apart from the earlier turn's. Distinct function calls inside one
2253+
assistant message can also share a call id, so each key holds every result in resolution order.
2254+
"""
2255+
declarations: dict[str, int] = {}
2256+
results_by_occurrence: dict[tuple[str, int], list[dict[str, Any]]] = {}
22442257
for msg in resolved_messages:
2245-
if get_role_value(msg) != "tool":
2258+
role = get_role_value(msg)
2259+
if role == "assistant":
2260+
for call_id in {
2261+
str(content.call_id)
2262+
for content in msg.contents or []
2263+
if content.type == "function_call" and content.call_id
2264+
}:
2265+
declarations[call_id] = declarations.get(call_id, 0) + 1
2266+
continue
2267+
if role != "tool":
22462268
continue
22472269
function_results = [
22482270
content for content in msg.contents or [] if content.type == "function_result" and content.call_id
@@ -2265,48 +2287,75 @@ def _resolved_tool_result_snapshot_messages(resolved_messages: list[Message]) ->
22652287
_model_items_for_agui_replay(content, llm_result),
22662288
)
22672289
)
2268-
result_by_call_id[call_id] = snapshot_message
2269-
return result_by_call_id
2290+
occurrence = (call_id, max(1, declarations.get(call_id, 0)))
2291+
results_by_occurrence.setdefault(occurrence, []).append(snapshot_message)
2292+
return results_by_occurrence
2293+
2294+
2295+
def _snapshot_tool_message_occurrences(snapshot_messages: list[dict[str, Any]]) -> dict[int, tuple[str, int]]:
2296+
"""Map each snapshot tool message to its call id and declaration occurrence."""
2297+
declarations: dict[str, int] = {}
2298+
occurrences: dict[int, tuple[str, int]] = {}
2299+
for message in snapshot_messages:
2300+
role = normalize_agui_role(message.get("role", ""))
2301+
if role == "assistant":
2302+
for call_id in set(_snapshot_tool_call_ids(message)):
2303+
declarations[call_id] = declarations.get(call_id, 0) + 1
2304+
elif role == "tool":
2305+
tool_call_id = message.get("toolCallId") or message.get("tool_call_id")
2306+
if tool_call_id:
2307+
call_id = str(tool_call_id)
2308+
occurrences[id(message)] = (call_id, max(1, declarations.get(call_id, 0)))
2309+
return occurrences
22702310

22712311

22722312
def _merge_resolved_approval_results_into_snapshot(
22732313
snapshot_messages: list[dict[str, Any]],
22742314
resolved_messages: list[Message],
22752315
) -> None:
2276-
"""Persist approval-resolved tool results under their original tool call ids."""
2277-
result_by_call_id = _resolved_tool_result_snapshot_messages(resolved_messages)
2278-
if not result_by_call_id:
2316+
"""Persist approval-resolved tool results under their original tool call ids.
2317+
2318+
The results of the n-th assistant message declaring a call id follow the n-th snapshot
2319+
assistant message declaring it, in resolution order, so a call id reused by a later turn
2320+
stays paired with its own call. Sibling occurrences hidden from the client share the
2321+
visible call's id, so one declared call can be followed by several results.
2322+
"""
2323+
results_by_occurrence = _resolved_tool_result_snapshot_messages(resolved_messages)
2324+
if not results_by_occurrence:
22792325
snapshot_messages[:] = [message for message in snapshot_messages if not message.get("function_approvals")]
22802326
return
22812327

2328+
tool_message_occurrences = _snapshot_tool_message_occurrences(snapshot_messages)
2329+
replaced_occurrences = set(results_by_occurrence)
2330+
retained_host_messages: set[int] = set()
22822331
for message in snapshot_messages:
2283-
if normalize_agui_role(message.get("role", "")) != "tool":
2284-
continue
2285-
tool_call_id = message.get("toolCallId") or message.get("tool_call_id")
2286-
if not tool_call_id or message.get(_AGUI_MCP_TOOL_RESULT_KEY) is not True:
2332+
occurrence = tool_message_occurrences.get(id(message))
2333+
if occurrence is None or message.get(_AGUI_MCP_TOOL_RESULT_KEY) is not True:
22872334
continue
2288-
replacement = result_by_call_id.get(str(tool_call_id))
2289-
if replacement is not None and replacement.get(_AGUI_MCP_TOOL_RESULT_KEY) is not True:
2290-
result_by_call_id.pop(str(tool_call_id))
2335+
queued_results = results_by_occurrence.get(occurrence)
2336+
if queued_results and queued_results[0].get(_AGUI_MCP_TOOL_RESULT_KEY) is not True:
2337+
# Keep the host-rendered result already in the snapshot for this occurrence.
2338+
queued_results.pop(0)
2339+
retained_host_messages.add(id(message))
22912340

2341+
declarations: dict[str, int] = {}
22922342
merged_messages: list[dict[str, Any]] = []
22932343
for message in snapshot_messages:
22942344
if message.get("function_approvals"):
22952345
continue
2296-
role = normalize_agui_role(message.get("role", ""))
2297-
if role == "tool":
2298-
tool_call_id = message.get("toolCallId") or message.get("tool_call_id")
2299-
if tool_call_id and str(tool_call_id) in result_by_call_id:
2300-
continue
2346+
if id(message) not in retained_host_messages and tool_message_occurrences.get(id(message)) in (
2347+
replaced_occurrences
2348+
):
2349+
continue
23012350
merged_messages.append(message)
2302-
if role != "assistant":
2351+
if normalize_agui_role(message.get("role", "")) != "assistant":
23032352
continue
2304-
for call_id in _snapshot_tool_call_ids(message):
2305-
result_message = result_by_call_id.pop(call_id, None)
2306-
if result_message is not None:
2307-
merged_messages.append(result_message)
2353+
for call_id in dict.fromkeys(_snapshot_tool_call_ids(message)):
2354+
declarations[call_id] = declarations.get(call_id, 0) + 1
2355+
merged_messages.extend(results_by_occurrence.pop((call_id, declarations[call_id]), []))
23082356

2309-
merged_messages.extend(result_by_call_id.values())
2357+
for queued_results in results_by_occurrence.values():
2358+
merged_messages.extend(queued_results)
23102359
snapshot_messages[:] = merged_messages
23112360

23122361

@@ -3686,6 +3735,7 @@ async def save_safe_point_snapshot() -> None:
36863735
telemetry_context = partial(_use_telemetry_conversation_id, telemetry_conversation_id)
36873736
stream_completed = False
36883737
native_approval_flow_result_ids: set[int] = set()
3738+
native_approval_results: list[Content] = []
36893739
try:
36903740
with telemetry_context():
36913741
for queued_executions in forwarded_executions.values():
@@ -3780,6 +3830,7 @@ async def save_safe_point_snapshot() -> None:
37803830
and intent.owner is not ApprovalExecutionOwner.HOSTED
37813831
):
37823832
native_approval_result = True
3833+
native_approval_results.append(content)
37833834
suppress_approval_call_end = intent.owner is ApprovalExecutionOwner.LOCAL
37843835
_replace_approval_contents_with_results(
37853836
messages, _collect_approval_responses(messages), [[content]]
@@ -4047,7 +4098,23 @@ async def save_safe_point_snapshot() -> None:
40474098
# Feature #5: Suppress intermediate snapshots for predictive tools without confirmation
40484099
_clean_resolved_approvals_from_snapshot(snapshot_messages, messages)
40494100
if native_approval_flow_result_ids:
4050-
_merge_resolved_approval_results_into_snapshot(snapshot_messages, messages)
4101+
# A decision collected in an earlier request has no approval response in this
4102+
# request's messages, so its result was never folded into them. Carry it into
4103+
# the merge so it is persisted next to its call like the rest of the batch.
4104+
# Match by result object, not call id: distinct occurrences can share a call id.
4105+
folded_result_ids = {
4106+
id(content)
4107+
for message in messages
4108+
if get_role_value(message) == "tool"
4109+
for content in message.contents or []
4110+
if content.type == "function_result"
4111+
}
4112+
carried_results = [
4113+
Message(role="tool", contents=[result])
4114+
for result in native_approval_results
4115+
if id(result) not in folded_result_ids
4116+
]
4117+
_merge_resolved_approval_results_into_snapshot(snapshot_messages, [*messages, *carried_results])
40514118
flow.tool_results[:] = [
40524119
result for result in flow.tool_results if id(result) not in native_approval_flow_result_ids
40534120
]

‎python/packages/ag-ui/agent_framework_ag_ui/_message_adapters.py‎

Lines changed: 63 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
import copy
1010
import json
1111
import logging
12+
from collections import Counter
1213
from typing import Any, cast, get_args
1314
from urllib.parse import urlsplit
1415

@@ -58,7 +59,8 @@ def _append_synthetic_tool_results(
5859
excluded_tool_call_ids: set[str] | None = None,
5960
) -> None:
6061
excluded_tool_call_ids = excluded_tool_call_ids or set()
61-
for pending_call_id in pending_tool_call_ids:
62+
# A call ID awaiting several occurrences still gets a single synthetic result.
63+
for pending_call_id in dict.fromkeys(pending_tool_call_ids):
6264
if pending_call_id in excluded_tool_call_ids:
6365
logger.info("Not injecting synthetic tool result for non-abandoned call_id=%s", pending_call_id)
6466
continue
@@ -76,26 +78,29 @@ def _append_synthetic_tool_results(
7678
)
7779

7880

79-
def _ordered_unique_tool_call_ids(contents: list[Content]) -> list[str]:
80-
tool_ids: list[str] = []
81-
seen: set[str] = set()
82-
for content in contents:
83-
if content.type != "function_call" or not content.call_id:
84-
continue
85-
tool_id = str(content.call_id)
86-
if tool_id in seen:
87-
continue
88-
tool_ids.append(tool_id)
89-
seen.add(tool_id)
90-
return tool_ids
81+
def _tool_call_ids_per_occurrence(contents: list[Content]) -> list[str]:
82+
"""Return one call ID per function call, so a call ID reused by several occurrences awaits each result."""
83+
return [str(content.call_id) for content in contents if content.type == "function_call" and content.call_id]
84+
9185

86+
def _answered_call_ids_from(messages: list[Message], start: int) -> set[str]:
87+
"""Return call IDs that get a function result from ``start`` on, before the call ID is declared again.
9288
93-
def _function_result_call_ids(messages: list[Message]) -> set[str]:
89+
A result from an earlier turn belongs to that turn's occurrence, and a result after the call ID
90+
is declared again belongs to the later occurrence, so neither answers the pending occurrence.
91+
"""
9492
result_ids: set[str] = set()
95-
for msg in messages:
93+
redeclared_ids: set[str] = set()
94+
for msg in messages[start:]:
95+
is_assistant = get_role_value(msg) == "assistant"
9696
for content in msg.contents or []:
97-
if content.type == "function_result" and content.call_id:
98-
result_ids.add(str(content.call_id))
97+
if not content.call_id:
98+
continue
99+
call_id = str(content.call_id)
100+
if content.type == "function_call" and is_assistant:
101+
redeclared_ids.add(call_id)
102+
elif content.type == "function_result" and call_id not in redeclared_ids:
103+
result_ids.add(call_id)
99104
return result_ids
100105

101106

@@ -108,9 +113,13 @@ def _sanitize_tool_history(
108113
sanitized: list[Message] = []
109114
pending_tool_call_ids: list[str] | None = None
110115
pending_confirm_changes_id: str | None = None
111-
non_abandoned_tool_call_ids = set(protected_tool_call_ids or set()) | _function_result_call_ids(messages)
116+
protected_ids = set(protected_tool_call_ids or set())
112117

113-
for msg in messages:
118+
def non_abandoned_tool_call_ids(start: int) -> set[str]:
119+
"""Protected calls and calls answered from ``start`` on; earlier turns do not count."""
120+
return protected_ids | _answered_call_ids_from(messages, start)
121+
122+
for idx, msg in enumerate(messages):
114123
role_value = get_role_value(msg)
115124

116125
if role_value == "assistant":
@@ -123,12 +132,12 @@ def _sanitize_tool_history(
123132
sanitized,
124133
pending_tool_call_ids,
125134
"Tool execution skipped - assistant continued before the tool result was available.",
126-
excluded_tool_call_ids=non_abandoned_tool_call_ids,
135+
excluded_tool_call_ids=non_abandoned_tool_call_ids(idx),
127136
)
128137
pending_tool_call_ids = None
129138
pending_confirm_changes_id = None
130139

131-
tool_ids = _ordered_unique_tool_call_ids(msg.contents or [])
140+
tool_ids = _tool_call_ids_per_occurrence(msg.contents or [])
132141
confirm_changes_call = None
133142
for content in msg.contents or []:
134143
if content.type == "function_call" and content.name == "confirm_changes":
@@ -250,7 +259,7 @@ def _sanitize_tool_history(
250259
sanitized,
251260
pending_tool_call_ids,
252261
"Tool execution skipped - user provided follow-up message",
253-
excluded_tool_call_ids=non_abandoned_tool_call_ids,
262+
excluded_tool_call_ids=non_abandoned_tool_call_ids(idx + 1),
254263
)
255264
pending_tool_call_ids = None
256265
pending_confirm_changes_id = None
@@ -268,12 +277,11 @@ def _sanitize_tool_history(
268277
call_id = str(content.call_id)
269278
if call_id in pending_tool_call_ids:
270279
keep = True
271-
# Remove the call_id from pending since we now have its result.
280+
# Remove one pending occurrence of the call_id since we now have its result.
272281
# This prevents duplicate synthetic "skipped" results from being
273-
# injected when a user message arrives later.
274-
pending_tool_call_ids = [
275-
pending_id for pending_id in pending_tool_call_ids if pending_id != call_id
276-
]
282+
# injected when a user message arrives later, while later occurrences
283+
# that reuse the call_id still wait for their own results.
284+
pending_tool_call_ids.remove(call_id)
277285
if call_id == pending_confirm_changes_id:
278286
pending_confirm_changes_id = None
279287
break
@@ -291,7 +299,7 @@ def _sanitize_tool_history(
291299
sanitized,
292300
pending_tool_call_ids,
293301
"Tool execution skipped - conversation continued before the tool result was available.",
294-
excluded_tool_call_ids=non_abandoned_tool_call_ids,
302+
excluded_tool_call_ids=non_abandoned_tool_call_ids(idx),
295303
)
296304

297305
sanitized.append(msg)
@@ -307,7 +315,7 @@ def _sanitize_tool_history(
307315
sanitized,
308316
pending_tool_call_ids,
309317
"Tool execution skipped - conversation ended before the tool result was available.",
310-
excluded_tool_call_ids=non_abandoned_tool_call_ids,
318+
excluded_tool_call_ids=non_abandoned_tool_call_ids(len(messages)),
311319
)
312320

313321
return sanitized
@@ -317,45 +325,57 @@ def _deduplicate_messages(messages: list[Message]) -> list[Message]:
317325
"""Remove duplicate messages while preserving order."""
318326
seen_keys: dict[Any, int] = {}
319327
unique_messages: list[Message] = []
328+
# A provider call ID can be reused by several function calls, in one assistant message or in a
329+
# later turn. Each declared occurrence opens one result slot, so every occurrence keeps its own
330+
# result while replayed results beyond the open slots are dropped.
331+
open_result_slots: Counter[str] = Counter()
332+
declared_call_ids: set[str] = set()
333+
tool_result_indexes: dict[str, list[int]] = {}
320334

321335
for idx, msg in enumerate(messages):
322336
role_value = get_role_value(msg)
323337

324338
if role_value == "tool" and msg.contents and msg.contents[0].type == "function_result":
325339
call_id = str(msg.contents[0].call_id)
326-
key: Any = (role_value, call_id)
327-
328-
if key in seen_keys:
329-
existing_idx = seen_keys[key]
330-
existing_msg = unique_messages[existing_idx]
340+
kept_indexes = tool_result_indexes.setdefault(call_id, [])
331341

332-
existing_result = None
333-
if existing_msg.contents and existing_msg.contents[0].type == "function_result":
334-
existing_result = existing_msg.contents[0].result
342+
if open_result_slots[call_id] > 0:
343+
open_result_slots[call_id] -= 1
344+
elif call_id in declared_call_ids or kept_indexes:
335345
new_result = msg.contents[0].result
336-
337-
if (not existing_result or existing_result == "") and new_result:
338-
logger.info(f"Replacing empty tool result at index {existing_idx} with data from index {idx}")
339-
unique_messages[existing_idx] = msg
346+
for existing_idx in kept_indexes:
347+
existing_msg = unique_messages[existing_idx]
348+
existing_result = None
349+
if existing_msg.contents and existing_msg.contents[0].type == "function_result":
350+
existing_result = existing_msg.contents[0].result
351+
352+
if (not existing_result or existing_result == "") and new_result:
353+
logger.info(f"Replacing empty tool result at index {existing_idx} with data from index {idx}")
354+
unique_messages[existing_idx] = msg
355+
break
340356
else:
341357
logger.info(f"Skipping duplicate tool result at index {idx}: call_id={call_id}")
342358
continue
343359

344-
seen_keys[key] = len(unique_messages)
360+
kept_indexes.append(len(unique_messages))
345361
unique_messages.append(msg)
346362

347363
elif role_value == "assistant" and msg.contents and any(c.type == "function_call" for c in msg.contents):
348364
tool_call_ids = tuple(
349365
sorted(str(c.call_id) for c in msg.contents if c.type == "function_call" and c.call_id)
350366
)
351-
key = (role_value, tool_call_ids)
367+
# As for other messages, a shared message_id marks a replay of the same message, while a
368+
# later turn reusing the same call IDs arrives as a distinct message.
369+
key: Any = ("id", msg.message_id) if msg.message_id else (role_value, tool_call_ids)
352370

353371
if key in seen_keys:
354372
logger.info(f"Skipping duplicate assistant tool call at index {idx}")
355373
continue
356374

357375
seen_keys[key] = len(unique_messages)
358376
unique_messages.append(msg)
377+
open_result_slots.update(tool_call_ids)
378+
declared_call_ids.update(tool_call_ids)
359379

360380
else:
361381
# Use message_id for deduplication when available — two messages with the

0 commit comments

Comments
 (0)