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 ab44b2249da..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) @@ -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 ] 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..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 @@ -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,38 +321,48 @@ 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}") @@ -352,6 +370,8 @@ def _deduplicate_messages(messages: list[Message]) -> list[Message]: 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 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 1a29960c1c6..a69ed8d3f12 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_endpoint.py +++ b/python/packages/ag-ui/tests/ag_ui/test_endpoint.py @@ -4709,6 +4709,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] = [] @@ -4760,6 +4761,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 @@ -6175,10 +6178,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"} @@ -6244,8 +6249,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( @@ -6306,6 +6324,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, *, @@ -6862,6 +6926,278 @@ 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"] + + +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 (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"), + ] + + +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 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..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 @@ -2360,6 +2386,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 +2515,63 @@ 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_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 2e6b8ed56ce..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"] + )[("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,8 +646,8 @@ 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) == [] assert flow.tool_results[-1]["content"] == "" @@ -683,8 +683,8 @@ 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] assert flow.tool_results[-1]["content"] == "Summary"