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 436607114a9..5353f49afed 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 @@ -8,7 +8,7 @@ import json import logging import uuid -from collections.abc import AsyncIterable, Awaitable, Collection, Mapping, Sequence +from collections.abc import AsyncIterable, Awaitable, Callable, Collection, Mapping, Sequence from dataclasses import dataclass, field from functools import partial from typing import TYPE_CHECKING, Any, cast @@ -61,7 +61,7 @@ _try_execute_function_call_groups, # type: ignore normalize_function_invocation_configuration, ) -from agent_framework._types import ResponseStream +from agent_framework._types import APPROVAL_REJECTION_RESULT_KEY, ResponseStream from agent_framework.exceptions import AgentInvalidResponseException from agent_framework.observability import ( _use_telemetry_conversation_id, # pyright: ignore[reportPrivateUsage] @@ -90,7 +90,13 @@ InRunPendingToolTransitionOwner, ) from ._approval_state import _APPROVAL_SCOPE_INPUT_KEY, InMemoryAGUIApprovalStateStore, approval_state_thread_id -from ._message_adapters import _APPROVAL_DECISION_IS_BOOLEAN_KEY, normalize_agui_input_messages +from ._message_adapters import ( + _APPROVAL_DECISION_IS_BOOLEAN_KEY, + _copy_backend_snapshot_origin, + _mark_backend_snapshot_origin, + _take_backend_snapshot_messages, + normalize_agui_input_messages, +) from ._predictive_state import PredictiveStateHandler from ._tooling import collect_server_tools, merge_tools from ._run_common import ( @@ -1764,6 +1770,7 @@ async def _resolve_approval_responses( pending_response_groups: dict[object, tuple[ApprovalOccurrence, list[Content]]] = {} intents_by_response_content_id: dict[int, AuthorizedExecution] = {} authenticated_non_grants: list[Content] = [] + authenticated_rejection_occurrence_ids: set[str] = set() for response in approval_responses: resp_id = response.id function_call_id = response.function_call.call_id if response.function_call else None @@ -1861,6 +1868,7 @@ async def _resolve_approval_responses( ) if decision_is_boolean and primary_response.approved is False and not server_label: authenticated_non_grants.append(_approval_observer_response(pending_entry)) + authenticated_rejection_occurrence_ids.add(pending_entry.function_call_id) if ( (primary_response.approved is True or pending_entry.owner is ApprovalExecutionOwner.DEFERRED) and lifecycle is not None @@ -1902,6 +1910,14 @@ async def _resolve_approval_responses( authenticated_non_grants, session=invocation_session ) + if valid_response_content_ids: + for message in messages: + if message.contents and all( + content.type == "function_approval_response" and id(content) in valid_response_content_ids + for content in message.contents + ): + _mark_backend_snapshot_origin(message) + if response_content_ids_to_strip: filtered_messages: list[Message] = [] for message in messages: @@ -2069,18 +2085,32 @@ async def execute_local_call(approval: Content = approval, call_id: str = call_i if id(response) in deferred_response_content_ids } ) - _replace_approval_contents_with_results(messages, normalization_responses, replacement_groups) + resolved_contents = _replace_approval_contents_with_results(messages, normalization_responses, replacement_groups) + authenticated_rejection_result_ids = { + id(content) + for content in resolved_contents + if content.type == "function_result" + and content.id in authenticated_rejection_occurrence_ids + and content.additional_properties.get(APPROVAL_REJECTION_RESULT_KEY) is True + } # Post-process: Convert user messages with function_result content to proper tool messages. # After _replace_approval_contents_with_results, approved tool calls have their results # placed in user messages. OpenAI requires tool results to be in role="tool" messages. # This transformation ensures the message history is valid for the LLM provider. - _convert_approval_results_to_tool_messages(messages) + _convert_approval_results_to_tool_messages( + messages, + backend_result_content_ids=authenticated_rejection_result_ids, + ) return approved_results -def _convert_approval_results_to_tool_messages(messages: list[Message]) -> None: +def _convert_approval_results_to_tool_messages( + messages: list[Message], + *, + backend_result_content_ids: set[int] | None = None, +) -> None: """Convert function_result content in user messages to proper tool messages. After approval processing, tool results end up in user messages. OpenAI and other @@ -2091,11 +2121,25 @@ def _convert_approval_results_to_tool_messages(messages: list[Message]) -> None: Args: messages: List of Message objects to process + backend_result_content_ids: Object identities of exact backend-synthesized + results whose derived tool messages retain snapshot provenance. """ result: list[Message] = [] + backend_result_ids = backend_result_content_ids or set() for msg in messages: if get_role_value(msg) != "user": + backend_contents = [content for content in msg.contents if id(content) in backend_result_ids] + if backend_contents and len(backend_contents) == len(msg.contents): + _mark_backend_snapshot_origin(msg) + elif backend_contents and get_role_value(msg) == "tool": + for content in msg.contents: + split_message = Message(role="tool", contents=[content]) + _copy_backend_snapshot_origin(msg, split_message) + if id(content) in backend_result_ids: + _mark_backend_snapshot_origin(split_message) + result.append(split_message) + continue result.append(msg) continue @@ -2113,11 +2157,17 @@ def _convert_approval_results_to_tool_messages(messages: list[Message]) -> None: # Tool messages first (right after the preceding assistant message per OpenAI requirements) for func_result in function_results: - result.append(Message(role="tool", contents=[func_result])) + tool_message = Message(role="tool", contents=[func_result]) + _copy_backend_snapshot_origin(msg, tool_message) + if id(func_result) in backend_result_ids: + _mark_backend_snapshot_origin(tool_message) + result.append(tool_message) # Then user message with remaining content (if any) if other_contents: - result.append(Message(role="user", contents=other_contents)) + user_message = Message(role="user", contents=other_contents) + _copy_backend_snapshot_origin(msg, user_message) + result.append(user_message) messages[:] = result @@ -2801,6 +2851,47 @@ def _split_service_session_input( return snapshot_messages[len(stored_snapshot_messages) :], snapshot_messages +def _mark_fides_state_covered_messages( + agent: SupportsAgentRun, + messages: Sequence[Message], +) -> None: + """Mark backend-restored replay messages for FIDES without importing the experimental module.""" + if not messages: + return + for provider in getattr(agent, "context_providers", ()): + marker = getattr(provider, "_mark_replay_state_covered_messages", None) + if callable(marker): + marker(messages) + + +def _collect_fides_replay_folders(agent: SupportsAgentRun) -> list[Callable[..., None]]: + """Return replay-folding hooks exposed by FIDES context providers.""" + folders: list[Callable[..., None]] = [] + for provider in getattr(agent, "context_providers", ()): + fold = getattr(provider, "_fold_uncovered_replay_messages", None) + if callable(fold): + folders.append(cast(Callable[..., None], fold)) + return folders + + +def _fold_fides_uncovered_replay_messages( + folders: Sequence[Callable[..., None]], + session: AgentSession, + messages: Sequence[Message], + *, + complete_replay_messages: Sequence[Message] | None = None, + replay_history_complete: bool, +) -> None: + """Fold uncovered replay into FIDES state before a snapshot-only return.""" + for fold in folders: + fold( + session, + messages, + complete_replay_messages=complete_replay_messages, + replay_history_complete=replay_history_complete, + ) + + def _legacy_tool_message_approval_resume( messages: Sequence[Mapping[str, Any]], *, @@ -3004,6 +3095,7 @@ async def _run_agent_stream( ) stored_snapshot = snapshot_session.stored + backend_snapshot_message_ids: set[int] = set() stored_pending_approval_interrupt_ids: set[str] = set() seeded_resume_from_snapshot = False if stored_snapshot is not None: @@ -3073,6 +3165,8 @@ async def _run_agent_stream( stored_interrupt=stored_snapshot.interrupt, ) raw_messages = provider_suffix + if not config.use_service_session: + backend_snapshot_message_ids = {id(message) for message in raw_messages[: len(stored_snapshot.messages)]} # Initialize flow state with stored state plus request-provided overrides; # endpoint-deferred defaults apply only to keys missing from both. @@ -3198,7 +3292,9 @@ async def _run_agent_stream( raw_messages, sanitize_tool_history=not (config.use_service_session and stored_snapshot is not None), protected_tool_call_ids=protected_tool_call_ids, + backend_snapshot_message_ids=backend_snapshot_message_ids, ) + _mark_fides_state_covered_messages(agent, _take_backend_snapshot_messages(messages)) if snapshot_seed_messages is not None: _, snapshot_messages = normalize_agui_input_messages( @@ -3361,6 +3457,29 @@ async def save_thread_snapshot( ), ) + fides_replay_folders = _collect_fides_replay_folders(agent) + + def fold_fides_replay_before_snapshot_save() -> None: + if not fides_replay_folders: + return + complete_replay_messages: Sequence[Message] | None = None + if not config.use_service_session: + complete_replay_snapshot = snapshot_messages + if resume_payload is not None and not seeded_resume_from_snapshot and snapshot_seed_messages is None: + complete_replay_snapshot = snapshot_session.resume_seeded_messages(complete_replay_snapshot) + complete_replay_messages, _ = normalize_agui_input_messages( + complete_replay_snapshot, + protected_tool_call_ids=protected_tool_call_ids, + ) + _fold_fides_uncovered_replay_messages( + fides_replay_folders, + session, + messages, + complete_replay_messages=complete_replay_messages, + replay_history_complete=not config.use_service_session, + ) + _mark_fides_state_covered_messages(agent, messages) + authenticated_cancellations = [ _approval_observer_response(occurrence, cancelled=True) for interrupt_id in cancelled_resume_ids @@ -3552,6 +3671,7 @@ async def save_thread_snapshot( forwarded_executions=forwarded_executions, ) resolved_approval_results = retained_approval_results + newly_resolved_approval_results + _mark_fides_state_covered_messages(agent, _take_backend_snapshot_messages(messages)) # Defense-in-depth: replace approval payloads in snapshot with actual tool results # so CopilotKit does not re-send stale approval content on subsequent turns. @@ -3613,6 +3733,7 @@ def merge_snapshot_interrupts( persisted_messages = snapshot_messages if resume_payload is not None and not seeded_resume_from_snapshot and snapshot_seed_messages is None: persisted_messages = snapshot_session.resume_seeded_messages(persisted_messages) + fold_fides_replay_before_snapshot_save() await save_thread_snapshot( persisted_messages=persisted_messages, state=cast(dict[str, Any], make_json_safe(flow.current_state)) if flow.current_state else None, @@ -3646,6 +3767,7 @@ def merge_snapshot_interrupts( if resume_payload is not None and not seeded_resume_from_snapshot and snapshot_seed_messages is None: persisted_messages = snapshot_session.resume_seeded_messages(persisted_messages) preserved_interrupts = remaining_stored_interrupts(retired_approval_interrupt_ids) + fold_fides_replay_before_snapshot_save() await save_thread_snapshot( persisted_messages=persisted_messages, state=cast(dict[str, Any], make_json_safe(flow.current_state)) if flow.current_state else None, @@ -3654,6 +3776,7 @@ def merge_snapshot_interrupts( # Handle confirm_changes response (state confirmation flow - emit confirmation and stop) if is_confirm_changes_response: + fold_fides_replay_before_snapshot_save() confirm_additional_properties = cast(dict[str, Any], messages[-1].additional_properties or {}) confirm_interrupt_id = confirm_additional_properties.get("tool_call_id") confirm_remaining_interrupts = remaining_stored_interrupts( 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 41adb4a7d2d..7bc8c4f1f4a 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 @@ -10,6 +10,7 @@ import json import logging from collections import Counter +from collections.abc import Sequence from typing import Any, cast, get_args from urllib.parse import urlsplit @@ -49,6 +50,35 @@ _ALLOWED_MULTIMODAL_URI_SCHEMES = frozenset({"data", "http", "https"}) _APPROVAL_DECISION_IS_BOOLEAN_KEY = "_ag_ui_approval_decision_is_boolean" _LEGACY_MEDIA_PLACEHOLDER = "[Unsupported legacy media omitted by AG-UI URI policy.]" +_BACKEND_SNAPSHOT_ORIGIN_KEY = "_agent_framework_ag_ui_backend_snapshot_origin" +_BACKEND_SNAPSHOT_ORIGIN_MARKER = object() + + +def _mark_backend_snapshot_origin(message: Message) -> None: + """Mark a provider message as derived from backend-owned snapshot history.""" + message.additional_properties[_BACKEND_SNAPSHOT_ORIGIN_KEY] = _BACKEND_SNAPSHOT_ORIGIN_MARKER + + +def _clear_backend_snapshot_origin(message: Message) -> None: + """Invalidate backend-snapshot provenance after client-controlled mutation.""" + message.additional_properties.pop(_BACKEND_SNAPSHOT_ORIGIN_KEY, None) + + +def _copy_backend_snapshot_origin(source: Message, target: Message) -> None: + """Preserve private backend-snapshot provenance on a derived provider message.""" + if source.additional_properties.get(_BACKEND_SNAPSHOT_ORIGIN_KEY) is _BACKEND_SNAPSHOT_ORIGIN_MARKER: + _mark_backend_snapshot_origin(target) + + +def _take_backend_snapshot_messages(messages: Sequence[Message]) -> list[Message]: + """Return normalized backend-snapshot messages and clear their private transport marker.""" + snapshot_messages: list[Message] = [] + for message in messages: + if message.additional_properties.get(_BACKEND_SNAPSHOT_ORIGIN_KEY) is not _BACKEND_SNAPSHOT_ORIGIN_MARKER: + continue + message.additional_properties.pop(_BACKEND_SNAPSHOT_ORIGIN_KEY, None) + snapshot_messages.append(message) + return snapshot_messages def _append_synthetic_tool_results( @@ -57,6 +87,7 @@ def _append_synthetic_tool_results( result: str, *, excluded_tool_call_ids: set[str] | None = None, + backend_snapshot_origin: bool = False, ) -> None: excluded_tool_call_ids = excluded_tool_call_ids or set() # A call ID awaiting several occurrences still gets a single synthetic result. @@ -65,17 +96,18 @@ def _append_synthetic_tool_results( logger.info("Not injecting synthetic tool result for non-abandoned call_id=%s", pending_call_id) continue logger.info("Injecting synthetic tool result for pending call_id=%s", pending_call_id) - sanitized.append( - Message( - role="tool", - contents=[ - Content.from_function_result( - call_id=pending_call_id, - result=result, - ) - ], - ) + synthetic_result = Message( + role="tool", + contents=[ + Content.from_function_result( + call_id=pending_call_id, + result=result, + ) + ], ) + if backend_snapshot_origin: + _mark_backend_snapshot_origin(synthetic_result) + sanitized.append(synthetic_result) def _tool_call_ids_per_occurrence(contents: list[Content]) -> list[str]: @@ -113,6 +145,7 @@ def _sanitize_tool_history( sanitized: list[Message] = [] pending_tool_call_ids: list[str] | None = None pending_confirm_changes_id: str | None = None + pending_tool_calls_from_backend_snapshot = False protected_ids = set(protected_tool_call_ids or set()) def non_abandoned_tool_call_ids(start: int) -> set[str]: @@ -133,9 +166,11 @@ def non_abandoned_tool_call_ids(start: int) -> set[str]: pending_tool_call_ids, "Tool execution skipped - assistant continued before the tool result was available.", excluded_tool_call_ids=non_abandoned_tool_call_ids(idx), + backend_snapshot_origin=pending_tool_calls_from_backend_snapshot, ) pending_tool_call_ids = None pending_confirm_changes_id = None + pending_tool_calls_from_backend_snapshot = False tool_ids = _tool_call_ids_per_occurrence(msg.contents or []) confirm_changes_call = None @@ -159,6 +194,7 @@ def non_abandoned_tool_call_ids(start: int) -> set[str]: if filtered_contents: # Create a new message without confirm_changes to avoid mutating the input filtered_msg = Message(role=msg.role, contents=filtered_contents) + _copy_backend_snapshot_origin(msg, filtered_msg) sanitized.append(filtered_msg) # If no contents left after filtering, don't append anything @@ -172,6 +208,9 @@ def non_abandoned_tool_call_ids(start: int) -> set[str]: sanitized.append(msg) pending_tool_call_ids = tool_ids if tool_ids else None + pending_tool_calls_from_backend_snapshot = bool(tool_ids) and ( + msg.additional_properties.get(_BACKEND_SNAPSHOT_ORIGIN_KEY) is _BACKEND_SNAPSHOT_ORIGIN_MARKER + ) pending_confirm_changes_id = ( str(confirm_changes_call.call_id) if confirm_changes_call and confirm_changes_call.call_id else None ) @@ -209,6 +248,8 @@ def non_abandoned_tool_call_ids(start: int) -> set[str]: ) ], ) + if pending_tool_calls_from_backend_snapshot: + _mark_backend_snapshot_origin(synthetic_result) sanitized.append(synthetic_result) if pending_tool_call_ids: pending_tool_call_ids = [ @@ -240,6 +281,8 @@ def non_abandoned_tool_call_ids(start: int) -> set[str]: ) ], ) + if pending_tool_calls_from_backend_snapshot: + _mark_backend_snapshot_origin(synthetic_result) sanitized.append(synthetic_result) if pending_tool_call_ids: pending_tool_call_ids = [ @@ -260,12 +303,16 @@ def non_abandoned_tool_call_ids(start: int) -> set[str]: pending_tool_call_ids, "Tool execution skipped - user provided follow-up message", excluded_tool_call_ids=non_abandoned_tool_call_ids(idx + 1), + backend_snapshot_origin=pending_tool_calls_from_backend_snapshot, ) pending_tool_call_ids = None pending_confirm_changes_id = None + pending_tool_calls_from_backend_snapshot = False sanitized.append(msg) pending_confirm_changes_id = None + if not pending_tool_call_ids: + pending_tool_calls_from_backend_snapshot = False continue if role_value == "tool": @@ -300,11 +347,13 @@ def non_abandoned_tool_call_ids(start: int) -> set[str]: pending_tool_call_ids, "Tool execution skipped - conversation continued before the tool result was available.", excluded_tool_call_ids=non_abandoned_tool_call_ids(idx), + backend_snapshot_origin=pending_tool_calls_from_backend_snapshot, ) sanitized.append(msg) pending_tool_call_ids = None pending_confirm_changes_id = None + pending_tool_calls_from_backend_snapshot = False if pending_tool_call_ids: logger.info( @@ -316,6 +365,7 @@ def non_abandoned_tool_call_ids(start: int) -> set[str]: pending_tool_call_ids, "Tool execution skipped - conversation ended before the tool result was available.", excluded_tool_call_ids=non_abandoned_tool_call_ids(len(messages)), + backend_snapshot_origin=pending_tool_calls_from_backend_snapshot, ) return sanitized @@ -661,6 +711,7 @@ def normalize_agui_input_messages( *, sanitize_tool_history: bool = True, protected_tool_call_ids: set[str] | None = None, + backend_snapshot_message_ids: set[int] | None = None, ) -> tuple[list[Message], list[dict[str, Any]]]: """Normalize raw AG-UI messages into provider and snapshot formats. @@ -671,8 +722,13 @@ def normalize_agui_input_messages( where pending-request responses must come explicitly from interrupt resume. protected_tool_call_ids: Server-owned tool calls that are still eligible to complete and must not receive synthetic skipped results. + backend_snapshot_message_ids: Object identities of raw messages restored + from the backend-owned Thread Snapshot. """ - provider_messages = agui_messages_to_agent_framework(messages) + provider_messages = agui_messages_to_agent_framework( + messages, + backend_snapshot_message_ids=backend_snapshot_message_ids, + ) if sanitize_tool_history: provider_messages = _sanitize_tool_history(provider_messages, protected_tool_call_ids=protected_tool_call_ids) provider_messages = _deduplicate_messages(provider_messages) @@ -698,11 +754,17 @@ def _deserialize_model_replay_items(serialized_items: Any) -> list[Content] | No return None -def agui_messages_to_agent_framework(messages: list[dict[str, Any]]) -> list[Message]: +def agui_messages_to_agent_framework( + messages: list[dict[str, Any]], + *, + backend_snapshot_message_ids: set[int] | None = None, +) -> list[Message]: """Convert AG-UI messages to Agent Framework format. Args: - messages: List of AG-UI messages + messages: List of AG-UI messages. + backend_snapshot_message_ids: Object identities of messages restored + from the backend-owned Thread Snapshot. Returns: List of Agent Framework Message objects @@ -732,14 +794,14 @@ def _update_tool_call_arguments( function_payload["arguments"] = modified_args return - def _find_matching_func_call(call_id: str) -> Content | None: + def _find_matching_func_call(call_id: str) -> tuple[Message, Content] | None: for prev_msg in result: role_val = prev_msg.role if hasattr(prev_msg.role, "value") else str(prev_msg.role) if role_val != "assistant": continue for content in prev_msg.contents or []: if content.type == "function_call" and content.call_id == call_id and content.name != "confirm_changes": - return content + return prev_msg, content return None def _parse_arguments(arguments: Any) -> dict[str, Any] | None: @@ -801,6 +863,12 @@ def _filter_modified_args( return {key: value for key, value in modified_args.items() if key in allowed_keys} result: list[Message] = [] + + def _append_message(raw_message: dict[str, Any], message: Message) -> None: + if backend_snapshot_message_ids is not None and id(raw_message) in backend_snapshot_message_ids: + message.additional_properties[_BACKEND_SNAPSHOT_ORIGIN_KEY] = _BACKEND_SNAPSHOT_ORIGIN_MARKER + result.append(message) + for msg in messages: # Handle standard tool result messages early (role="tool") to preserve provider invariants # This path maps AG‑UI tool messages to function_result content with the correct tool_call_id @@ -850,7 +918,7 @@ def _filter_modified_args( ) if "id" in msg: chat_msg.message_id = msg["id"] - result.append(chat_msg) + _append_message(msg, chat_msg) continue if msg.get(_AGUI_HOST_PAYLOAD_OMITTED_KEY) is True: @@ -861,7 +929,7 @@ def _filter_modified_args( ) if "id" in msg: chat_msg.message_id = msg["id"] - result.append(chat_msg) + _append_message(msg, chat_msg) continue is_approval = parsed is not None and "accepted" in parsed @@ -883,9 +951,10 @@ def _filter_modified_args( resolved_call_id = _resolve_approval_call_id(tool_call_id, parsed) if resolved_call_id: approval_call_id = resolved_call_id - matching_func_call = _find_matching_func_call(approval_call_id) + matching_call = _find_matching_func_call(approval_call_id) - if matching_func_call: + if matching_call: + matching_message, matching_func_call = matching_call # Remove any existing tool result for this call_id since the framework # will re-execute the tool after approval. Keeping old results causes # OpenAI API errors ("tool message must follow assistant with tool_calls"). @@ -908,6 +977,7 @@ def _filter_modified_args( filtered_args = _filter_modified_args(modified_args, original_args) state_args: dict[str, Any] | None = None if filtered_args: + _clear_backend_snapshot_origin(matching_message) original_args = original_args or {} merged_args: dict[str, Any] if isinstance(original_args, dict) and original_args: @@ -993,7 +1063,7 @@ def _filter_modified_args( ) if "id" in msg: chat_msg.message_id = msg["id"] - result.append(chat_msg) + _append_message(msg, chat_msg) continue # Cast result_content to acceptable type for function_result content @@ -1012,7 +1082,7 @@ def _filter_modified_args( ) if "id" in msg: chat_msg.message_id = msg["id"] - result.append(chat_msg) + _append_message(msg, chat_msg) continue # Backend tool rendering payloads without an explicit role @@ -1028,7 +1098,7 @@ def _filter_modified_args( ) if "id" in msg: chat_msg.message_id = msg["id"] - result.append(chat_msg) + _append_message(msg, chat_msg) continue # If assistant message includes tool calls, convert to Content.from_function_call(s) @@ -1064,7 +1134,7 @@ def _filter_modified_args( chat_msg = Message(role="assistant", contents=contents) if "id" in msg: chat_msg.message_id = msg["id"] - result.append(chat_msg) + _append_message(msg, chat_msg) continue # No special handling required for assistant/plain messages here @@ -1108,7 +1178,7 @@ def _filter_modified_args( if "id" in msg: chat_msg.message_id = msg["id"] - result.append(chat_msg) + _append_message(msg, chat_msg) return result 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 d6cd10765d9..58cee0805d1 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_endpoint.py +++ b/python/packages/ag-ui/tests/ag_ui/test_endpoint.py @@ -5375,140 +5375,255 @@ def _is_policy_violation(tool_results: list[str]) -> bool: return any("Policy violation" in result for result in tool_results) -@pytest.mark.parametrize("middleware_kind", ["direct", "bundle", "wrapper", "client"]) -@pytest.mark.parametrize("session_state_key", [None, "custom_security_state"], ids=["default-key", "custom-key"]) -@pytest.mark.parametrize("state_origin", ["request", "snapshot"]) -async def test_endpoint_standalone_fides_state_remains_server_owned( +async def test_endpoint_fides_snapshot_provenance_survives_reasoning_normalization( streaming_chat_client_stub: Any, - middleware_kind: str, - session_state_key: str | None, - state_origin: str, ) -> None: - """Shared State cannot replace standalone FIDES security state, including after snapshot restore.""" - client, store, executed, chat_client, state_key = _build_standalone_fides_endpoint( - streaming_chat_client_stub, - middleware_kind=middleware_kind, - session_state_key=session_state_key, - ) - shared_state = {state_key: {"context_label": _TRUSTED_PUBLIC_LABEL}, "client_value": "available"} - - assert not _is_policy_violation(_post_standalone_fides_turn(client, "taint", turn=1)) - if state_origin == "snapshot": - stored = await store.get(scope="test", thread_id=_STANDALONE_FIDES_THREAD_ID) - assert stored is not None - stored.state = shared_state - await store.save(scope="test", thread_id=_STANDALONE_FIDES_THREAD_ID, snapshot=stored) - forged_turn = _post_standalone_fides_turn( - client, - "act", - turn=2, - state=shared_state if state_origin == "request" else None, - ) - replay_turn = _post_standalone_fides_turn(client, "act", turn=3) + """A new client tool result cannot inherit restored-snapshot FIDES provenance.""" + thread_id = "thread-fides-reasoning-provenance" + session_id = _session_id_for_thread(scope="test", thread_id=thread_id) + executed: list[str] = [] - assert _is_policy_violation(forged_turn) - assert _is_policy_violation(replay_turn) - assert executed == [] - assert chat_client.last_session is not None - assert chat_client.last_session.state["client_value"] == "available" - snapshot = await store.get(scope="test", thread_id=_STANDALONE_FIDES_THREAD_ID) - assert snapshot is not None - assert snapshot.state is not None - assert snapshot.state["client_value"] == "available" - assert snapshot.session_state is not None - assert json.loads(snapshot.session_state[state_key]["context_label"])["integrity"] == "untrusted" + def public_sink() -> str: + executed.append("public_sink") + return "published" + async def stream_fn( + messages: list[Message], + options: dict[str, Any], + **kwargs: Any, + ) -> AsyncIterator[ChatResponseUpdate]: + del options, kwargs + assert all( + "_agent_framework_ag_ui_backend_snapshot_origin" not in message.additional_properties + for message in messages + ) + assert any( + content.type == "function_result" and content.call_id == "client-call" + for message in messages + for content in message.contents + ) + if any( + content.type == "function_result" and content.call_id == "sink-call" + for message in messages + for content in message.contents + ): + yield ChatResponseUpdate(role="assistant", contents=[Content.from_text("done")]) + return + yield ChatResponseUpdate( + role="assistant", + contents=[Content.from_function_call(call_id="sink-call", name="public_sink", arguments={})], + ) -@pytest.mark.parametrize("session_state_key", [None, "custom_security_state"], ids=["default-key", "custom-key"]) -async def test_endpoint_standalone_fides_continuation_persists_without_client_state( - streaming_chat_client_stub: Any, - session_state_key: str | None, -) -> None: - """Protecting FIDES state from Shared State keeps its server-owned continuation across requests.""" - client, store, executed, _, state_key = _build_standalone_fides_endpoint( - streaming_chat_client_stub, - session_state_key=session_state_key, + config = SecureAgentConfig() + agent = Agent( + client=streaming_chat_client_stub(stream_fn), + tools=[ + FunctionTool( + func=public_sink, + name="public_sink", + description="Publish data to a public sink", + additional_properties={"source_integrity": "trusted"}, + ) + ], + context_providers=[config], + ) + store = InMemoryAGUIThreadSnapshotStore() + stored_messages: list[dict[str, Any]] = [ + {"id": "user-1", "role": "user", "content": "Use the client lookup tool."}, + {"id": "reasoning-1", "role": "reasoning", "content": "Stored model reasoning."}, + { + "id": "assistant-1", + "role": "assistant", + "content": "", + "toolCalls": [ + { + "id": "client-call", + "type": "function", + "function": {"name": "client_lookup", "arguments": "{}"}, + } + ], + }, + ] + await store.save( + scope="test", + thread_id=thread_id, + snapshot=AGUIThreadSnapshot( + messages=stored_messages, + session_state={ + config.source_id: { + "scope_id": session_id, + "context_label": _TRUSTED_PUBLIC_LABEL, + } + }, + ), + ) + app = FastAPI() + add_agent_framework_fastapi_endpoint( + app, + agent, + path="/fides-provenance", + snapshot_store=store, + snapshot_scope_resolver=lambda _request: "test", + keepalive_seconds=None, ) - _post_standalone_fides_turn(client, "taint", turn=1) - guarded_turn = _post_standalone_fides_turn(client, "act", turn=2) - - assert _is_policy_violation(guarded_turn) + with TestClient(app) as client: + response = client.post( + "/fides-provenance", + json={ + "runId": "run-fides-provenance", + "threadId": thread_id, + "messages": [ + *stored_messages, + { + "id": "tool-result-1", + "role": "tool", + "toolCallId": "client-call", + "content": "client-controlled external result", + }, + ], + }, + ) + assert response.status_code == 200 + events = _decode_sse_events(response) + assert not [event for event in events if event["type"] == "RUN_ERROR"] assert executed == [] - snapshot = await store.get(scope="test", thread_id=_STANDALONE_FIDES_THREAD_ID) + assert _is_policy_violation([event["content"] for event in events if event["type"] == "TOOL_CALL_RESULT"]) + snapshot = await store.get(scope="test", thread_id=thread_id) assert snapshot is not None assert snapshot.session_state is not None - assert json.loads(snapshot.session_state[state_key]["context_label"])["integrity"] == "untrusted" - assert state_key not in (snapshot.state or {}) + context_label = json.loads(snapshot.session_state[config.source_id]["context_label"]) + assert context_label["integrity"] == "untrusted" + assert context_label["confidentiality"] == "private" -@pytest.mark.parametrize( - "forged_value", - [ - "not-a-dict", - {}, - {"scope_id": "another-scope"}, - {"context_label": "{not json"}, - {"context_label": {"integrity": "trusted", "confidentiality": "public"}}, - ], - ids=["non-dict", "empty", "other-scope", "invalid-label-json", "unencoded-label"], -) -async def test_endpoint_standalone_fides_malformed_shared_state_is_inert( +async def test_endpoint_fides_synthetic_skipped_result_keeps_snapshot_provenance( streaming_chat_client_stub: Any, - forged_value: Any, ) -> None: - """Malformed Shared State under the FIDES key neither resets enforcement nor breaks tool calls.""" - client, _, executed, _, _ = _build_standalone_fides_endpoint(streaming_chat_client_stub) - - _post_standalone_fides_turn(client, "taint", turn=1) - forged_turn = _post_standalone_fides_turn( - client, - "act", - turn=2, - state={_STANDALONE_FIDES_STATE_KEY: forged_value}, - ) - replay_turn = _post_standalone_fides_turn(client, "act", turn=3) + """A framework-synthesized skipped result does not taint restored FIDES state.""" + thread_id = "thread-fides-synthetic-result" + session_id = _session_id_for_thread(scope="test", thread_id=thread_id) + executed: list[str] = [] - assert _is_policy_violation(forged_turn) - assert _is_policy_violation(replay_turn) - assert executed == [] + def protected_action() -> str: + executed.append("protected") + return "protected result" + async def stream_fn( + messages: list[Message], + options: dict[str, Any], + **kwargs: Any, + ) -> AsyncIterator[ChatResponseUpdate]: + del options, kwargs + assert all( + "_agent_framework_ag_ui_backend_snapshot_origin" not in message.additional_properties + for message in messages + ) + assert any( + content.type == "function_result" + and content.call_id == "client-call" + and "Tool execution skipped" in (content.result or "") + for message in messages + for content in message.contents + ) + if any( + content.type == "function_result" and content.call_id == "protected-call" + for message in messages + for content in message.contents + ): + yield ChatResponseUpdate(role="assistant", contents=[Content.from_text("done")]) + return + yield ChatResponseUpdate( + role="assistant", + contents=[Content.from_function_call(call_id="protected-call", name="protected_action", arguments={})], + ) -async def test_endpoint_standalone_fides_poisoned_snapshot_state_does_not_reset_enforcement( - streaming_chat_client_stub: Any, -) -> None: - """A snapshot whose Shared State already carries a FIDES scope no longer overrides server-owned state.""" - client, store, executed, _, _ = _build_standalone_fides_endpoint(streaming_chat_client_stub) + config = SecureAgentConfig() + agent = Agent( + client=streaming_chat_client_stub(stream_fn), + tools=[ + FunctionTool( + func=protected_action, + name="protected_action", + description="Run an unrelated protected action", + additional_properties={ + "source_integrity": "trusted", + "max_allowed_confidentiality": "public", + }, + ) + ], + context_providers=[config], + ) + store = InMemoryAGUIThreadSnapshotStore() await store.save( scope="test", - thread_id=_STANDALONE_FIDES_THREAD_ID, - snapshot=AGUIThreadSnapshot(state={_STANDALONE_FIDES_STATE_KEY: {"context_label": _TRUSTED_PUBLIC_LABEL}}), + thread_id=thread_id, + snapshot=AGUIThreadSnapshot( + messages=[ + {"id": "user-1", "role": "user", "content": "Use the client lookup tool."}, + { + "id": "assistant-1", + "role": "assistant", + "content": "", + "toolCalls": [ + { + "id": "client-call", + "type": "function", + "function": {"name": "client_lookup", "arguments": "{}"}, + } + ], + }, + ], + session_state={ + config.source_id: { + "scope_id": session_id, + "context_label": _TRUSTED_PUBLIC_LABEL, + } + }, + ), + ) + app = FastAPI() + add_agent_framework_fastapi_endpoint( + app, + agent, + path="/fides-synthetic-result", + snapshot_store=store, + snapshot_scope_resolver=lambda _request: "test", + keepalive_seconds=None, ) - _post_standalone_fides_turn(client, "taint", turn=1) - guarded_turn = _post_standalone_fides_turn(client, "act", turn=2) + with TestClient(app) as client: + response = client.post( + "/fides-synthetic-result", + json={ + "runId": "run-fides-synthetic-result", + "threadId": thread_id, + "messages": [{"id": "user-2", "role": "user", "content": "Continue without the client tool."}], + }, + ) - assert _is_policy_violation(guarded_turn) - assert executed == [] + assert response.status_code == 200 + events = _decode_sse_events(response) + assert not [event for event in events if event["type"] == "RUN_ERROR"] + assert executed == ["protected"] + snapshot = await store.get(scope="test", thread_id=thread_id) + assert snapshot is not None + assert snapshot.session_state is not None + context_label = json.loads(snapshot.session_state[config.source_id]["context_label"]) + assert context_label["integrity"] == "trusted" + assert context_label["confidentiality"] == "public" -def _build_tool_approval_auto_endpoint( +async def test_endpoint_fides_normal_continuation_reconstructs_missing_context_label( streaming_chat_client_stub: Any, -) -> tuple[TestClient, list[str], list[Message], dict[str, str]]: +) -> None: + """A normal continuation reconstructs backend history when FIDES state is absent.""" + thread_id = "thread-fides-normal-missing-state" executed: list[str] = [] - messages_received: list[Message] = [] - state = {"phase": "pause"} - - def auto_tool() -> str: - executed.append("auto") - return "auto result" - - def manual_tool() -> str: - executed.append("manual") - return "manual result" - def auto_approve_auto_tool(function_call: Content) -> bool: - return function_call.name == "auto_tool" + def protected_action() -> str: + executed.append("protected") + return "protected result" async def stream_fn( messages: list[Message], @@ -5516,2113 +5631,3209 @@ async def stream_fn( **kwargs: Any, ) -> AsyncIterator[ChatResponseUpdate]: del options, kwargs - if state["phase"] == "pause": - yield ChatResponseUpdate( - contents=[ - Content.from_function_call(call_id="call_auto", name="auto_tool", arguments="{}"), - Content.from_function_call(call_id="call_manual", name="manual_tool", arguments="{}"), - ], - role="assistant", - ) + assert any( + content.type == "function_result" + and content.call_id == "external-call" + and content.result == "stored external result" + for message in messages + for content in message.contents + ) + if any( + content.type == "function_result" and content.call_id == "protected-call" + for message in messages + for content in message.contents + ): + yield ChatResponseUpdate(role="assistant", contents=[Content.from_text("done")]) return - messages_received[:] = list(messages) - yield ChatResponseUpdate(contents=[Content.from_text(text="Done.")], role="assistant") + yield ChatResponseUpdate( + role="assistant", + contents=[Content.from_function_call(call_id="protected-call", name="protected_action", arguments={})], + ) + config = SecureAgentConfig() agent = Agent( - name="test_agent", - instructions="Test", client=streaming_chat_client_stub(stream_fn), tools=[ - FunctionTool(name="auto_tool", description="Auto tool", func=auto_tool, approval_mode="always_require"), FunctionTool( - name="manual_tool", description="Manual tool", func=manual_tool, approval_mode="always_require" - ), + func=protected_action, + name="protected_action", + description="Run an unrelated protected action", + additional_properties={ + "source_integrity": "trusted", + "max_allowed_confidentiality": "public", + }, + ) ], - middleware=[ToolApprovalMiddleware(auto_approval_rules=[auto_approve_auto_tool])], + context_providers=[config], + ) + store = InMemoryAGUIThreadSnapshotStore() + await store.save( + scope="test", + thread_id=thread_id, + snapshot=AGUIThreadSnapshot( + messages=[ + {"id": "user-1", "role": "user", "content": "Look up the external record."}, + { + "id": "assistant-external-call", + "role": "assistant", + "content": "", + "toolCalls": [ + { + "id": "external-call", + "type": "function", + "function": {"name": "client_lookup", "arguments": "{}"}, + } + ], + }, + { + "id": "external-result", + "role": "tool", + "toolCallId": "external-call", + "content": "stored external result", + }, + ], + ), ) app = FastAPI() add_agent_framework_fastapi_endpoint( app, - AgentFrameworkAgent(agent=agent, require_confirmation=False), - path="/approval", + agent, + path="/fides-normal-missing-state", + snapshot_store=store, + snapshot_scope_resolver=lambda _request: "test", + keepalive_seconds=None, ) - return TestClient(app), executed, messages_received, state - - -async def test_endpoint_agent_approval_resume_entry_executes_approved_tool(): - """A resolved canonical approval resume should execute the pending approved tool.""" - client, _, executed_cities = _build_weather_approval_endpoint() - response = client.post( - "/approval", - json={ - "runId": "run-resume", - "threadId": "thread-weather", - "messages": [], - "resume": [ - { - "interruptId": "call_get_weather", - "status": "resolved", - "payload": {"accepted": True}, - } - ], - }, - ) + with TestClient(app) as client: + response = client.post( + "/fides-normal-missing-state", + json={ + "runId": "run-fides-normal-missing-state", + "threadId": thread_id, + "messages": [{"id": "user-2", "role": "user", "content": "Run the protected action."}], + }, + ) assert response.status_code == 200 events = _decode_sse_events(response) - tool_results = [event for event in events if event.get("type") == "TOOL_CALL_RESULT"] - assert len(tool_results) == 1 - assert tool_results[0]["toolCallId"] == "call_get_weather" - assert tool_results[0]["content"] == "Sunny in Seattle" - assert executed_cities == ["Seattle"] - assert "outcome" not in [event for event in events if event.get("type") == "RUN_FINISHED"][-1] + assert not [event for event in events if event["type"] == "RUN_ERROR"] + assert executed == [] + assert _is_policy_violation([event["content"] for event in events if event["type"] == "TOOL_CALL_RESULT"]) + snapshot = await store.get(scope="test", thread_id=thread_id) + assert snapshot is not None + assert snapshot.session_state is not None + context_label = json.loads(snapshot.session_state[config.source_id]["context_label"]) + assert context_label["integrity"] == "untrusted" + assert context_label["confidentiality"] == "private" -async def test_endpoint_agent_legacy_tool_message_uses_unique_pending_call_id( - caplog: pytest.LogCaptureFixture, +async def test_endpoint_fides_confirm_changes_folds_fresh_client_result( + streaming_chat_client_stub: Any, ) -> None: - """Legacy tool messages remain usable only while their provider call id is unambiguous.""" + """A confirm-only return folds a fresh client result before snapshot persistence.""" + thread_id = "thread-fides-confirm-client-result" executed: list[str] = [] + model_calls = 0 - def get_weather(city: str) -> str: - executed.append(city) - return f"Sunny in {city}" + def protected_action() -> str: + executed.append("protected") + return "protected result" - weather_tool = FunctionTool( - name="get_weather", - description="Get the weather for a city", - func=get_weather, - approval_mode="always_require", - ) - function_call = Content.from_function_call( - call_id="provider-weather", - name="get_weather", - arguments={"city": "Seattle"}, - id="af-call-weather", - ) - approval_request = Content.from_function_approval_request( - id="af-call-weather", - function_call=function_call, + async def stream_fn( + messages: list[Message], + options: dict[str, Any], + **kwargs: Any, + ) -> AsyncIterator[ChatResponseUpdate]: + nonlocal model_calls + del options, kwargs + model_calls += 1 + assert any( + content.type == "function_result" and content.call_id == "client-call" + for message in messages + for content in message.contents + ) + if any( + content.type == "function_result" and content.call_id == "protected-call" + for message in messages + for content in message.contents + ): + yield ChatResponseUpdate(role="assistant", contents=[Content.from_text("done")]) + return + yield ChatResponseUpdate( + role="assistant", + contents=[Content.from_function_call(call_id="protected-call", name="protected_action", arguments={})], + ) + + config = SecureAgentConfig() + agent = Agent( + client=streaming_chat_client_stub(stream_fn), + tools=[ + FunctionTool( + func=protected_action, + name="protected_action", + description="Run an unrelated protected action", + additional_properties={ + "source_integrity": "trusted", + "max_allowed_confidentiality": "public", + }, + ) + ], + context_providers=[config], ) - agent = StubAgent( - updates=[AgentResponseUpdate(contents=[approval_request], role="assistant")], - default_options={"tools": [weather_tool]}, + store = InMemoryAGUIThreadSnapshotStore() + await store.save( + scope="test", + thread_id=thread_id, + snapshot=AGUIThreadSnapshot( + messages=[ + {"id": "user-1", "role": "user", "content": "Draft the plan"}, + { + "id": "assistant-1", + "role": "assistant", + "content": "", + "toolCalls": [ + { + "id": "client-call", + "type": "function", + "function": {"name": "client_lookup", "arguments": "{}"}, + }, + { + "id": "confirm-call", + "type": "function", + "function": {"name": "confirm_changes", "arguments": "{}"}, + }, + ], + }, + ], + session_state={ + config.source_id: { + "scope_id": _session_id_for_thread(scope="test", thread_id=thread_id), + "context_label": _TRUSTED_PUBLIC_LABEL, + } + }, + ), ) - wrapped_agent = AgentFrameworkAgent(agent=agent, require_confirmation=False) app = FastAPI() - add_agent_framework_fastapi_endpoint(app, wrapped_agent, path="/approval") - client = TestClient(app) - pause = client.post( - "/approval", - json={ - "runId": "run-pause", - "threadId": "thread-legacy-message", - "messages": [{"role": "user", "content": "Weather?"}], - }, + add_agent_framework_fastapi_endpoint( + app, + agent, + path="/fides-confirm-client-result", + state_schema={"steps": {"type": "array", "items": {"type": "object"}}}, + predict_state_config={"steps": {"tool": "draft_steps", "tool_argument": "steps"}}, + snapshot_store=store, + snapshot_scope_resolver=lambda _request: "test", + keepalive_seconds=None, ) - assert pause.status_code == 200 - agent.updates = [AgentResponseUpdate(contents=[Content.from_text(text="Done.")], role="assistant")] - with caplog.at_level(logging.WARNING, logger="agent_framework"): - response = client.post( - "/approval", + with TestClient(app) as client: + confirm_response = client.post( + "/fides-confirm-client-result", json={ - "runId": "run-resume", - "threadId": "thread-legacy-message", + "runId": "run-fides-confirm-resume", + "threadId": thread_id, "messages": [ { - "role": "assistant", - "toolCalls": [ - { - "id": "provider-weather", - "type": "function", - "function": { - "name": "get_weather", - "arguments": '{"city":"Seattle"}', - }, - } - ], + "id": "client-result-1", + "role": "tool", + "toolCallId": "client-call", + "content": "fresh external client result", }, { + "id": "confirm-result-1", "role": "tool", - "toolCallId": "provider-weather", - "content": '{"accepted":true}', + "toolCallId": "confirm-call", + "content": json.dumps( + { + "accepted": True, + "steps": [], + "function_call_id": "unknown-confirm-target", + } + ), }, ], }, ) - assert response.status_code == 200 - events = _decode_sse_events(response) - assert executed == ["Seattle"], events - assert "Translated a legacy AG-UI tool-message approval" in caplog.text + assert confirm_response.status_code == 200 + assert model_calls == 0 + snapshot = await store.get(scope="test", thread_id=thread_id) + assert snapshot is not None + assert snapshot.session_state is not None + context_label = json.loads(snapshot.session_state[config.source_id]["context_label"]) + assert context_label["integrity"] == "untrusted" + assert context_label["confidentiality"] == "private" + protected_response = client.post( + "/fides-confirm-client-result", + json={ + "runId": "run-fides-confirm-protected", + "threadId": thread_id, + "messages": [{"id": "user-2", "role": "user", "content": "Run the protected action"}], + }, + ) -async def test_endpoint_agent_legacy_tool_message_reuses_historical_confirm_changes_call_id() -> None: - """A sole pending local call may reuse an older synthetic confirmation call id.""" + assert protected_response.status_code == 200 + protected_events = _decode_sse_events(protected_response) + assert not [event for event in protected_events if event["type"] == "RUN_ERROR"] + assert executed == [] + assert _is_policy_violation([event["content"] for event in protected_events if event["type"] == "TOOL_CALL_RESULT"]) + + +@pytest.mark.parametrize("confirmation_transport", ["message", "resume"]) +async def test_endpoint_fides_confirm_changes_reconstructs_missing_context_label( + streaming_chat_client_stub: Any, + confirmation_transport: str, +) -> None: + """A confirm-only return reconstructs stored history when FIDES state lacks its label.""" + thread_id = f"thread-fides-confirm-missing-state-{confirmation_transport}" executed: list[str] = [] + model_calls = 0 - def guarded_tool(value: str) -> str: - executed.append(value) - return value + def protected_action() -> str: + executed.append("protected") + return "protected result" - tool = FunctionTool(name="guarded_tool", description="Guarded", func=guarded_tool) - agent = StubAgent( - updates=[AgentResponseUpdate(contents=[Content.from_text(text="Done.")], role="assistant")], - default_options={"tools": [tool]}, - ) - wrapped_agent = AgentFrameworkAgent(agent=agent, require_confirmation=False) - wrapped_agent._approval_state_store.lifecycle.register( - owner=ApprovalExecutionOwner.LOCAL, - thread_id=_default_approval_thread_id("thread-reused-confirm-id"), - interrupt_id="af-call-current", - call_id="provider-reused", - name="guarded_tool", - arguments='{"value":"current"}', - ) - app = FastAPI() - add_agent_framework_fastapi_endpoint(app, wrapped_agent, path="/approval") + async def stream_fn( + messages: list[Message], + options: dict[str, Any], + **kwargs: Any, + ) -> AsyncIterator[ChatResponseUpdate]: + nonlocal model_calls + del messages, options, kwargs + model_calls += 1 + yield ChatResponseUpdate( + role="assistant", + contents=[Content.from_function_call(call_id="protected-call", name="protected_action", arguments={})], + ) - response = TestClient(app).post( - "/approval", - json={ - "runId": "run-reused-confirm-id", - "threadId": "thread-reused-confirm-id", - "messages": [ + config = SecureAgentConfig() + agent = Agent( + client=streaming_chat_client_stub(stream_fn), + tools=[ + FunctionTool( + func=protected_action, + name="protected_action", + description="Run an unrelated protected action", + additional_properties={ + "source_integrity": "trusted", + "max_allowed_confidentiality": "public", + }, + ) + ], + context_providers=[config], + ) + store = InMemoryAGUIThreadSnapshotStore() + confirm_payload = {"accepted": True, "steps": []} + await store.save( + scope="test", + thread_id=thread_id, + snapshot=AGUIThreadSnapshot( + messages=[ + {"id": "user-1", "role": "user", "content": "Look up the external record"}, { + "id": "assistant-external-call", "role": "assistant", + "content": "", "toolCalls": [ { - "id": "provider-reused", + "id": "external-call", "type": "function", - "function": {"name": "confirm_changes", "arguments": "{}"}, + "function": {"name": "client_lookup", "arguments": "{}"}, } ], }, - {"role": "tool", "toolCallId": "provider-reused", "content": '{"accepted":true}'}, { + "id": "external-result", + "role": "tool", + "toolCallId": "external-call", + "content": "stored external result", + }, + { + "id": "assistant-confirm-call", "role": "assistant", + "content": "", "toolCalls": [ { - "id": "provider-reused", + "id": "confirm-call", "type": "function", - "function": {"name": "guarded_tool", "arguments": '{"value":"current"}'}, + "function": { + "name": "confirm_changes", + "arguments": json.dumps({"function_call_id": "unknown-confirm-target"}), + }, } ], }, - {"role": "tool", "toolCallId": "provider-reused", "content": '{"accepted":true}'}, ], - }, - ) - - assert response.status_code == 200 - events = _decode_sse_events(response) - assert executed == ["current"], events - assert not [event for event in events if event.get("type") == "RUN_ERROR"] - - -async def test_endpoint_agent_legacy_tool_message_rejects_reused_call_id( - caplog: pytest.LogCaptureFixture, -) -> None: - """A provider call id shared by retained occurrences cannot authorize either one.""" - executed: list[str] = [] - - def first_tool() -> str: - executed.append("first") - return "first" - - def second_tool() -> str: - executed.append("second") - return "second" - - tools = [ - FunctionTool(name="first_tool", description="First", func=first_tool), - FunctionTool(name="second_tool", description="Second", func=second_tool), - ] - agent = StubAgent( - updates=[AgentResponseUpdate(contents=[Content.from_text(text="Done.")], role="assistant")], - default_options={"tools": tools}, - ) - wrapped_agent = AgentFrameworkAgent(agent=agent, require_confirmation=False) - lifecycle = wrapped_agent._approval_state_store.lifecycle - approval_thread_id = _default_approval_thread_id("thread-legacy-reused") - lifecycle.register( - owner=ApprovalExecutionOwner.LOCAL, - thread_id=approval_thread_id, - interrupt_id="approval-first", - call_id="provider-reused", - name="first_tool", - arguments="{}", - ) - lifecycle.register( - owner=ApprovalExecutionOwner.LOCAL, - thread_id=approval_thread_id, - interrupt_id="approval-second", - call_id="provider-reused", - name="second_tool", - arguments="{}", + session_state={ + config.source_id: { + "scope_id": _session_id_for_thread(scope="test", thread_id=thread_id), + } + }, + ), ) app = FastAPI() - add_agent_framework_fastapi_endpoint(app, wrapped_agent, path="/approval") + add_agent_framework_fastapi_endpoint( + app, + agent, + path="/fides-confirm-missing-state", + state_schema={"steps": {"type": "array", "items": {"type": "object"}}}, + predict_state_config={"steps": {"tool": "draft_steps", "tool_argument": "steps"}}, + snapshot_store=store, + snapshot_scope_resolver=lambda _request: "test", + keepalive_seconds=None, + ) + confirmation_request: dict[str, Any] = { + "runId": f"run-fides-confirm-missing-{confirmation_transport}", + "threadId": thread_id, + "messages": [], + } + if confirmation_transport == "message": + confirmation_request["messages"] = [ + { + "id": "confirm-result", + "role": "tool", + "toolCallId": "confirm-call", + "content": json.dumps(confirm_payload), + } + ] + else: + confirmation_request["resume"] = [ + { + "interruptId": "confirm-call", + "status": "resolved", + "payload": json.dumps(confirm_payload), + } + ] - with caplog.at_level(logging.WARNING, logger="agent_framework"): - response = TestClient(app).post( - "/approval", + with TestClient(app) as client: + confirm_response = client.post("/fides-confirm-missing-state", json=confirmation_request) + + assert confirm_response.status_code == 200 + assert not [event for event in _decode_sse_events(confirm_response) if event["type"] == "RUN_ERROR"] + assert model_calls == 0 + snapshot = await store.get(scope="test", thread_id=thread_id) + assert snapshot is not None + assert snapshot.session_state is not None + context_label = json.loads(snapshot.session_state[config.source_id]["context_label"]) + assert context_label["integrity"] == "untrusted" + assert context_label["confidentiality"] == "private" + + protected_response = client.post( + "/fides-confirm-missing-state", json={ - "runId": "run-legacy-reused", - "threadId": "thread-legacy-reused", - "messages": [ - { - "role": "assistant", - "toolCalls": [ - { - "id": "provider-reused", - "type": "function", - "function": {"name": "first_tool", "arguments": "{}"}, - } - ], - }, - { - "role": "tool", - "toolCallId": "provider-reused", - "content": '{"accepted":true}', - }, - ], + "runId": f"run-fides-confirm-missing-protected-{confirmation_transport}", + "threadId": thread_id, + "messages": [{"id": "user-2", "role": "user", "content": "Run the protected action"}], }, ) - assert response.status_code == 200 - events = _decode_sse_events(response) + assert protected_response.status_code == 200 + protected_events = _decode_sse_events(protected_response) + assert not [event for event in protected_events if event["type"] == "RUN_ERROR"] assert executed == [] - assert [event for event in events if event.get("type") == "RUN_ERROR"] - assert "does not identify exactly one retained pending local occurrence" in events[-1]["message"] + assert _is_policy_violation([event["content"] for event in protected_events if event["type"] == "TOOL_CALL_RESULT"]) -async def test_endpoint_agent_legacy_tool_message_cannot_collide_with_interrupt_id() -> None: - """An unknown provider call id cannot be reinterpreted as a canonical interrupt id.""" - executed: list[str] = [] +async def test_endpoint_fides_service_confirmation_does_not_reconstruct_incomplete_history( + streaming_chat_client_stub: Any, +) -> None: + """A service-owned continuation cannot rebuild missing FIDES state from the AG-UI snapshot alone.""" + thread_id = "thread-fides-service-confirm-missing-state" + model_calls = 0 - def guarded_tool() -> str: - executed.append("ran") - return "done" + async def stream_fn( + messages: list[Message], + options: dict[str, Any], + **kwargs: Any, + ) -> AsyncIterator[ChatResponseUpdate]: + nonlocal model_calls + del messages, options, kwargs + model_calls += 1 + yield ChatResponseUpdate(role="assistant", contents=[Content.from_text("unexpected")]) - tool = FunctionTool(name="guarded_tool", description="Guarded", func=guarded_tool) - agent = StubAgent( - updates=[AgentResponseUpdate(contents=[Content.from_text(text="Done.")], role="assistant")], - default_options={"tools": [tool]}, - ) - wrapped_agent = AgentFrameworkAgent(agent=agent, require_confirmation=False) - wrapped_agent._approval_state_store.lifecycle.register( - owner=ApprovalExecutionOwner.LOCAL, - thread_id=_default_approval_thread_id("thread-legacy-collision"), - interrupt_id="af-call-secret", - call_id="provider-real", - name="guarded_tool", - arguments="{}", + security = SecureAgentConfig() + agent = Agent( + client=streaming_chat_client_stub(stream_fn), + context_providers=[security], ) - app = FastAPI() - add_agent_framework_fastapi_endpoint(app, wrapped_agent, path="/approval") - - response = TestClient(app).post( - "/approval", - json={ - "runId": "run-legacy-collision", - "threadId": "thread-legacy-collision", - "messages": [ + store = InMemoryAGUIThreadSnapshotStore() + await store.save( + scope="test", + thread_id=thread_id, + snapshot=AGUIThreadSnapshot( + messages=[ + {"id": "user-1", "role": "user", "content": "Look up the remote record"}, { + "id": "assistant-external-call", "role": "assistant", + "content": "", "toolCalls": [ { - "id": "af-call-secret", + "id": "external-call", "type": "function", - "function": {"name": "guarded_tool", "arguments": "{}"}, + "function": {"name": "client_lookup", "arguments": "{}"}, } ], }, { + "id": "external-result", "role": "tool", - "actionExecutionId": "af-call-secret", - "content": None, - "result": {"accepted": True}, + "toolCallId": "external-call", + "content": "snapshot-visible external result", }, - ], - }, - ) - - assert response.status_code == 200 - assert executed == [] - events = _decode_sse_events(response) - assert events[-1]["code"] == "APPROVAL_RESUME_REQUIRED" - - -async def test_endpoint_agent_legacy_tool_message_rejects_duplicate_decisions() -> None: - """Conflicting legacy decisions cannot select an earlier approval.""" - executed: list[str] = [] - - def guarded_tool() -> str: - executed.append("ran") - return "done" - - tool = FunctionTool(name="guarded_tool", description="Guarded", func=guarded_tool) - agent = StubAgent( - updates=[AgentResponseUpdate(contents=[Content.from_text(text="Done.")], role="assistant")], - default_options={"tools": [tool]}, - ) - wrapped_agent = AgentFrameworkAgent(agent=agent, require_confirmation=False) - wrapped_agent._approval_state_store.lifecycle.register( - owner=ApprovalExecutionOwner.LOCAL, - thread_id=_default_approval_thread_id("thread-legacy-duplicate"), - interrupt_id="af-call-current", - call_id="provider-call", - name="guarded_tool", - arguments="{}", - ) - app = FastAPI() - add_agent_framework_fastapi_endpoint(app, wrapped_agent, path="/approval") - - response = TestClient(app).post( - "/approval", - json={ - "runId": "run-legacy-duplicate", - "threadId": "thread-legacy-duplicate", - "messages": [ { + "id": "assistant-confirm-call", "role": "assistant", + "content": "", "toolCalls": [ { - "id": "provider-call", + "id": "confirm-call", "type": "function", - "function": {"name": "guarded_tool", "arguments": "{}"}, + "function": {"name": "confirm_changes", "arguments": "{}"}, } ], }, - {"role": "tool", "toolCallId": "provider-call", "content": '{"accepted":true}'}, - {"role": "tool", "toolCallId": "provider-call", "content": '{"accepted":false}'}, ], - }, + session_state={ + security.source_id: { + "scope_id": _session_id_for_thread(scope="test", thread_id=thread_id), + "history_reconstruction_complete": True, + } + }, + ), + ) + wrapped = AgentFrameworkAgent( + agent=agent, + state_schema={"steps": {"type": "array", "items": {"type": "object"}}}, + predict_state_config={"steps": {"tool": "draft_steps", "tool_argument": "steps"}}, + use_service_session=True, + service_session_id_from_thread_id=True, + snapshot_store=store, + ) + app = FastAPI() + add_agent_framework_fastapi_endpoint( + app, + wrapped, + path="/fides-service-confirm-missing-state", + snapshot_scope_resolver=lambda _request: "test", + keepalive_seconds=None, ) + with TestClient(app) as client: + response = client.post( + "/fides-service-confirm-missing-state", + json={ + "runId": "run-fides-service-confirm-missing-state", + "threadId": thread_id, + "messages": [ + { + "id": "confirm-result", + "role": "tool", + "toolCallId": "confirm-call", + "content": json.dumps({"accepted": True, "steps": []}), + } + ], + }, + ) + assert response.status_code == 200 - events = _decode_sse_events(response) - assert executed == [] - assert events[-1]["code"] == "APPROVAL_RESUME_INVALID" - assert "repeats call_id" in events[-1]["message"] + errors = [event for event in _decode_sse_events(response) if event.get("type") == "RUN_ERROR"] + assert [event["code"] for event in errors] == ["RuntimeError"] + assert model_calls == 0 + snapshot = await store.get(scope="test", thread_id=thread_id) + assert snapshot is not None + assert snapshot.session_state is not None + assert "context_label" not in snapshot.session_state[security.source_id] -async def test_endpoint_agent_historical_legacy_approval_cannot_authorize_newer_reused_call() -> None: - """A historical approval outside the submitted turn suffix remains inert.""" - executed: list[str] = [] +@pytest.mark.parametrize("middleware_kind", ["direct", "bundle", "wrapper", "client"]) +@pytest.mark.parametrize("session_state_key", [None, "custom_security_state"], ids=["default-key", "custom-key"]) +@pytest.mark.parametrize("state_origin", ["request", "snapshot"]) +async def test_endpoint_standalone_fides_state_remains_server_owned( + streaming_chat_client_stub: Any, + middleware_kind: str, + session_state_key: str | None, + state_origin: str, +) -> None: + """Shared State cannot replace standalone FIDES security state, including after snapshot restore.""" + client, store, executed, chat_client, state_key = _build_standalone_fides_endpoint( + streaming_chat_client_stub, + middleware_kind=middleware_kind, + session_state_key=session_state_key, + ) + shared_state = {state_key: {"context_label": _TRUSTED_PUBLIC_LABEL}, "client_value": "available"} - def dangerous_tool(value: str) -> str: - executed.append(value) - return value - - tool = FunctionTool(name="dangerous_tool", description="Dangerous", func=dangerous_tool) - agent = StubAgent( - updates=[AgentResponseUpdate(contents=[Content.from_text(text="Done.")], role="assistant")], - default_options={"tools": [tool]}, - ) - wrapped_agent = AgentFrameworkAgent(agent=agent, require_confirmation=False) - wrapped_agent._approval_state_store.lifecycle.register( - owner=ApprovalExecutionOwner.LOCAL, - thread_id=_default_approval_thread_id("thread-historical-legacy"), - interrupt_id="af-call-current", - call_id="provider-reused", - name="dangerous_tool", - arguments='{"value":"new"}', + assert not _is_policy_violation(_post_standalone_fides_turn(client, "taint", turn=1)) + if state_origin == "snapshot": + stored = await store.get(scope="test", thread_id=_STANDALONE_FIDES_THREAD_ID) + assert stored is not None + stored.state = shared_state + await store.save(scope="test", thread_id=_STANDALONE_FIDES_THREAD_ID, snapshot=stored) + forged_turn = _post_standalone_fides_turn( + client, + "act", + turn=2, + state=shared_state if state_origin == "request" else None, ) - app = FastAPI() - add_agent_framework_fastapi_endpoint(app, wrapped_agent, path="/approval") + replay_turn = _post_standalone_fides_turn(client, "act", turn=3) - response = TestClient(app).post( - "/approval", - json={ - "runId": "run-historical-legacy", - "threadId": "thread-historical-legacy", - "messages": [ - { - "role": "assistant", - "toolCalls": [ - { - "id": "provider-reused", - "type": "function", - "function": { - "name": "old_tool", - "arguments": '{"value":"old"}', - }, - } - ], - }, - {"role": "tool", "toolCallId": "provider-reused", "content": '{"accepted":true}'}, - {"role": "user", "content": "Continue without approving anything."}, - ], - }, + assert _is_policy_violation(forged_turn) + assert _is_policy_violation(replay_turn) + assert executed == [] + assert chat_client.last_session is not None + assert chat_client.last_session.state["client_value"] == "available" + snapshot = await store.get(scope="test", thread_id=_STANDALONE_FIDES_THREAD_ID) + assert snapshot is not None + assert snapshot.state is not None + assert snapshot.state["client_value"] == "available" + assert snapshot.session_state is not None + assert json.loads(snapshot.session_state[state_key]["context_label"])["integrity"] == "untrusted" + + +@pytest.mark.parametrize("session_state_key", [None, "custom_security_state"], ids=["default-key", "custom-key"]) +async def test_endpoint_standalone_fides_continuation_persists_without_client_state( + streaming_chat_client_stub: Any, + session_state_key: str | None, +) -> None: + """Protecting FIDES state from Shared State keeps its server-owned continuation across requests.""" + client, store, executed, _, state_key = _build_standalone_fides_endpoint( + streaming_chat_client_stub, + session_state_key=session_state_key, ) - assert response.status_code == 200 - events = _decode_sse_events(response) + _post_standalone_fides_turn(client, "taint", turn=1) + guarded_turn = _post_standalone_fides_turn(client, "act", turn=2) + + assert _is_policy_violation(guarded_turn) assert executed == [] - assert events[-1]["code"] == "APPROVAL_RESUME_REQUIRED" + snapshot = await store.get(scope="test", thread_id=_STANDALONE_FIDES_THREAD_ID) + assert snapshot is not None + assert snapshot.session_state is not None + assert json.loads(snapshot.session_state[state_key]["context_label"])["integrity"] == "untrusted" + assert state_key not in (snapshot.state or {}) -async def test_endpoint_agent_approval_resume_remains_retryable_when_local_tool_is_temporarily_unavailable(): - """A local approval can be retried after its executor disappears before resume.""" - client, agent, executed_cities = _build_weather_approval_endpoint(snapshot_store=InMemoryAGUIThreadSnapshotStore()) - weather_tool = agent.default_options["tools"][0] - agent.default_options["tools"] = [] +@pytest.mark.parametrize( + "forged_value", + [ + "not-a-dict", + {}, + {"scope_id": "another-scope"}, + {"context_label": "{not json"}, + {"context_label": {"integrity": "trusted", "confidentiality": "public"}}, + ], + ids=["non-dict", "empty", "other-scope", "invalid-label-json", "unencoded-label"], +) +async def test_endpoint_standalone_fides_malformed_shared_state_is_inert( + streaming_chat_client_stub: Any, + forged_value: Any, +) -> None: + """Malformed Shared State under the FIDES key neither resets enforcement nor breaks tool calls.""" + client, _, executed, _, _ = _build_standalone_fides_endpoint(streaming_chat_client_stub) - unavailable_response = client.post( - "/approval", - json={ - "runId": "run-unavailable", - "threadId": "thread-weather", - "messages": [], - "resume": [{"interruptId": "call_get_weather", "status": "resolved", "payload": {"accepted": True}}], - }, + _post_standalone_fides_turn(client, "taint", turn=1) + forged_turn = _post_standalone_fides_turn( + client, + "act", + turn=2, + state={_STANDALONE_FIDES_STATE_KEY: forged_value}, ) + replay_turn = _post_standalone_fides_turn(client, "act", turn=3) - assert unavailable_response.status_code == 200 - unavailable_events = _decode_sse_events(unavailable_response) - run_errors = [event for event in unavailable_events if event.get("type") == "RUN_ERROR"] - assert len(run_errors) == 1 - assert run_errors[0]["code"] == "APPROVAL_TOOL_UNAVAILABLE" - assert "temporarily unavailable" in run_errors[0]["message"] - assert executed_cities == [] + assert _is_policy_violation(forged_turn) + assert _is_policy_violation(replay_turn) + assert executed == [] - agent.default_options["tools"] = [weather_tool] - retry_response = client.post( - "/approval", - json={ - "runId": "run-retry", - "threadId": "thread-weather", - "messages": [], - "resume": [{"interruptId": "call_get_weather", "status": "resolved", "payload": {"accepted": True}}], - }, + +async def test_endpoint_standalone_fides_poisoned_snapshot_state_does_not_reset_enforcement( + streaming_chat_client_stub: Any, +) -> None: + """A snapshot whose Shared State already carries a FIDES scope no longer overrides server-owned state.""" + client, store, executed, _, _ = _build_standalone_fides_endpoint(streaming_chat_client_stub) + await store.save( + scope="test", + thread_id=_STANDALONE_FIDES_THREAD_ID, + snapshot=AGUIThreadSnapshot(state={_STANDALONE_FIDES_STATE_KEY: {"context_label": _TRUSTED_PUBLIC_LABEL}}), ) - assert retry_response.status_code == 200 - retry_events = _decode_sse_events(retry_response) - assert not [event for event in retry_events if event.get("type") == "RUN_ERROR"] - assert [ - (event["toolCallId"], event["content"]) for event in retry_events if event.get("type") == "TOOL_CALL_RESULT" - ] == [("call_get_weather", "Sunny in Seattle")] - assert executed_cities == ["Seattle"] + _post_standalone_fides_turn(client, "taint", turn=1) + guarded_turn = _post_standalone_fides_turn(client, "act", turn=2) + assert _is_policy_violation(guarded_turn) + assert executed == [] -@pytest.mark.parametrize("a2ui_mode", ["none", "manual", "automatic"]) -async def test_endpoint_agent_approval_resume_remains_pending_when_invocation_is_disabled( - streaming_chat_client_stub, - a2ui_mode: str, -) -> None: - """A disabled local executor blocks an approved call without consuming its authority.""" - call_id = "call_local_action" + +def _build_tool_approval_auto_endpoint( + streaming_chat_client_stub: Any, +) -> tuple[TestClient, list[str], list[Message], dict[str, str]]: + executed: list[str] = [] + messages_received: list[Message] = [] state = {"phase": "pause"} - local_executions: list[str] = [] - provider_resume_calls: list[str] = [] - function_call = Content.from_function_call( - call_id=call_id, - name="local_action", - arguments={"document": "Approved draft"}, - ) - def local_action(document: str) -> str: - local_executions.append(document) - return "Draft applied" + def auto_tool() -> str: + executed.append("auto") + return "auto result" - local_tool = FunctionTool( - name="local_action", - description="Apply an approved draft.", - func=local_action, - approval_mode="always_require", - ) + def manual_tool() -> str: + executed.append("manual") + return "manual result" + + def auto_approve_auto_tool(function_call: Content) -> bool: + return function_call.name == "auto_tool" async def stream_fn( messages: list[Message], options: dict[str, Any], **kwargs: Any, ) -> AsyncIterator[ChatResponseUpdate]: - del messages, options, kwargs + del options, kwargs if state["phase"] == "pause": yield ChatResponseUpdate( - contents=[Content.from_function_approval_request(id=call_id, function_call=function_call)], + contents=[ + Content.from_function_call(call_id="call_auto", name="auto_tool", arguments="{}"), + Content.from_function_call(call_id="call_manual", name="manual_tool", arguments="{}"), + ], role="assistant", ) return - provider_resume_calls.append("resumed") + messages_received[:] = list(messages) yield ChatResponseUpdate(contents=[Content.from_text(text="Done.")], role="assistant") - chat_client = cast(Any, streaming_chat_client_stub(stream_fn)) - agent = Agent(name="test_agent", instructions="Test", client=chat_client, tools=[local_tool]) - endpoint_agent: SupportsAgentRun = agent - if a2ui_mode != "none": - pytest.importorskip("ag_ui_a2ui_toolkit") - if a2ui_mode == "manual": - from agent_framework_ag_ui._a2ui import enable_a2ui - - endpoint_agent = cast(SupportsAgentRun, enable_a2ui(agent, object())) - wrapped_agent = AgentFrameworkAgent( - agent=endpoint_agent, - require_confirmation=False, - a2ui_config={"inject_a2ui_tool": True} if a2ui_mode == "automatic" else None, + agent = Agent( + name="test_agent", + instructions="Test", + client=streaming_chat_client_stub(stream_fn), + tools=[ + FunctionTool(name="auto_tool", description="Auto tool", func=auto_tool, approval_mode="always_require"), + FunctionTool( + name="manual_tool", description="Manual tool", func=manual_tool, approval_mode="always_require" + ), + ], + middleware=[ToolApprovalMiddleware(auto_approval_rules=[auto_approve_auto_tool])], ) app = FastAPI() add_agent_framework_fastapi_endpoint( app, - wrapped_agent, + AgentFrameworkAgent(agent=agent, require_confirmation=False), path="/approval", - snapshot_store=InMemoryAGUIThreadSnapshotStore(), - snapshot_scope_resolver=lambda _request: "tenant-a", ) - client = TestClient(app) - thread_id = f"thread-disabled-resume-{a2ui_mode}" - forwarded_props = {"injectA2UITool": True} if a2ui_mode == "automatic" else {} + return TestClient(app), executed, messages_received, state - pause_response = client.post( + +async def test_endpoint_agent_approval_resume_entry_executes_approved_tool(): + """A resolved canonical approval resume should execute the pending approved tool.""" + client, _, executed_cities = _build_weather_approval_endpoint() + + response = client.post( "/approval", json={ - "runId": "run-pause", - "threadId": thread_id, - "messages": [{"role": "user", "content": "Apply the draft"}], - "forwardedProps": forwarded_props, - }, - ) - pause_finished = [event for event in _decode_sse_events(pause_response) if event.get("type") == "RUN_FINISHED"] - assert _run_finished_interrupts(pause_finished[-1])[0]["id"] == call_id - - state["phase"] = "resume" - chat_client.function_invocation_configuration["enabled"] = False - resume = [{"interruptId": call_id, "status": "resolved", "payload": {"accepted": True}}] - blocked_response = client.post( - "/approval", - json={ - "runId": "run-blocked", - "threadId": thread_id, + "runId": "run-resume", + "threadId": "thread-weather", "messages": [], - "resume": resume, - "forwardedProps": forwarded_props, + "resume": [ + { + "interruptId": "call_get_weather", + "status": "resolved", + "payload": {"accepted": True}, + } + ], }, ) - blocked_events = _decode_sse_events(blocked_response) - blocked_errors = [event for event in blocked_events if event.get("type") == "RUN_ERROR"] - assert [(event["code"], event["message"]) for event in blocked_errors] == [ - ( - "APPROVAL_INVOCATION_DISABLED", - "Function invocation is disabled; the approved tool remains pending for an explicit retry.", - ) - ] - assert not [event for event in blocked_events if event.get("type") == "TOOL_CALL_RESULT"] - assert local_executions == [] - assert provider_resume_calls == [] - - repeated_blocked_response = client.post( - "/approval", - json={ - "runId": "run-blocked-again", - "threadId": thread_id, - "messages": [], - "resume": resume, - "forwardedProps": forwarded_props, - }, - ) - repeated_blocked_events = _decode_sse_events(repeated_blocked_response) - assert [event["code"] for event in repeated_blocked_events if event.get("type") == "RUN_ERROR"] == [ - "APPROVAL_INVOCATION_DISABLED" - ] - assert not [event for event in repeated_blocked_events if event.get("type") == "TOOL_CALL_RESULT"] - assert local_executions == [] - assert provider_resume_calls == [] + assert response.status_code == 200 + events = _decode_sse_events(response) + tool_results = [event for event in events if event.get("type") == "TOOL_CALL_RESULT"] + assert len(tool_results) == 1 + assert tool_results[0]["toolCallId"] == "call_get_weather" + assert tool_results[0]["content"] == "Sunny in Seattle" + assert executed_cities == ["Seattle"] + assert "outcome" not in [event for event in events if event.get("type") == "RUN_FINISHED"][-1] - chat_client.function_invocation_configuration["enabled"] = True - retry_response = client.post( - "/approval", - json={ - "runId": "run-retry", - "threadId": thread_id, - "messages": [], - "resume": resume, - "forwardedProps": forwarded_props, - }, - ) - retry_events = _decode_sse_events(retry_response) - assert not [event for event in retry_events if event.get("type") == "RUN_ERROR"] - assert [ - (event["toolCallId"], event["content"]) for event in retry_events if event.get("type") == "TOOL_CALL_RESULT" - ] == [(call_id, "Draft applied")] - assert local_executions == ["Approved draft"] - assert provider_resume_calls == ["resumed"] +async def test_endpoint_agent_legacy_tool_message_uses_unique_pending_call_id( + caplog: pytest.LogCaptureFixture, +) -> None: + """Legacy tool messages remain usable only while their provider call id is unambiguous.""" + executed: list[str] = [] + def get_weather(city: str) -> str: + executed.append(city) + return f"Sunny in {city}" -async def test_endpoint_agent_disabled_resume_releases_all_unstarted_local_grants( - streaming_chat_client_stub, -) -> None: - """A disabled mixed batch leaves every unstarted local grant retryable.""" - client, executed, _, state = _build_mixed_approval_batch_endpoint(streaming_chat_client_stub) - pause_response = client.post( + weather_tool = FunctionTool( + name="get_weather", + description="Get the weather for a city", + func=get_weather, + approval_mode="always_require", + ) + function_call = Content.from_function_call( + call_id="provider-weather", + name="get_weather", + arguments={"city": "Seattle"}, + id="af-call-weather", + ) + approval_request = Content.from_function_approval_request( + id="af-call-weather", + function_call=function_call, + ) + agent = StubAgent( + updates=[AgentResponseUpdate(contents=[approval_request], role="assistant")], + default_options={"tools": [weather_tool]}, + ) + wrapped_agent = AgentFrameworkAgent(agent=agent, require_confirmation=False) + app = FastAPI() + add_agent_framework_fastapi_endpoint(app, wrapped_agent, path="/approval") + client = TestClient(app) + pause = client.post( "/approval", json={ "runId": "run-pause", - "threadId": "thread-disabled-mixed-grants", - "messages": [{"role": "user", "content": "Run both tools"}], + "threadId": "thread-legacy-message", + "messages": [{"role": "user", "content": "Weather?"}], }, ) - pause_finished = [event for event in _decode_sse_events(pause_response) if event.get("type") == "RUN_FINISHED"] - approval_id = _run_finished_interrupts(pause_finished[-1])[0]["id"] - resume = [{"interruptId": approval_id, "status": "resolved", "payload": {"accepted": True}}] + assert pause.status_code == 200 + agent.updates = [AgentResponseUpdate(contents=[Content.from_text(text="Done.")], role="assistant")] - state["phase"] = "resume" - chat_client = state["chat_client"] - chat_client.function_invocation_configuration["enabled"] = False - blocked_response = client.post( - "/approval", - json={ - "runId": "run-blocked", - "threadId": "thread-disabled-mixed-grants", - "messages": [], - "resume": resume, - }, - ) + with caplog.at_level(logging.WARNING, logger="agent_framework"): + response = client.post( + "/approval", + json={ + "runId": "run-resume", + "threadId": "thread-legacy-message", + "messages": [ + { + "role": "assistant", + "toolCalls": [ + { + "id": "provider-weather", + "type": "function", + "function": { + "name": "get_weather", + "arguments": '{"city":"Seattle"}', + }, + } + ], + }, + { + "role": "tool", + "toolCallId": "provider-weather", + "content": '{"accepted":true}', + }, + ], + }, + ) - blocked_events = _decode_sse_events(blocked_response) - assert [event["code"] for event in blocked_events if event.get("type") == "RUN_ERROR"] == [ - "APPROVAL_INVOCATION_DISABLED" - ] - assert not [event for event in blocked_events if event.get("type") == "TOOL_CALL_RESULT"] - assert executed == [] + assert response.status_code == 200 + events = _decode_sse_events(response) + assert executed == ["Seattle"], events + assert "Translated a legacy AG-UI tool-message approval" in caplog.text - chat_client.function_invocation_configuration["enabled"] = True - retry_response = client.post( + +async def test_endpoint_agent_legacy_tool_message_reuses_historical_confirm_changes_call_id() -> None: + """A sole pending local call may reuse an older synthetic confirmation call id.""" + executed: list[str] = [] + + def guarded_tool(value: str) -> str: + executed.append(value) + return value + + tool = FunctionTool(name="guarded_tool", description="Guarded", func=guarded_tool) + agent = StubAgent( + updates=[AgentResponseUpdate(contents=[Content.from_text(text="Done.")], role="assistant")], + default_options={"tools": [tool]}, + ) + wrapped_agent = AgentFrameworkAgent(agent=agent, require_confirmation=False) + wrapped_agent._approval_state_store.lifecycle.register( + owner=ApprovalExecutionOwner.LOCAL, + thread_id=_default_approval_thread_id("thread-reused-confirm-id"), + interrupt_id="af-call-current", + call_id="provider-reused", + name="guarded_tool", + arguments='{"value":"current"}', + ) + app = FastAPI() + add_agent_framework_fastapi_endpoint(app, wrapped_agent, path="/approval") + + response = TestClient(app).post( "/approval", json={ - "runId": "run-retry", - "threadId": "thread-disabled-mixed-grants", - "messages": [], - "resume": resume, + "runId": "run-reused-confirm-id", + "threadId": "thread-reused-confirm-id", + "messages": [ + { + "role": "assistant", + "toolCalls": [ + { + "id": "provider-reused", + "type": "function", + "function": {"name": "confirm_changes", "arguments": "{}"}, + } + ], + }, + {"role": "tool", "toolCallId": "provider-reused", "content": '{"accepted":true}'}, + { + "role": "assistant", + "toolCalls": [ + { + "id": "provider-reused", + "type": "function", + "function": {"name": "guarded_tool", "arguments": '{"value":"current"}'}, + } + ], + }, + {"role": "tool", "toolCallId": "provider-reused", "content": '{"accepted":true}'}, + ], }, ) - retry_events = _decode_sse_events(retry_response) - assert not [event for event in retry_events if event.get("type") == "RUN_ERROR"] - assert [ - (event["toolCallId"], event["content"]) for event in retry_events if event.get("type") == "TOOL_CALL_RESULT" - ] == [ - ("call_sensitive", "Sensitive action in Seattle"), - ("call_weather", "Weather in Seattle"), - ] - assert executed == ["sensitive:Seattle", "weather:Seattle"] + assert response.status_code == 200 + events = _decode_sse_events(response) + assert executed == ["current"], events + assert not [event for event in events if event.get("type") == "RUN_ERROR"] -@pytest.mark.parametrize("blocking_reason", ["disabled", "provider_failure", "provider_disable"]) -async def test_endpoint_disabled_resume_keeps_local_and_hosted_grants_retryable( - streaming_chat_client_stub, blocking_reason: str +async def test_endpoint_agent_legacy_tool_message_rejects_reused_call_id( + caplog: pytest.LogCaptureFixture, ) -> None: - """Blocking preparation must not strand any unstarted approval owner.""" + """A provider call id shared by retained occurrences cannot authorize either one.""" executed: list[str] = [] - phase = "pause" - local_call = Content.from_function_call( - id="local-occurrence", call_id="local-call", name="local_action", arguments={} - ) - hosted_call = Content.from_function_call( - call_id="hosted-call", - name="hosted_action", - arguments={}, - additional_properties={"server_label": "remote"}, - ) - class PreparationProvider(ContextProvider): - async def before_run( - self, - *, - agent: SupportsAgentRun, - session: AgentSession, - context: SessionContext, - state: dict[str, Any], - ) -> None: - if phase == "blocked" and blocking_reason == "provider_failure": - raise RuntimeError("Provider is temporarily unavailable") - if phase == "blocked" and blocking_reason == "provider_disable": - invocation = getattr(agent, "client") - assert isinstance(invocation, FunctionInvocationLayer) - invocation.function_invocation_configuration["enabled"] = False + def first_tool() -> str: + executed.append("first") + return "first" - def local_action() -> str: - executed.append("local") - return "Local completed" + def second_tool() -> str: + executed.append("second") + return "second" - async def stream_fn( - messages: list[Message], options: dict[str, Any], **kwargs: Any - ) -> AsyncIterator[ChatResponseUpdate]: - if phase == "pause": - yield ChatResponseUpdate( - role="assistant", - contents=[ - Content.from_function_approval_request(id="local-occurrence", function_call=local_call), - Content.from_function_approval_request(id="hosted-approval", function_call=hosted_call), + tools = [ + FunctionTool(name="first_tool", description="First", func=first_tool), + FunctionTool(name="second_tool", description="Second", func=second_tool), + ] + agent = StubAgent( + updates=[AgentResponseUpdate(contents=[Content.from_text(text="Done.")], role="assistant")], + default_options={"tools": tools}, + ) + wrapped_agent = AgentFrameworkAgent(agent=agent, require_confirmation=False) + lifecycle = wrapped_agent._approval_state_store.lifecycle + approval_thread_id = _default_approval_thread_id("thread-legacy-reused") + lifecycle.register( + owner=ApprovalExecutionOwner.LOCAL, + thread_id=approval_thread_id, + interrupt_id="approval-first", + call_id="provider-reused", + name="first_tool", + arguments="{}", + ) + lifecycle.register( + owner=ApprovalExecutionOwner.LOCAL, + thread_id=approval_thread_id, + interrupt_id="approval-second", + call_id="provider-reused", + name="second_tool", + arguments="{}", + ) + app = FastAPI() + add_agent_framework_fastapi_endpoint(app, wrapped_agent, path="/approval") + + with caplog.at_level(logging.WARNING, logger="agent_framework"): + response = TestClient(app).post( + "/approval", + json={ + "runId": "run-legacy-reused", + "threadId": "thread-legacy-reused", + "messages": [ + { + "role": "assistant", + "toolCalls": [ + { + "id": "provider-reused", + "type": "function", + "function": {"name": "first_tool", "arguments": "{}"}, + } + ], + }, + { + "role": "tool", + "toolCallId": "provider-reused", + "content": '{"accepted":true}', + }, ], - ) - return - for message in messages: - for content in message.contents: - if content.type == "function_approval_response" and content.id == "hosted-approval": - assert content.approved is True - executed.append("hosted") - yield ChatResponseUpdate(role="assistant", contents=[Content.from_text("Done.")]) + }, + ) - chat_client = streaming_chat_client_stub(stream_fn) - agent = Agent( - client=chat_client, - context_providers=[PreparationProvider("preparation")], - tools=[ - FunctionTool( - name="local_action", description="Local action", func=local_action, approval_mode="always_require" - ) - ], + assert response.status_code == 200 + events = _decode_sse_events(response) + assert executed == [] + assert [event for event in events if event.get("type") == "RUN_ERROR"] + assert "does not identify exactly one retained pending local occurrence" in events[-1]["message"] + + +async def test_endpoint_agent_legacy_tool_message_cannot_collide_with_interrupt_id() -> None: + """An unknown provider call id cannot be reinterpreted as a canonical interrupt id.""" + executed: list[str] = [] + + def guarded_tool() -> str: + executed.append("ran") + return "done" + + tool = FunctionTool(name="guarded_tool", description="Guarded", func=guarded_tool) + agent = StubAgent( + updates=[AgentResponseUpdate(contents=[Content.from_text(text="Done.")], role="assistant")], + default_options={"tools": [tool]}, + ) + wrapped_agent = AgentFrameworkAgent(agent=agent, require_confirmation=False) + wrapped_agent._approval_state_store.lifecycle.register( + owner=ApprovalExecutionOwner.LOCAL, + thread_id=_default_approval_thread_id("thread-legacy-collision"), + interrupt_id="af-call-secret", + call_id="provider-real", + name="guarded_tool", + arguments="{}", ) - assert isinstance(chat_client, FunctionInvocationLayer) app = FastAPI() - add_agent_framework_fastapi_endpoint(app, agent, path="/approval") - client = TestClient(app) - pause = client.post( + add_agent_framework_fastapi_endpoint(app, wrapped_agent, path="/approval") + + response = TestClient(app).post( "/approval", json={ - "runId": "pause", - "threadId": "cross-owner", - "messages": [{"role": "user", "content": "Run both actions"}], + "runId": "run-legacy-collision", + "threadId": "thread-legacy-collision", + "messages": [ + { + "role": "assistant", + "toolCalls": [ + { + "id": "af-call-secret", + "type": "function", + "function": {"name": "guarded_tool", "arguments": "{}"}, + } + ], + }, + { + "role": "tool", + "actionExecutionId": "af-call-secret", + "content": None, + "result": {"accepted": True}, + }, + ], }, ) - finished = [event for event in _decode_sse_events(pause) if event["type"] == "RUN_FINISHED"] - interrupts = _run_finished_interrupts(finished[-1]) - assert {interrupt["id"] for interrupt in interrupts} == {"local-occurrence", "hosted-approval"} - resume = [ - {"interruptId": interrupt["id"], "status": "resolved", "payload": {"approved": True}} - for interrupt in interrupts - ] - phase = "blocked" - chat_client.function_invocation_configuration["enabled"] = blocking_reason != "disabled" - for run_id in ("blocked", "blocked-again"): - response = client.post( - "/approval", - json={"runId": run_id, "threadId": "cross-owner", "messages": [], "resume": resume}, - ) - assert [event["code"] for event in _decode_sse_events(response) if event["type"] == "RUN_ERROR"] == [ - "RuntimeError" if blocking_reason == "provider_failure" else "APPROVAL_INVOCATION_DISABLED" - ] - assert executed == [] - phase = "resume" - chat_client.function_invocation_configuration["enabled"] = True - for run_id in ("retry", "completed-replay"): - response = client.post( - "/approval", - json={"runId": run_id, "threadId": "cross-owner", "messages": [], "resume": resume}, - ) - assert not [event for event in _decode_sse_events(response) if event["type"] == "RUN_ERROR"] - assert executed == ["local", "hosted"] + assert response.status_code == 200 + assert executed == [] + events = _decode_sse_events(response) + assert events[-1]["code"] == "APPROVAL_RESUME_REQUIRED" -@pytest.mark.parametrize("failure_point", ["middleware", "tool"]) -async def test_endpoint_fatal_approval_preserves_unstarted_siblings( - streaming_chat_client_stub, failure_point: str -) -> None: - """A denied local approval must not strand untouched local or hosted siblings.""" +async def test_endpoint_agent_legacy_tool_message_rejects_duplicate_decisions() -> None: + """Conflicting legacy decisions cannot select an earlier approval.""" executed: list[str] = [] - phase = "pause" - - class DenyFirst(FunctionMiddleware): - async def process(self, context: Any, call_next: Any) -> None: - if context.function.name == "denied_action" and failure_point == "middleware": - raise MiddlewareFailure("Action denied") - if context.function.name == "allowed_action" and phase == "resume": - # Keep this sibling outside invocation while core cancels the failed parallel batch. - await asyncio.Event().wait() - await call_next() - def denied_action() -> str: - executed.append("denied") - raise MiddlewareFailure("Action failed after starting") - - def allowed_action() -> str: - executed.append("allowed") - return "Allowed action completed" + def guarded_tool() -> str: + executed.append("ran") + return "done" - calls = [ - Content.from_function_call(id="denied", call_id="denied-call", name="denied_action", arguments={}), - Content.from_function_call(id="allowed", call_id="allowed-call", name="allowed_action", arguments={}), - Content.from_function_call( - id="hosted", - call_id="hosted-call", - name="hosted_action", - arguments={}, - additional_properties={"server_label": "remote"}, - ), - ] + tool = FunctionTool(name="guarded_tool", description="Guarded", func=guarded_tool) + agent = StubAgent( + updates=[AgentResponseUpdate(contents=[Content.from_text(text="Done.")], role="assistant")], + default_options={"tools": [tool]}, + ) + wrapped_agent = AgentFrameworkAgent(agent=agent, require_confirmation=False) + wrapped_agent._approval_state_store.lifecycle.register( + owner=ApprovalExecutionOwner.LOCAL, + thread_id=_default_approval_thread_id("thread-legacy-duplicate"), + interrupt_id="af-call-current", + call_id="provider-call", + name="guarded_tool", + arguments="{}", + ) + app = FastAPI() + add_agent_framework_fastapi_endpoint(app, wrapped_agent, path="/approval") - async def stream_fn( - messages: list[Message], options: dict[str, Any], **kwargs: Any - ) -> AsyncIterator[ChatResponseUpdate]: - if phase == "pause": - yield ChatResponseUpdate( - role="assistant", - contents=[ - Content.from_function_approval_request(id=identity, function_call=call) - for identity, call in zip(("denied", "allowed", "hosted"), calls, strict=True) - ], - ) - return - for message in messages: - for content in message.contents: - if content.type == "function_approval_response" and content.id == "hosted": - assert content.approved is True - executed.append("hosted") - yield ChatResponseUpdate(role="assistant", contents=[Content.from_text("Done.")]) + response = TestClient(app).post( + "/approval", + json={ + "runId": "run-legacy-duplicate", + "threadId": "thread-legacy-duplicate", + "messages": [ + { + "role": "assistant", + "toolCalls": [ + { + "id": "provider-call", + "type": "function", + "function": {"name": "guarded_tool", "arguments": "{}"}, + } + ], + }, + {"role": "tool", "toolCallId": "provider-call", "content": '{"accepted":true}'}, + {"role": "tool", "toolCallId": "provider-call", "content": '{"accepted":false}'}, + ], + }, + ) - agent = Agent( - client=streaming_chat_client_stub(stream_fn), - tools=[ - FunctionTool( - name="denied_action", description="Denied action", func=denied_action, approval_mode="always_require" - ), - FunctionTool( - name="allowed_action", description="Allowed action", func=allowed_action, approval_mode="always_require" - ), - ], - middleware=[DenyFirst()], + assert response.status_code == 200 + events = _decode_sse_events(response) + assert executed == [] + assert events[-1]["code"] == "APPROVAL_RESUME_INVALID" + assert "repeats call_id" in events[-1]["message"] + + +async def test_endpoint_agent_historical_legacy_approval_cannot_authorize_newer_reused_call() -> None: + """A historical approval outside the submitted turn suffix remains inert.""" + executed: list[str] = [] + + def dangerous_tool(value: str) -> str: + executed.append(value) + return value + + tool = FunctionTool(name="dangerous_tool", description="Dangerous", func=dangerous_tool) + agent = StubAgent( + updates=[AgentResponseUpdate(contents=[Content.from_text(text="Done.")], role="assistant")], + default_options={"tools": [tool]}, + ) + wrapped_agent = AgentFrameworkAgent(agent=agent, require_confirmation=False) + wrapped_agent._approval_state_store.lifecycle.register( + owner=ApprovalExecutionOwner.LOCAL, + thread_id=_default_approval_thread_id("thread-historical-legacy"), + interrupt_id="af-call-current", + call_id="provider-reused", + name="dangerous_tool", + arguments='{"value":"new"}', ) app = FastAPI() - add_agent_framework_fastapi_endpoint(app, agent, path="/approval") - client = TestClient(app) - pause = client.post( + add_agent_framework_fastapi_endpoint(app, wrapped_agent, path="/approval") + + response = TestClient(app).post( "/approval", - json={"runId": "pause", "threadId": "fatal-batch", "messages": [{"role": "user", "content": "Act"}]}, + json={ + "runId": "run-historical-legacy", + "threadId": "thread-historical-legacy", + "messages": [ + { + "role": "assistant", + "toolCalls": [ + { + "id": "provider-reused", + "type": "function", + "function": { + "name": "old_tool", + "arguments": '{"value":"old"}', + }, + } + ], + }, + {"role": "tool", "toolCallId": "provider-reused", "content": '{"accepted":true}'}, + {"role": "user", "content": "Continue without approving anything."}, + ], + }, ) - finished = [event for event in _decode_sse_events(pause) if event["type"] == "RUN_FINISHED"] - assert {item["id"] for item in _run_finished_interrupts(finished[-1])} == {"denied", "allowed", "hosted"} - phase = "resume" - failed = client.post( + + assert response.status_code == 200 + events = _decode_sse_events(response) + assert executed == [] + assert events[-1]["code"] == "APPROVAL_RESUME_REQUIRED" + + +async def test_endpoint_agent_approval_resume_remains_retryable_when_local_tool_is_temporarily_unavailable(): + """A local approval can be retried after its executor disappears before resume.""" + client, agent, executed_cities = _build_weather_approval_endpoint(snapshot_store=InMemoryAGUIThreadSnapshotStore()) + weather_tool = agent.default_options["tools"][0] + agent.default_options["tools"] = [] + + unavailable_response = client.post( "/approval", json={ - "runId": "denied", - "threadId": "fatal-batch", + "runId": "run-unavailable", + "threadId": "thread-weather", "messages": [], - "resume": [ - {"interruptId": identity, "status": "resolved", "payload": {"approved": True}} - for identity in ("denied", "allowed", "hosted") - ], + "resume": [{"interruptId": "call_get_weather", "status": "resolved", "payload": {"accepted": True}}], }, ) - assert [event["code"] for event in _decode_sse_events(failed) if event["type"] == "RUN_ERROR"] == [ - "MiddlewareFailure" - ] - failed_effects = ["denied"] if failure_point == "tool" else [] - assert executed == failed_effects - denied_retry = client.post( + assert unavailable_response.status_code == 200 + unavailable_events = _decode_sse_events(unavailable_response) + run_errors = [event for event in unavailable_events if event.get("type") == "RUN_ERROR"] + assert len(run_errors) == 1 + assert run_errors[0]["code"] == "APPROVAL_TOOL_UNAVAILABLE" + assert "temporarily unavailable" in run_errors[0]["message"] + assert executed_cities == [] + + agent.default_options["tools"] = [weather_tool] + retry_response = client.post( "/approval", json={ - "runId": "unsafe-retry", - "threadId": "fatal-batch", + "runId": "run-retry", + "threadId": "thread-weather", "messages": [], - "resume": [ - {"interruptId": identity, "status": "resolved", "payload": {"approved": True}} - for identity in ("denied", "allowed", "hosted") - ], + "resume": [{"interruptId": "call_get_weather", "status": "resolved", "payload": {"accepted": True}}], }, ) - assert [event for event in _decode_sse_events(denied_retry) if event["type"] == "RUN_ERROR"] - assert executed == failed_effects - phase = "retry" - for run_id in ("retry-siblings", "replay-completed"): - response = client.post( - "/approval", - json={ - "runId": run_id, - "threadId": "fatal-batch", - "messages": [], - "resume": [ - *([{"interruptId": "denied", "status": "cancelled"}] if failure_point == "middleware" else []), - *[ - {"interruptId": identity, "status": "resolved", "payload": {"approved": True}} - for identity in ("allowed", "hosted") - ], - ], - }, - ) - assert not [event for event in _decode_sse_events(response) if event["type"] == "RUN_ERROR"] - assert executed == [*failed_effects, "allowed", "hosted"] + assert retry_response.status_code == 200 + retry_events = _decode_sse_events(retry_response) + assert not [event for event in retry_events if event.get("type") == "RUN_ERROR"] + assert [ + (event["toolCallId"], event["content"]) for event in retry_events if event.get("type") == "TOOL_CALL_RESULT" + ] == [("call_get_weather", "Sunny in Seattle")] + assert executed_cities == ["Seattle"] -@pytest.mark.parametrize("reuse_call_id", [False, True]) -async def test_endpoint_retained_and_pending_occurrences_resume_independently( - streaming_chat_client_stub, reuse_call_id: bool +@pytest.mark.parametrize("a2ui_mode", ["none", "manual", "automatic"]) +async def test_endpoint_agent_approval_resume_remains_pending_when_invocation_is_disabled( + streaming_chat_client_stub, + a2ui_mode: str, ) -> None: - """A retained result cannot invalidate a distinct pending approval with the same provider id.""" - phase = "first" - executed: list[str] = [] + """A disabled local executor blocks an approved call without consuming its authority.""" + call_id = "call_local_action" + state = {"phase": "pause"} + local_executions: list[str] = [] + provider_resume_calls: list[str] = [] + function_call = Content.from_function_call( + call_id=call_id, + name="local_action", + arguments={"document": "Approved draft"}, + ) - def record(value: str) -> str: - executed.append(value) - return f"Recorded {value}" + def local_action(document: str) -> str: + local_executions.append(document) + return "Draft applied" + + local_tool = FunctionTool( + name="local_action", + description="Apply an approved draft.", + func=local_action, + approval_mode="always_require", + ) async def stream_fn( - messages: list[Message], options: dict[str, Any], **kwargs: Any + messages: list[Message], + options: dict[str, Any], + **kwargs: Any, ) -> AsyncIterator[ChatResponseUpdate]: - if phase == "complete": - yield ChatResponseUpdate(role="assistant", contents=[Content.from_text("Done.")]) + del messages, options, kwargs + if state["phase"] == "pause": + yield ChatResponseUpdate( + contents=[Content.from_function_approval_request(id=call_id, function_call=function_call)], + role="assistant", + ) return - call = Content.from_function_call( - id=phase, - call_id="shared" if reuse_call_id else f"provider-{phase}", - name="record", - arguments={"value": phase}, - ) - yield ChatResponseUpdate( - role="assistant", contents=[Content.from_function_approval_request(id=phase, function_call=call)] - ) + provider_resume_calls.append("resumed") + yield ChatResponseUpdate(contents=[Content.from_text(text="Done.")], role="assistant") - agent = Agent( - client=streaming_chat_client_stub(stream_fn), - tools=[FunctionTool(name="record", description="Record", func=record, approval_mode="always_require")], + chat_client = cast(Any, streaming_chat_client_stub(stream_fn)) + agent = Agent(name="test_agent", instructions="Test", client=chat_client, tools=[local_tool]) + endpoint_agent: SupportsAgentRun = agent + if a2ui_mode != "none": + pytest.importorskip("ag_ui_a2ui_toolkit") + if a2ui_mode == "manual": + from agent_framework_ag_ui._a2ui import enable_a2ui + + endpoint_agent = cast(SupportsAgentRun, enable_a2ui(agent, object())) + wrapped_agent = AgentFrameworkAgent( + agent=endpoint_agent, + require_confirmation=False, + a2ui_config={"inject_a2ui_tool": True} if a2ui_mode == "automatic" else None, ) app = FastAPI() - add_agent_framework_fastapi_endpoint(app, agent, path="/approval") - client = TestClient(app) - for occurrence_id in ("first", "second"): - phase = occurrence_id - pause = client.post( - "/approval", - json={ - "runId": f"pause-{occurrence_id}", - "threadId": "retained-occurrences", - "messages": [{"role": "user", "content": f"Record {occurrence_id}"}], - }, - ) - finished = [event for event in _decode_sse_events(pause) if event["type"] == "RUN_FINISHED"] - assert [item["id"] for item in _run_finished_interrupts(finished[-1])] == [occurrence_id] - if occurrence_id == "first": - phase = "complete" - completed = client.post( - "/approval", - json={ - "runId": "complete-first", - "threadId": "retained-occurrences", - "messages": [], - "resume": [{"interruptId": "first", "status": "resolved", "payload": {"approved": True}}], - }, - ) - assert not [event for event in _decode_sse_events(completed) if event["type"] == "RUN_ERROR"] - assert executed == ["first"] - - phase = "complete" - for run_id in ("complete-both", "repeat-completed"): - response = client.post( - "/approval", - json={ - "runId": run_id, - "threadId": "retained-occurrences", - "messages": [], - "resume": [ - {"interruptId": identity, "status": "resolved", "payload": {"approved": True}} - for identity in ("first", "second") - ], - }, - ) - events = _decode_sse_events(response) - assert not [event for event in events if event["type"] == "RUN_ERROR"] - assert executed == ["first", "second"] - assert [event["content"] for event in events if event["type"] == "TOOL_CALL_RESULT"] == [ - "Recorded first", - "Recorded second", - ] - - -async def test_endpoint_agent_disabled_mixed_resume_preserves_terminal_progress_for_retry() -> None: - """A blocked mixed resume keeps grants retryable without reviving terminal siblings.""" - executed: list[str] = [] - observed_non_grants: list[tuple[str | None, bool]] = [] - - class ObserveNonGrants(FunctionMiddleware): - async def process(self, context: Any, call_next: Any) -> None: - del context - await call_next() - - def _on_approval_responses( - self, - responses: Sequence[Content], - *, - session: AgentSession | None, - ) -> None: - del session - observed_non_grants.extend( - ( - response.function_call.call_id if response.function_call else None, - response.additional_properties.get("cancelled") is True, - ) - for response in responses - ) - - def record_city(city: str) -> str: - executed.append(city) - return f"Recorded {city}" - - tool = FunctionTool( - name="record_city", - description="Record a city", - func=record_city, - approval_mode="always_require", - ) - approval_requests = [ - Content.from_function_approval_request( - id=f"approval-{city.lower()}", - function_call=Content.from_function_call( - call_id=f"call-{city.lower()}", - name="record_city", - arguments={"city": city}, - ), - ) - for city in ("Seattle", "Portland", "Tokyo") - ] - client_config = Mock(function_invocation_configuration={"enabled": True}) - - class ObservedAgent(StubAgent): - middleware = [ObserveNonGrants()] - - agent = ObservedAgent( - updates=[AgentResponseUpdate(contents=approval_requests, role="assistant")], - default_options={"tools": [tool]}, - client=client_config, - ) - wrapped_agent = AgentFrameworkAgent(agent=agent, require_confirmation=False) - snapshot_store = InMemoryAGUIThreadSnapshotStore() - app = FastAPI() add_agent_framework_fastapi_endpoint( app, wrapped_agent, path="/approval", - snapshot_store=snapshot_store, + snapshot_store=InMemoryAGUIThreadSnapshotStore(), snapshot_scope_resolver=lambda _request: "tenant-a", ) client = TestClient(app) - thread_id = "thread-disabled-mixed-decisions" - pause = client.post( + thread_id = f"thread-disabled-resume-{a2ui_mode}" + forwarded_props = {"injectA2UITool": True} if a2ui_mode == "automatic" else {} + + pause_response = client.post( "/approval", json={ "runId": "run-pause", "threadId": thread_id, - "messages": [{"role": "user", "content": "Record three cities"}], + "messages": [{"role": "user", "content": "Apply the draft"}], + "forwardedProps": forwarded_props, }, ) - pause_finished = [event for event in _decode_sse_events(pause) if event.get("type") == "RUN_FINISHED"] - assert {interrupt["id"] for interrupt in _run_finished_interrupts(pause_finished[-1])} == { - "call-seattle", - "call-portland", - "call-tokyo", - } + pause_finished = [event for event in _decode_sse_events(pause_response) if event.get("type") == "RUN_FINISHED"] + assert _run_finished_interrupts(pause_finished[-1])[0]["id"] == call_id - agent.updates = [AgentResponseUpdate(contents=[Content.from_text(text="Done.")], role="assistant")] - mixed_resume = [ - { - "interruptId": "call-seattle", - "status": "resolved", - "payload": {"approved": True}, - }, - { - "interruptId": "call-portland", - "status": "resolved", - "payload": {"approved": False}, + state["phase"] = "resume" + chat_client.function_invocation_configuration["enabled"] = False + resume = [{"interruptId": call_id, "status": "resolved", "payload": {"accepted": True}}] + blocked_response = client.post( + "/approval", + json={ + "runId": "run-blocked", + "threadId": thread_id, + "messages": [], + "resume": resume, + "forwardedProps": forwarded_props, }, - {"interruptId": "call-tokyo", "status": "cancelled"}, - ] - client_config.function_invocation_configuration["enabled"] = False + ) - for run_id in ("run-blocked", "run-blocked-again"): - blocked = client.post( - "/approval", - json={"runId": run_id, "threadId": thread_id, "messages": [], "resume": mixed_resume}, + blocked_events = _decode_sse_events(blocked_response) + blocked_errors = [event for event in blocked_events if event.get("type") == "RUN_ERROR"] + assert [(event["code"], event["message"]) for event in blocked_errors] == [ + ( + "APPROVAL_INVOCATION_DISABLED", + "Function invocation is disabled; the approved tool remains pending for an explicit retry.", ) - blocked_events = _decode_sse_events(blocked) - assert [event["code"] for event in blocked_events if event.get("type") == "RUN_ERROR"] == [ - "APPROVAL_INVOCATION_DISABLED" - ] - assert not [event for event in blocked_events if event.get("type") == "TOOL_CALL_RESULT"] - assert executed == [] + ] + assert not [event for event in blocked_events if event.get("type") == "TOOL_CALL_RESULT"] + assert local_executions == [] + assert provider_resume_calls == [] - assert set(observed_non_grants) == { - ("call-portland", False), - ("call-tokyo", True), - } - snapshot = await snapshot_store.get(scope="tenant-a", thread_id=thread_id) - assert snapshot is not None and snapshot.interrupt is not None - assert {interrupt["id"] for interrupt in snapshot.interrupt} == {"call-seattle"} + repeated_blocked_response = client.post( + "/approval", + json={ + "runId": "run-blocked-again", + "threadId": thread_id, + "messages": [], + "resume": resume, + "forwardedProps": forwarded_props, + }, + ) + repeated_blocked_events = _decode_sse_events(repeated_blocked_response) + assert [event["code"] for event in repeated_blocked_events if event.get("type") == "RUN_ERROR"] == [ + "APPROVAL_INVOCATION_DISABLED" + ] + assert not [event for event in repeated_blocked_events if event.get("type") == "TOOL_CALL_RESULT"] + assert local_executions == [] + assert provider_resume_calls == [] - client_config.function_invocation_configuration["enabled"] = True - retry = client.post( + chat_client.function_invocation_configuration["enabled"] = True + retry_response = client.post( "/approval", - json={"runId": "run-retry", "threadId": thread_id, "messages": [], "resume": mixed_resume}, + json={ + "runId": "run-retry", + "threadId": thread_id, + "messages": [], + "resume": resume, + "forwardedProps": forwarded_props, + }, ) - retry_events = _decode_sse_events(retry) + + retry_events = _decode_sse_events(retry_response) assert not [event for event in retry_events if event.get("type") == "RUN_ERROR"] assert [ (event["toolCallId"], event["content"]) for event in retry_events if event.get("type") == "TOOL_CALL_RESULT" - ] == [("call-seattle", "Recorded Seattle")] - assert executed == ["Seattle"] - + ] == [(call_id, "Draft applied")] + assert local_executions == ["Approved draft"] + assert provider_resume_calls == ["resumed"] -async def test_endpoint_agent_mixed_retry_reprojects_completed_sibling_without_reexecution() -> None: - """A completed sibling is replayed while only the unfinished approval executes.""" - executed: list[str] = [] - def record_city(city: str) -> str: - executed.append(city) - return f"Recorded {city}" +async def test_endpoint_agent_disabled_resume_releases_all_unstarted_local_grants( + streaming_chat_client_stub, +) -> None: + """A disabled mixed batch leaves every unstarted local grant retryable.""" + client, executed, _, state = _build_mixed_approval_batch_endpoint(streaming_chat_client_stub) + pause_response = client.post( + "/approval", + json={ + "runId": "run-pause", + "threadId": "thread-disabled-mixed-grants", + "messages": [{"role": "user", "content": "Run both tools"}], + }, + ) + pause_finished = [event for event in _decode_sse_events(pause_response) if event.get("type") == "RUN_FINISHED"] + approval_id = _run_finished_interrupts(pause_finished[-1])[0]["id"] + resume = [{"interruptId": approval_id, "status": "resolved", "payload": {"accepted": True}}] - tool = FunctionTool(name="record_city", description="Record a city", func=record_city) - agent = StubAgent( - updates=[AgentResponseUpdate(contents=[Content.from_text(text="Done.")], role="assistant")], - default_options={"tools": [tool]}, - ) - wrapped_agent = AgentFrameworkAgent(agent=agent, require_confirmation=False) - lifecycle = wrapped_agent._approval_state_store.lifecycle - approval_thread_id = _default_approval_thread_id("thread-completed-sibling") - completed = lifecycle.register( - owner=ApprovalExecutionOwner.LOCAL, - thread_id=approval_thread_id, - interrupt_id="approval-seattle", - call_id="call-seattle", - name="record_city", - arguments='{"city":"Seattle"}', - ) - completed_intent = lifecycle.claim( - thread_id=approval_thread_id, - decision=ResumeDecision( - interrupt_id="approval-seattle", - accepted=True, - arguments='{"city":"Seattle"}', - name="record_city", - original_arguments='{"city":"Seattle"}', - ), - ) - lifecycle.begin_execution(completed_intent, owner=ApprovalExecutionOwner.LOCAL) - lifecycle.settle( - completed_intent, - [Content.from_function_result(call_id=completed.identity.call_id, result="Recorded Seattle")], - ) - lifecycle.register( - owner=ApprovalExecutionOwner.LOCAL, - thread_id=approval_thread_id, - interrupt_id="approval-portland", - call_id="call-portland", - name="record_city", - arguments='{"city":"Portland"}', - ) - app = FastAPI() - add_agent_framework_fastapi_endpoint(app, wrapped_agent, path="/approval") - client = TestClient(app) - resume = [ - { - "interruptId": "approval-seattle", - "status": "resolved", - "payload": {"approved": True}, - }, - { - "interruptId": "approval-portland", - "status": "resolved", - "payload": {"approved": True}, - }, - ] - - response = client.post( - "/approval", - json={ - "runId": "run-completed-sibling", - "threadId": "thread-completed-sibling", - "messages": [], - "resume": resume, - }, + state["phase"] = "resume" + chat_client = state["chat_client"] + chat_client.function_invocation_configuration["enabled"] = False + blocked_response = client.post( + "/approval", + json={ + "runId": "run-blocked", + "threadId": "thread-disabled-mixed-grants", + "messages": [], + "resume": resume, + }, ) - events = _decode_sse_events(response) - assert not [event for event in events if event.get("type") == "RUN_ERROR"] - assert [(event["toolCallId"], event["content"]) for event in events if event.get("type") == "TOOL_CALL_RESULT"] == [ - ("call-seattle", "Recorded Seattle"), - ("call-portland", "Recorded Portland"), + blocked_events = _decode_sse_events(blocked_response) + assert [event["code"] for event in blocked_events if event.get("type") == "RUN_ERROR"] == [ + "APPROVAL_INVOCATION_DISABLED" ] - assert executed == ["Portland"] - - -async def test_endpoint_agent_approval_resume_releases_already_approved_sibling(streaming_chat_client_stub): - """Resuming a visible approval should also complete never-require siblings from the same batch.""" - client, executed, messages_received, state = _build_mixed_approval_batch_endpoint(streaming_chat_client_stub) - - pause_response = client.post( - "/approval", - json={ - "runId": "run-pause", - "threadId": "thread-mixed-batch", - "messages": [{"role": "user", "content": "Run both tools"}], - }, - ) - - assert pause_response.status_code == 200 - pause_events = _decode_sse_events(pause_response) - pause_finished = [event for event in pause_events if event.get("type") == "RUN_FINISHED"] - interrupts = _run_finished_interrupts(pause_finished[-1]) - assert len(interrupts) == 1 - approval_id = interrupts[0]["id"] - assert approval_id.startswith("af-call-") - assert interrupts[0]["toolCallId"] == "call_sensitive" - assert not [event for event in pause_events if event.get("type") == "TOOL_CALL_RESULT"] + assert not [event for event in blocked_events if event.get("type") == "TOOL_CALL_RESULT"] + assert executed == [] - state["phase"] = "resume" - resume_response = client.post( + chat_client.function_invocation_configuration["enabled"] = True + retry_response = client.post( "/approval", json={ - "runId": "run-resume", - "threadId": "thread-mixed-batch", + "runId": "run-retry", + "threadId": "thread-disabled-mixed-grants", "messages": [], - "resume": [{"interruptId": approval_id, "status": "resolved", "payload": {"accepted": True}}], + "resume": resume, }, ) - assert resume_response.status_code == 200 - resume_events = _decode_sse_events(resume_response) - tool_results = [event for event in resume_events if event.get("type") == "TOOL_CALL_RESULT"] - assert [(event["toolCallId"], event["content"]) for event in tool_results] == [ + retry_events = _decode_sse_events(retry_response) + assert not [event for event in retry_events if event.get("type") == "RUN_ERROR"] + assert [ + (event["toolCallId"], event["content"]) for event in retry_events if event.get("type") == "TOOL_CALL_RESULT" + ] == [ ("call_sensitive", "Sensitive action in Seattle"), ("call_weather", "Weather in Seattle"), ] assert executed == ["sensitive:Seattle", "weather:Seattle"] - assert not [ - event - for event in resume_events - if event.get("type") == "TOOL_CALL_START" and event.get("toolCallId") == "call_weather" - ] - replayed_results = [ - content for message in messages_received for content in message.contents if content.type == "function_result" - ] - replayed_call_ids = [content.call_id for content in replayed_results if content.call_id is not None] - assert sorted(replayed_call_ids) == ["call_sensitive", "call_weather"] -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.""" +@pytest.mark.parametrize("blocking_reason", ["disabled", "provider_failure", "provider_disable"]) +async def test_endpoint_disabled_resume_keeps_local_and_hosted_grants_retryable( + streaming_chat_client_stub, blocking_reason: str +) -> None: + """Blocking preparation must not strand any unstarted approval owner.""" executed: list[str] = [] - state = {"phase": "pause"} - - def guarded_tool() -> str: - executed.append("guarded") - return "guarded result" + phase = "pause" + local_call = Content.from_function_call( + id="local-occurrence", call_id="local-call", name="local_action", arguments={} + ) + hosted_call = Content.from_function_call( + call_id="hosted-call", + name="hosted_action", + arguments={}, + additional_properties={"server_label": "remote"}, + ) - def first_safe_tool() -> str: - executed.append("first-safe") - return "first safe result" + class PreparationProvider(ContextProvider): + async def before_run( + self, + *, + agent: SupportsAgentRun, + session: AgentSession, + context: SessionContext, + state: dict[str, Any], + ) -> None: + if phase == "blocked" and blocking_reason == "provider_failure": + raise RuntimeError("Provider is temporarily unavailable") + if phase == "blocked" and blocking_reason == "provider_disable": + invocation = getattr(agent, "client") + assert isinstance(invocation, FunctionInvocationLayer) + invocation.function_invocation_configuration["enabled"] = False - def second_safe_tool() -> str: - executed.append("second-safe") - return "second safe result" + def local_action() -> str: + executed.append("local") + return "Local completed" async def stream_fn( - messages: list[Message], - options: dict[str, Any], - **kwargs: Any, + messages: list[Message], options: dict[str, Any], **kwargs: Any ) -> AsyncIterator[ChatResponseUpdate]: - del messages, options, kwargs - if state["phase"] == "pause": + if phase == "pause": yield ChatResponseUpdate( + role="assistant", contents=[ - Content.from_function_call( - id="guarded-occurrence", - call_id="provider-shared", - name="guarded_tool", - arguments="{}", - ), - Content.from_function_call( - id="first-safe-occurrence", - call_id="provider-shared", - name="first_safe_tool", - arguments="{}", - ), - Content.from_function_call( - id="second-safe-occurrence", - call_id="provider-shared", - name="second_safe_tool", - arguments="{}", - ), + Content.from_function_approval_request(id="local-occurrence", function_call=local_call), + Content.from_function_approval_request(id="hosted-approval", function_call=hosted_call), ], - role="assistant", ) return - yield ChatResponseUpdate(contents=[Content.from_text(text="Done.")], role="assistant") + for message in messages: + for content in message.contents: + if content.type == "function_approval_response" and content.id == "hosted-approval": + assert content.approved is True + executed.append("hosted") + yield ChatResponseUpdate(role="assistant", contents=[Content.from_text("Done.")]) + chat_client = streaming_chat_client_stub(stream_fn) agent = Agent( - name="test_agent", - instructions="Test", - client=streaming_chat_client_stub(stream_fn), + client=chat_client, + context_providers=[PreparationProvider("preparation")], tools=[ FunctionTool( - name="guarded_tool", - description="Guarded tool", - func=guarded_tool, - approval_mode="always_require", - ), - FunctionTool(name="first_safe_tool", description="First safe tool", func=first_safe_tool), - FunctionTool(name="second_safe_tool", description="Second safe tool", func=second_safe_tool), + name="local_action", description="Local action", func=local_action, approval_mode="always_require" + ) ], ) - wrapped_agent = AgentFrameworkAgent(agent=agent, require_confirmation=False) + assert isinstance(chat_client, FunctionInvocationLayer) app = FastAPI() - 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( + add_agent_framework_fastapi_endpoint(app, agent, path="/approval") + client = TestClient(app) + pause = client.post( "/approval", json={ - "runId": "run-pause", - "threadId": thread_id, - "messages": [{"role": "user", "content": "Run all three tools"}], + "runId": "pause", + "threadId": "cross-owner", + "messages": [{"role": "user", "content": "Run both actions"}], }, ) - - assert pause_response.status_code == 200 - pause_events = _decode_sse_events(pause_response) - pause_finished = [event for event in pause_events if event.get("type") == "RUN_FINISHED"] - interrupts = _run_finished_interrupts(pause_finished[-1]) - assert [(interrupt["id"], interrupt["toolCallId"]) for interrupt in interrupts] == [ - ("guarded-occurrence", "provider-shared") + finished = [event for event in _decode_sse_events(pause) if event["type"] == "RUN_FINISHED"] + interrupts = _run_finished_interrupts(finished[-1]) + assert {interrupt["id"] for interrupt in interrupts} == {"local-occurrence", "hosted-approval"} + resume = [ + {"interruptId": interrupt["id"], "status": "resolved", "payload": {"approved": True}} + for interrupt in interrupts ] - assert executed == [] - - state["phase"] = "resume" - resume_response = client.post( - "/approval", - json={ - "runId": "run-resume", - "threadId": thread_id, - "messages": [], - "resume": [ - { - "interruptId": "guarded-occurrence", - "status": "resolved", - "payload": {"accepted": True}, - } - ], - }, - ) - - assert resume_response.status_code == 200 - resume_events = _decode_sse_events(resume_response) - assert not [event for event in resume_events if event.get("type") == "RUN_ERROR"] - assert Counter(executed) == {"guarded": 1, "first-safe": 1, "second-safe": 1} - tool_results = [event for event in resume_events if event.get("type") == "TOOL_CALL_RESULT"] - assert Counter(event["content"] for event in tool_results) == { - "guarded result": 1, - "first safe result": 1, - "second safe result": 1, - } - assert {event["toolCallId"] for event in tool_results} == {"provider-shared"} - - occurrences = wrapped_agent._approval_state_store.lifecycle.occurrences_for_thread( - thread_id=_default_approval_thread_id(thread_id) - ) - assert {occurrence.identity.interrupt_id for occurrence in occurrences} == { - "guarded-occurrence", - "first-safe-occurrence", - "second-safe-occurrence", - } - assert {occurrence.identity.call_id for occurrence in occurrences} == {"provider-shared"} - assert len({occurrence.identity.occurrence_id for occurrence in occurrences}) == 3 - 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" + phase = "blocked" + chat_client.function_invocation_configuration["enabled"] = blocking_reason != "disabled" + for run_id in ("blocked", "blocked-again"): + response = client.post( + "/approval", + json={"runId": run_id, "threadId": "cross-owner", "messages": [], "resume": resume}, + ) + assert [event["code"] for event in _decode_sse_events(response) if event["type"] == "RUN_ERROR"] == [ + "RuntimeError" if blocking_reason == "provider_failure" else "APPROVAL_INVOCATION_DISABLED" ] + assert executed == [] - 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] + phase = "resume" + chat_client.function_invocation_configuration["enabled"] = True + for run_id in ("retry", "completed-replay"): + response = client.post( + "/approval", + json={"runId": run_id, "threadId": "cross-owner", "messages": [], "resume": resume}, + ) + assert not [event for event in _decode_sse_events(response) if event["type"] == "RUN_ERROR"] + assert executed == ["local", "hosted"] -def _build_fides_policy_approval_endpoint( - streaming_chat_client_stub: Any, - *, - multiple_guarded: bool = False, -) -> tuple[TestClient, InMemoryAGUIThreadSnapshotStore, SecureAgentConfig, list[str], AgentFrameworkAgent]: - provider_calls = 0 +@pytest.mark.parametrize("failure_point", ["middleware", "tool"]) +async def test_endpoint_fatal_approval_preserves_unstarted_siblings( + streaming_chat_client_stub, failure_point: str +) -> None: + """A denied local approval must not strand untouched local or hosted siblings.""" executed: list[str] = [] + phase = "pause" - def fetch_external() -> str: - return "untrusted content" + class DenyFirst(FunctionMiddleware): + async def process(self, context: Any, call_next: Any) -> None: + if context.function.name == "denied_action" and failure_point == "middleware": + raise MiddlewareFailure("Action denied") + if context.function.name == "allowed_action" and phase == "resume": + # Keep this sibling outside invocation while core cancels the failed parallel batch. + await asyncio.Event().wait() + await call_next() - def guarded_action() -> str: - executed.append("guarded") - return "guarded result" + def denied_action() -> str: + executed.append("denied") + raise MiddlewareFailure("Action failed after starting") - def guarded_action_two() -> str: - executed.append("guarded-two") - return "guarded result two" + def allowed_action() -> str: + executed.append("allowed") + return "Allowed action completed" + + calls = [ + Content.from_function_call(id="denied", call_id="denied-call", name="denied_action", arguments={}), + Content.from_function_call(id="allowed", call_id="allowed-call", name="allowed_action", arguments={}), + Content.from_function_call( + id="hosted", + call_id="hosted-call", + name="hosted_action", + arguments={}, + additional_properties={"server_label": "remote"}, + ), + ] async def stream_fn( - messages: list[Message], - options: dict[str, Any], - **kwargs: Any, + messages: list[Message], options: dict[str, Any], **kwargs: Any ) -> AsyncIterator[ChatResponseUpdate]: - nonlocal provider_calls - del messages, options, kwargs - provider_calls += 1 - if provider_calls == 1: + if phase == "pause": yield ChatResponseUpdate( role="assistant", - contents=[Content.from_function_call(call_id="call-fetch", name="fetch_external", arguments={})], + contents=[ + Content.from_function_approval_request(id=identity, function_call=call) + for identity, call in zip(("denied", "allowed", "hosted"), calls, strict=True) + ], ) - elif provider_calls == 2: - guarded_calls = [Content.from_function_call(call_id="call-guarded", name="guarded_action", arguments={})] - if multiple_guarded: - guarded_calls.append( - Content.from_function_call( - call_id="call-guarded-two", - name="guarded_action_two", - arguments={}, - ) - ) - yield ChatResponseUpdate(role="assistant", contents=guarded_calls) - else: - yield ChatResponseUpdate(role="assistant", contents=[Content.from_text("done")]) + return + for message in messages: + for content in message.contents: + if content.type == "function_approval_response" and content.id == "hosted": + assert content.approved is True + executed.append("hosted") + yield ChatResponseUpdate(role="assistant", contents=[Content.from_text("Done.")]) - security = SecureAgentConfig( - auto_hide_untrusted=False, - allow_untrusted_tools={"fetch_external"}, - approval_on_violation=True, - ) agent = Agent( - name="fides-approval-agent", - instructions="Test FIDES approval cleanup", client=streaming_chat_client_stub(stream_fn), - tools=[fetch_external, guarded_action, guarded_action_two], - context_providers=[security], + tools=[ + FunctionTool( + name="denied_action", description="Denied action", func=denied_action, approval_mode="always_require" + ), + FunctionTool( + name="allowed_action", description="Allowed action", func=allowed_action, approval_mode="always_require" + ), + ], + middleware=[DenyFirst()], ) - snapshot_store = InMemoryAGUIThreadSnapshotStore() app = FastAPI() - wrapped_agent = AgentFrameworkAgent(agent=agent, require_confirmation=False) - add_agent_framework_fastapi_endpoint( - app, - wrapped_agent, - path="/approval", - snapshot_store=snapshot_store, - snapshot_scope_resolver=lambda _request: "tenant-a", - ) - return TestClient(app), snapshot_store, security, executed, wrapped_agent - - -@pytest.mark.parametrize("cancelled", [False, True], ids=["rejected", "cancelled"]) -async def test_endpoint_fides_non_grant_cleans_authenticated_fixed_scope_only( - streaming_chat_client_stub: Any, - cancelled: bool, -) -> None: - """AG-UI lifecycle-authenticated non-grants clean only the owning FIDES occurrence.""" - client, snapshot_store, security, executed, wrapped_agent = _build_fides_policy_approval_endpoint( - streaming_chat_client_stub - ) - thread_id = f"thread-fides-cleanup-{cancelled}" + add_agent_framework_fastapi_endpoint(app, agent, path="/approval") + client = TestClient(app) pause = client.post( + "/approval", + json={"runId": "pause", "threadId": "fatal-batch", "messages": [{"role": "user", "content": "Act"}]}, + ) + finished = [event for event in _decode_sse_events(pause) if event["type"] == "RUN_FINISHED"] + assert {item["id"] for item in _run_finished_interrupts(finished[-1])} == {"denied", "allowed", "hosted"} + phase = "resume" + failed = client.post( "/approval", json={ - "runId": "run-pause", - "threadId": thread_id, - "messages": [{"id": "user-1", "role": "user", "content": "Run guarded action"}], + "runId": "denied", + "threadId": "fatal-batch", + "messages": [], + "resume": [ + {"interruptId": identity, "status": "resolved", "payload": {"approved": True}} + for identity in ("denied", "allowed", "hosted") + ], }, ) - assert pause.status_code == 200 - finished = [event for event in _decode_sse_events(pause) if event.get("type") == "RUN_FINISHED"] - interrupts = _run_finished_interrupts(finished[-1]) - assert len(interrupts) == 1 - approval_id = interrupts[0]["id"] - - snapshot = await snapshot_store.get(scope="tenant-a", thread_id=thread_id) - assert snapshot is not None - assert snapshot.session_state is not None - security_state = snapshot.session_state[security.source_id] - pending = security_state["pending_policy_approvals"] - assert set(pending) == {approval_id} - pending["unrelated-occurrence"] = json.loads(json.dumps(pending[approval_id])) - await snapshot_store.save(scope="tenant-a", thread_id=thread_id, snapshot=snapshot) + assert [event["code"] for event in _decode_sse_events(failed) if event["type"] == "RUN_ERROR"] == [ + "MiddlewareFailure" + ] + failed_effects = ["denied"] if failure_point == "tool" else [] + assert executed == failed_effects - unauthorized = client.post( + denied_retry = client.post( "/approval", json={ - "runId": "run-unauthorized", - "threadId": thread_id, + "runId": "unsafe-retry", + "threadId": "fatal-batch", "messages": [], - "resume": [{"interruptId": "unknown-approval", "status": "cancelled"}], + "resume": [ + {"interruptId": identity, "status": "resolved", "payload": {"approved": True}} + for identity in ("denied", "allowed", "hosted") + ], }, ) - unauthorized_errors = [event for event in _decode_sse_events(unauthorized) if event.get("type") == "RUN_ERROR"] - assert [error["code"] for error in unauthorized_errors] == ["APPROVAL_RESUME_NOT_FOUND"] - unchanged = await snapshot_store.get(scope="tenant-a", thread_id=thread_id) - assert unchanged is not None and unchanged.session_state is not None - assert set(unchanged.session_state[security.source_id]["pending_policy_approvals"]) == { - approval_id, - "unrelated-occurrence", - } + assert [event for event in _decode_sse_events(denied_retry) if event["type"] == "RUN_ERROR"] + assert executed == failed_effects - resume_entry: dict[str, Any] = {"interruptId": approval_id} - if cancelled: - resume_entry["status"] = "cancelled" - else: - resume_entry.update({"status": "resolved", "payload": {"approved": False}}) - completed = client.post( - "/approval", - json={ - "runId": "run-non-grant", - "threadId": thread_id, - "messages": [], - "resume": [resume_entry], - }, + phase = "retry" + for run_id in ("retry-siblings", "replay-completed"): + response = client.post( + "/approval", + json={ + "runId": run_id, + "threadId": "fatal-batch", + "messages": [], + "resume": [ + *([{"interruptId": "denied", "status": "cancelled"}] if failure_point == "middleware" else []), + *[ + {"interruptId": identity, "status": "resolved", "payload": {"approved": True}} + for identity in ("allowed", "hosted") + ], + ], + }, + ) + assert not [event for event in _decode_sse_events(response) if event["type"] == "RUN_ERROR"] + assert executed == [*failed_effects, "allowed", "hosted"] + + +@pytest.mark.parametrize("reuse_call_id", [False, True]) +async def test_endpoint_retained_and_pending_occurrences_resume_independently( + streaming_chat_client_stub, reuse_call_id: bool +) -> None: + """A retained result cannot invalidate a distinct pending approval with the same provider id.""" + phase = "first" + executed: list[str] = [] + + def record(value: str) -> str: + executed.append(value) + return f"Recorded {value}" + + async def stream_fn( + messages: list[Message], options: dict[str, Any], **kwargs: Any + ) -> AsyncIterator[ChatResponseUpdate]: + if phase == "complete": + yield ChatResponseUpdate(role="assistant", contents=[Content.from_text("Done.")]) + return + call = Content.from_function_call( + id=phase, + call_id="shared" if reuse_call_id else f"provider-{phase}", + name="record", + arguments={"value": phase}, + ) + yield ChatResponseUpdate( + role="assistant", contents=[Content.from_function_approval_request(id=phase, function_call=call)] + ) + + agent = Agent( + client=streaming_chat_client_stub(stream_fn), + tools=[FunctionTool(name="record", description="Record", func=record, approval_mode="always_require")], ) - assert completed.status_code == 200 - assert not [event for event in _decode_sse_events(completed) if event.get("type") == "RUN_ERROR"] + app = FastAPI() + add_agent_framework_fastapi_endpoint(app, agent, path="/approval") + client = TestClient(app) + for occurrence_id in ("first", "second"): + phase = occurrence_id + pause = client.post( + "/approval", + json={ + "runId": f"pause-{occurrence_id}", + "threadId": "retained-occurrences", + "messages": [{"role": "user", "content": f"Record {occurrence_id}"}], + }, + ) + finished = [event for event in _decode_sse_events(pause) if event["type"] == "RUN_FINISHED"] + assert [item["id"] for item in _run_finished_interrupts(finished[-1])] == [occurrence_id] + if occurrence_id == "first": + phase = "complete" + completed = client.post( + "/approval", + json={ + "runId": "complete-first", + "threadId": "retained-occurrences", + "messages": [], + "resume": [{"interruptId": "first", "status": "resolved", "payload": {"approved": True}}], + }, + ) + assert not [event for event in _decode_sse_events(completed) if event["type"] == "RUN_ERROR"] + assert executed == ["first"] - cleaned = await snapshot_store.get(scope="tenant-a", thread_id=thread_id) - assert cleaned is not None and cleaned.session_state is not None - assert set(cleaned.session_state[security.source_id]["pending_policy_approvals"]) == {"unrelated-occurrence"} - assert executed == [] + phase = "complete" + for run_id in ("complete-both", "repeat-completed"): + response = client.post( + "/approval", + json={ + "runId": run_id, + "threadId": "retained-occurrences", + "messages": [], + "resume": [ + {"interruptId": identity, "status": "resolved", "payload": {"approved": True}} + for identity in ("first", "second") + ], + }, + ) + events = _decode_sse_events(response) + assert not [event for event in events if event["type"] == "RUN_ERROR"] + assert executed == ["first", "second"] + assert [event["content"] for event in events if event["type"] == "TOOL_CALL_RESULT"] == [ + "Recorded first", + "Recorded second", + ] -async def test_endpoint_fides_mixed_cancel_and_reject_clean_each_authenticated_occurrence( - streaming_chat_client_stub: Any, -) -> None: - """A mixed AG-UI resume notifies FIDES for both cancelled and rejected occurrences.""" - client, snapshot_store, security, executed, _ = _build_fides_policy_approval_endpoint( - streaming_chat_client_stub, - multiple_guarded=True, +async def test_endpoint_agent_disabled_mixed_resume_preserves_terminal_progress_for_retry() -> None: + """A blocked mixed resume keeps grants retryable without reviving terminal siblings.""" + executed: list[str] = [] + observed_non_grants: list[tuple[str | None, bool]] = [] + + class ObserveNonGrants(FunctionMiddleware): + async def process(self, context: Any, call_next: Any) -> None: + del context + await call_next() + + def _on_approval_responses( + self, + responses: Sequence[Content], + *, + session: AgentSession | None, + ) -> None: + del session + observed_non_grants.extend( + ( + response.function_call.call_id if response.function_call else None, + response.additional_properties.get("cancelled") is True, + ) + for response in responses + ) + + def record_city(city: str) -> str: + executed.append(city) + return f"Recorded {city}" + + tool = FunctionTool( + name="record_city", + description="Record a city", + func=record_city, + approval_mode="always_require", ) - thread_id = "thread-fides-mixed-non-grants" + approval_requests = [ + Content.from_function_approval_request( + id=f"approval-{city.lower()}", + function_call=Content.from_function_call( + call_id=f"call-{city.lower()}", + name="record_city", + arguments={"city": city}, + ), + ) + for city in ("Seattle", "Portland", "Tokyo") + ] + client_config = Mock(function_invocation_configuration={"enabled": True}) + + class ObservedAgent(StubAgent): + middleware = [ObserveNonGrants()] + + agent = ObservedAgent( + updates=[AgentResponseUpdate(contents=approval_requests, role="assistant")], + default_options={"tools": [tool]}, + client=client_config, + ) + wrapped_agent = AgentFrameworkAgent(agent=agent, require_confirmation=False) + snapshot_store = InMemoryAGUIThreadSnapshotStore() + app = FastAPI() + add_agent_framework_fastapi_endpoint( + app, + wrapped_agent, + path="/approval", + snapshot_store=snapshot_store, + snapshot_scope_resolver=lambda _request: "tenant-a", + ) + client = TestClient(app) + thread_id = "thread-disabled-mixed-decisions" pause = client.post( "/approval", json={ "runId": "run-pause", "threadId": thread_id, - "messages": [{"id": "user-1", "role": "user", "content": "Run guarded actions"}], + "messages": [{"role": "user", "content": "Record three cities"}], }, ) - interrupts = _run_finished_interrupts( - [event for event in _decode_sse_events(pause) if event.get("type") == "RUN_FINISHED"][-1] - ) - assert len(interrupts) == 2 - snapshot = await snapshot_store.get(scope="tenant-a", thread_id=thread_id) - assert snapshot is not None and snapshot.session_state is not None - pending = snapshot.session_state[security.source_id]["pending_policy_approvals"] - interrupt_ids = [interrupt["id"] for interrupt in interrupts] - assert set(pending) == set(interrupt_ids) - pending["unrelated-occurrence"] = json.loads(json.dumps(pending[interrupt_ids[0]])) - await snapshot_store.save(scope="tenant-a", thread_id=thread_id, snapshot=snapshot) + pause_finished = [event for event in _decode_sse_events(pause) if event.get("type") == "RUN_FINISHED"] + assert {interrupt["id"] for interrupt in _run_finished_interrupts(pause_finished[-1])} == { + "call-seattle", + "call-portland", + "call-tokyo", + } - mixed = client.post( - "/approval", - json={ - "runId": "run-mixed", - "threadId": thread_id, - "messages": [], - "resume": [ - {"interruptId": interrupt_ids[0], "status": "cancelled"}, - { - "interruptId": interrupt_ids[1], - "status": "resolved", - "payload": {"approved": False}, - }, - ], + agent.updates = [AgentResponseUpdate(contents=[Content.from_text(text="Done.")], role="assistant")] + mixed_resume = [ + { + "interruptId": "call-seattle", + "status": "resolved", + "payload": {"approved": True}, }, - ) + { + "interruptId": "call-portland", + "status": "resolved", + "payload": {"approved": False}, + }, + {"interruptId": "call-tokyo", "status": "cancelled"}, + ] + client_config.function_invocation_configuration["enabled"] = False - assert mixed.status_code == 200 - assert not [event for event in _decode_sse_events(mixed) if event.get("type") == "RUN_ERROR"] - assert executed == [] - cleaned = await snapshot_store.get(scope="tenant-a", thread_id=thread_id) - assert cleaned is not None and cleaned.session_state is not None - assert set(cleaned.session_state[security.source_id]["pending_policy_approvals"]) == {"unrelated-occurrence"} + for run_id in ("run-blocked", "run-blocked-again"): + blocked = client.post( + "/approval", + json={"runId": run_id, "threadId": thread_id, "messages": [], "resume": mixed_resume}, + ) + blocked_events = _decode_sse_events(blocked) + assert [event["code"] for event in blocked_events if event.get("type") == "RUN_ERROR"] == [ + "APPROVAL_INVOCATION_DISABLED" + ] + assert not [event for event in blocked_events if event.get("type") == "TOOL_CALL_RESULT"] + assert executed == [] + assert set(observed_non_grants) == { + ("call-portland", False), + ("call-tokyo", True), + } + snapshot = await snapshot_store.get(scope="tenant-a", thread_id=thread_id) + assert snapshot is not None and snapshot.interrupt is not None + assert {interrupt["id"] for interrupt in snapshot.interrupt} == {"call-seattle"} -async def test_endpoint_fides_malformed_legacy_non_grant_does_not_clean_policy_authority( - streaming_chat_client_stub: Any, -) -> None: - """A non-boolean legacy control cannot trigger FIDES cleanup.""" - client, snapshot_store, security, executed, wrapped_agent = _build_fides_policy_approval_endpoint( - streaming_chat_client_stub - ) - thread_id = "thread-fides-malformed-legacy" - pause = client.post( + client_config.function_invocation_configuration["enabled"] = True + retry = client.post( "/approval", - json={ - "runId": "run-pause", - "threadId": thread_id, - "messages": [{"id": "user-1", "role": "user", "content": "Run guarded action"}], - }, + json={"runId": "run-retry", "threadId": thread_id, "messages": [], "resume": mixed_resume}, ) - interrupt = _run_finished_interrupts( - [event for event in _decode_sse_events(pause) if event.get("type") == "RUN_FINISHED"][-1] - )[0] - approval_id = interrupt["id"] + retry_events = _decode_sse_events(retry) + assert not [event for event in retry_events if event.get("type") == "RUN_ERROR"] + assert [ + (event["toolCallId"], event["content"]) for event in retry_events if event.get("type") == "TOOL_CALL_RESULT" + ] == [("call-seattle", "Recorded Seattle")] + assert executed == ["Seattle"] - malformed = client.post( - "/approval", - json={ - "runId": "run-malformed", - "threadId": thread_id, - "messages": [ - { - "id": "malformed-response", - "role": "user", - "function_approvals": [ - { - "id": approval_id, - "call_id": interrupt["toolCallId"], - "name": "guarded_action", - "approved": "false", - "arguments": {}, - } - ], - } - ], - }, - ) - assert malformed.status_code == 200 - assert executed == [] - snapshot = await snapshot_store.get(scope="tenant-a", thread_id=thread_id) - assert snapshot is not None and snapshot.session_state is not None - assert set(snapshot.session_state[security.source_id]["pending_policy_approvals"]) == {approval_id} +async def test_endpoint_agent_mixed_retry_reprojects_completed_sibling_without_reexecution() -> None: + """A completed sibling is replayed while only the unfinished approval executes.""" + executed: list[str] = [] + def record_city(city: str) -> str: + executed.append(city) + return f"Recorded {city}" -async def test_endpoint_fides_replacement_rotates_lifecycle_generation( - streaming_chat_client_stub: Any, -) -> None: - """An expired FIDES grant becomes a visible fresh AG-UI interrupt before execution.""" - client, snapshot_store, security, executed, wrapped_agent = _build_fides_policy_approval_endpoint( - streaming_chat_client_stub + tool = FunctionTool(name="record_city", description="Record a city", func=record_city) + agent = StubAgent( + updates=[AgentResponseUpdate(contents=[Content.from_text(text="Done.")], role="assistant")], + default_options={"tools": [tool]}, ) - thread_id = "thread-fides-replacement-generation" - pause = client.post( - "/approval", - json={ - "runId": "run-pause", - "threadId": thread_id, - "messages": [{"id": "user-1", "role": "user", "content": "Run guarded action"}], - }, + wrapped_agent = AgentFrameworkAgent(agent=agent, require_confirmation=False) + lifecycle = wrapped_agent._approval_state_store.lifecycle + approval_thread_id = _default_approval_thread_id("thread-completed-sibling") + completed = lifecycle.register( + owner=ApprovalExecutionOwner.LOCAL, + thread_id=approval_thread_id, + interrupt_id="approval-seattle", + call_id="call-seattle", + name="record_city", + arguments='{"city":"Seattle"}', ) - original_interrupt = _run_finished_interrupts( - [event for event in _decode_sse_events(pause) if event.get("type") == "RUN_FINISHED"][-1] - )[0] - original_id = original_interrupt["id"] - snapshot = await snapshot_store.get(scope="tenant-a", thread_id=thread_id) - assert snapshot is not None and snapshot.session_state is not None - snapshot.session_state[security.source_id]["pending_policy_approvals"][original_id]["created_at"] = 0.0 - assert snapshot.interrupt is not None - snapshot.interrupt.append( - { - "id": "other-input", - "reason": "input_required", - "message": "Choose another value.", - } + completed_intent = lifecycle.claim( + thread_id=approval_thread_id, + decision=ResumeDecision( + interrupt_id="approval-seattle", + accepted=True, + arguments='{"city":"Seattle"}', + name="record_city", + original_arguments='{"city":"Seattle"}', + ), ) - await snapshot_store.save(scope="tenant-a", thread_id=thread_id, snapshot=snapshot) + lifecycle.begin_execution(completed_intent, owner=ApprovalExecutionOwner.LOCAL) + lifecycle.settle( + completed_intent, + [Content.from_function_result(call_id=completed.identity.call_id, result="Recorded Seattle")], + ) + lifecycle.register( + owner=ApprovalExecutionOwner.LOCAL, + thread_id=approval_thread_id, + interrupt_id="approval-portland", + call_id="call-portland", + name="record_city", + arguments='{"city":"Portland"}', + ) + app = FastAPI() + add_agent_framework_fastapi_endpoint(app, wrapped_agent, path="/approval") + client = TestClient(app) + resume = [ + { + "interruptId": "approval-seattle", + "status": "resolved", + "payload": {"approved": True}, + }, + { + "interruptId": "approval-portland", + "status": "resolved", + "payload": {"approved": True}, + }, + ] - stale = client.post( + response = client.post( "/approval", json={ - "runId": "run-stale", - "threadId": thread_id, + "runId": "run-completed-sibling", + "threadId": "thread-completed-sibling", "messages": [], - "resume": [{"interruptId": original_id, "status": "resolved", "payload": {"approved": True}}], + "resume": resume, }, ) - assert stale.status_code == 200 - replacement_interrupts = _run_finished_interrupts( - [event for event in _decode_sse_events(stale) if event.get("type") == "RUN_FINISHED"][-1] - ) - assert len(replacement_interrupts) == 1 - replacement_id = replacement_interrupts[0]["id"] - assert replacement_id != original_id - assert replacement_interrupts[0]["toolCallId"] == "call-guarded" - assert executed == [] - replaced_snapshot = await snapshot_store.get(scope="tenant-a", thread_id=thread_id) - assert replaced_snapshot is not None - assert replaced_snapshot.interrupt is not None - assert {interrupt["id"] for interrupt in replaced_snapshot.interrupt} == {replacement_id, "other-input"} - approval_state = wrapped_agent._approval_state_store.get_tool_approval_state( - approval_state_thread_id(scope="tenant-a", thread_id=thread_id) - ) - assert approval_state is not None - pending_snapshots = approval_state["pending_approval_requests"] - assert [(item["id"], item["function_call"]["id"]) for item in pending_snapshots] == [(replacement_id, original_id)] + events = _decode_sse_events(response) + assert not [event for event in events if event.get("type") == "RUN_ERROR"] + assert [(event["toolCallId"], event["content"]) for event in events if event.get("type") == "TOOL_CALL_RESULT"] == [ + ("call-seattle", "Recorded Seattle"), + ("call-portland", "Recorded Portland"), + ] + assert executed == ["Portland"] - legacy_stale = client.post( - "/approval", - json={ - "runId": "run-legacy-stale", - "threadId": thread_id, - "messages": [ - { - "role": "tool", - "toolCallId": "call-guarded", - "content": json.dumps({"accepted": True}), - } - ], - }, - ) - legacy_errors = [event for event in _decode_sse_events(legacy_stale) if event.get("type") == "RUN_ERROR"] - assert [error["code"] for error in legacy_errors] == ["APPROVAL_RESUME_INVALID"] - assert executed == [] - stale_again = client.post( +async def test_endpoint_agent_approval_resume_releases_already_approved_sibling(streaming_chat_client_stub): + """Resuming a visible approval should also complete never-require siblings from the same batch.""" + client, executed, messages_received, state = _build_mixed_approval_batch_endpoint(streaming_chat_client_stub) + + pause_response = client.post( "/approval", json={ - "runId": "run-stale-again", - "threadId": thread_id, - "messages": [], - "resume": [{"interruptId": original_id, "status": "resolved", "payload": {"approved": True}}], + "runId": "run-pause", + "threadId": "thread-mixed-batch", + "messages": [{"role": "user", "content": "Run both tools"}], }, ) - stale_errors = [event for event in _decode_sse_events(stale_again) if event.get("type") == "RUN_ERROR"] - assert [error["code"] for error in stale_errors] == ["APPROVAL_RESUME_NOT_FOUND"] - assert executed == [] - approved = client.post( + assert pause_response.status_code == 200 + pause_events = _decode_sse_events(pause_response) + pause_finished = [event for event in pause_events if event.get("type") == "RUN_FINISHED"] + interrupts = _run_finished_interrupts(pause_finished[-1]) + assert len(interrupts) == 1 + approval_id = interrupts[0]["id"] + assert approval_id.startswith("af-call-") + assert interrupts[0]["toolCallId"] == "call_sensitive" + assert not [event for event in pause_events if event.get("type") == "TOOL_CALL_RESULT"] + + state["phase"] = "resume" + resume_response = client.post( "/approval", json={ - "runId": "run-fresh", - "threadId": thread_id, + "runId": "run-resume", + "threadId": "thread-mixed-batch", "messages": [], - "resume": [{"interruptId": replacement_id, "status": "resolved", "payload": {"approved": True}}], + "resume": [{"interruptId": approval_id, "status": "resolved", "payload": {"accepted": True}}], }, ) - results = [event for event in _decode_sse_events(approved) if event.get("type") == "TOOL_CALL_RESULT"] - assert [(event["toolCallId"], event["content"]) for event in results] == [("call-guarded", "guarded result")] - assert executed == ["guarded"] + + assert resume_response.status_code == 200 + resume_events = _decode_sse_events(resume_response) + tool_results = [event for event in resume_events if event.get("type") == "TOOL_CALL_RESULT"] + assert [(event["toolCallId"], event["content"]) for event in tool_results] == [ + ("call_sensitive", "Sensitive action in Seattle"), + ("call_weather", "Weather in Seattle"), + ] + assert executed == ["sensitive:Seattle", "weather:Seattle"] + assert not [ + event + for event in resume_events + if event.get("type") == "TOOL_CALL_START" and event.get("toolCallId") == "call_weather" + ] + replayed_results = [ + content for message in messages_received for content in message.contents if content.type == "function_result" + ] + replayed_call_ids = [content.call_id for content in replayed_results if content.call_id is not None] + assert sorted(replayed_call_ids) == ["call_sensitive", "call_weather"] -async def test_endpoint_fides_approval_uses_lifecycle_bound_request_generation( +def _build_reused_provider_call_id_endpoint( streaming_chat_client_stub: Any, -) -> None: - """AG-UI preserves trusted request generation while rebinding to occurrence identity.""" - client, snapshot_store, security, executed, wrapped_agent = _build_fides_policy_approval_endpoint( - streaming_chat_client_stub - ) - thread_id = "thread-fides-approved-generation" - pause = client.post( - "/approval", - json={ - "runId": "run-pause", - "threadId": thread_id, - "messages": [{"id": "user-1", "role": "user", "content": "Run guarded action"}], - }, - ) - finished = [event for event in _decode_sse_events(pause) if event.get("type") == "RUN_FINISHED"] - approval_id = _run_finished_interrupts(finished[-1])[0]["id"] + *, + 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"} - approved = client.post( - "/approval", - json={ - "runId": "run-approved", - "threadId": thread_id, - "messages": [], - "resume": [{"interruptId": approval_id, "status": "resolved", "payload": {"approved": True}}], - }, - ) + def guarded_tool() -> str: + executed.append("guarded") + return "guarded result" - assert approved.status_code == 200 - results = [event for event in _decode_sse_events(approved) if event.get("type") == "TOOL_CALL_RESULT"] - assert [(event["toolCallId"], event["content"]) for event in results] == [("call-guarded", "guarded result")] - assert executed == ["guarded"] - snapshot = await snapshot_store.get(scope="tenant-a", thread_id=thread_id) - assert snapshot is not None and snapshot.session_state is not None - assert snapshot.session_state[security.source_id]["pending_policy_approvals"] == {} + def first_safe_tool() -> str: + executed.append("first-safe") + return "first safe result" + def second_safe_tool() -> str: + executed.append("second-safe") + return "second safe result" -async def test_endpoint_agent_approval_resume_persists_replayable_tool_results(streaming_chat_client_stub): - """Approved batches should hydrate with real results under original tool call ids.""" - client, executed, messages_received, state = _build_mixed_approval_batch_endpoint( - streaming_chat_client_stub, - snapshot_store=InMemoryAGUIThreadSnapshotStore(), + async def stream_fn( + messages: list[Message], + options: dict[str, Any], + **kwargs: Any, + ) -> AsyncIterator[ChatResponseUpdate]: + del messages, options, kwargs + if state["phase"] == "pause": + yield ChatResponseUpdate( + contents=[ + Content.from_function_call( + id="guarded-occurrence", + call_id="provider-shared", + name="guarded_tool", + arguments="{}", + ), + Content.from_function_call( + id="first-safe-occurrence", + call_id="provider-shared", + name="first_safe_tool", + arguments="{}", + ), + Content.from_function_call( + id="second-safe-occurrence", + call_id="provider-shared", + name="second_safe_tool", + 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="guarded_tool", + description="Guarded tool", + func=guarded_tool, + approval_mode="always_require", + ), + FunctionTool(name="first_safe_tool", description="First safe tool", func=first_safe_tool), + FunctionTool(name="second_safe_tool", description="Second safe tool", func=second_safe_tool), + ], + ) + wrapped_agent = AgentFrameworkAgent(agent=agent, require_confirmation=False) + app = FastAPI() + 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( "/approval", json={ "runId": "run-pause", - "threadId": "thread-mixed-replay", - "messages": [{"id": "user-1", "role": "user", "content": "Run both tools"}], + "threadId": thread_id, + "messages": [{"role": "user", "content": "Run all three tools"}], }, ) + assert pause_response.status_code == 200 - pause_finished = [event for event in _decode_sse_events(pause_response) if event.get("type") == "RUN_FINISHED"] + pause_events = _decode_sse_events(pause_response) + pause_finished = [event for event in pause_events if event.get("type") == "RUN_FINISHED"] interrupts = _run_finished_interrupts(pause_finished[-1]) - assert len(interrupts) == 1 - approval_id = interrupts[0]["id"] - assert approval_id.startswith("af-call-") - assert interrupts[0]["toolCallId"] == "call_sensitive" + assert [(interrupt["id"], interrupt["toolCallId"]) for interrupt in interrupts] == [ + ("guarded-occurrence", "provider-shared") + ] + assert executed == [] state["phase"] = "resume" resume_response = client.post( "/approval", json={ "runId": "run-resume", - "threadId": "thread-mixed-replay", + "threadId": thread_id, "messages": [], - "resume": [{"interruptId": approval_id, "status": "resolved", "payload": {"accepted": True}}], + "resume": [ + { + "interruptId": "guarded-occurrence", + "status": "resolved", + "payload": {"accepted": True}, + } + ], }, ) assert resume_response.status_code == 200 resume_events = _decode_sse_events(resume_response) - live_results = [ - (event["toolCallId"], event["content"]) for event in resume_events if event.get("type") == "TOOL_CALL_RESULT" - ] - assert live_results == [ - ("call_sensitive", "Sensitive action in Seattle"), - ("call_weather", "Weather in Seattle"), - ] - assert Counter(call_id for call_id, _ in live_results) == {"call_sensitive": 1, "call_weather": 1} - assert executed == ["sensitive:Seattle", "weather:Seattle"] + assert not [event for event in resume_events if event.get("type") == "RUN_ERROR"] + assert Counter(executed) == {"guarded": 1, "first-safe": 1, "second-safe": 1} + tool_results = [event for event in resume_events if event.get("type") == "TOOL_CALL_RESULT"] + assert Counter(event["content"] for event in tool_results) == { + "guarded result": 1, + "first safe result": 1, + "second safe result": 1, + } + assert {event["toolCallId"] for event in tool_results} == {"provider-shared"} + + occurrences = wrapped_agent._approval_state_store.lifecycle.occurrences_for_thread( + thread_id=_default_approval_thread_id(thread_id) + ) + assert {occurrence.identity.interrupt_id for occurrence in occurrences} == { + "guarded-occurrence", + "first-safe-occurrence", + "second-safe-occurrence", + } + assert {occurrence.identity.call_id for occurrence in occurrences} == {"provider-shared"} + assert len({occurrence.identity.occurrence_id for occurrence in occurrences}) == 3 + 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, + *, + multiple_guarded: bool = False, +) -> tuple[TestClient, InMemoryAGUIThreadSnapshotStore, SecureAgentConfig, list[str], AgentFrameworkAgent]: + provider_calls = 0 + executed: list[str] = [] + + def fetch_external() -> str: + return "untrusted content" + + def guarded_action() -> str: + executed.append("guarded") + return "guarded result" + + def guarded_action_two() -> str: + executed.append("guarded-two") + return "guarded result two" + + async def stream_fn( + messages: list[Message], + options: dict[str, Any], + **kwargs: Any, + ) -> AsyncIterator[ChatResponseUpdate]: + nonlocal provider_calls + del messages, options, kwargs + provider_calls += 1 + if provider_calls == 1: + yield ChatResponseUpdate( + role="assistant", + contents=[Content.from_function_call(call_id="call-fetch", name="fetch_external", arguments={})], + ) + elif provider_calls == 2: + guarded_calls = [Content.from_function_call(call_id="call-guarded", name="guarded_action", arguments={})] + if multiple_guarded: + guarded_calls.append( + Content.from_function_call( + call_id="call-guarded-two", + name="guarded_action_two", + arguments={}, + ) + ) + yield ChatResponseUpdate(role="assistant", contents=guarded_calls) + else: + yield ChatResponseUpdate(role="assistant", contents=[Content.from_text("done")]) + + security = SecureAgentConfig( + auto_hide_untrusted=False, + allow_untrusted_tools={"fetch_external"}, + approval_on_violation=True, + ) + agent = Agent( + name="fides-approval-agent", + instructions="Test FIDES approval cleanup", + client=streaming_chat_client_stub(stream_fn), + tools=[fetch_external, guarded_action, guarded_action_two], + context_providers=[security], + ) + snapshot_store = InMemoryAGUIThreadSnapshotStore() + app = FastAPI() + wrapped_agent = AgentFrameworkAgent(agent=agent, require_confirmation=False) + add_agent_framework_fastapi_endpoint( + app, + wrapped_agent, + path="/approval", + snapshot_store=snapshot_store, + snapshot_scope_resolver=lambda _request: "tenant-a", + ) + return TestClient(app), snapshot_store, security, executed, wrapped_agent + + +@pytest.mark.parametrize("cancelled", [False, True], ids=["rejected", "cancelled"]) +async def test_endpoint_fides_non_grant_cleans_authenticated_fixed_scope_only( + streaming_chat_client_stub: Any, + cancelled: bool, +) -> None: + """AG-UI lifecycle-authenticated non-grants clean only the owning FIDES occurrence.""" + client, snapshot_store, security, executed, wrapped_agent = _build_fides_policy_approval_endpoint( + streaming_chat_client_stub + ) + thread_id = f"thread-fides-cleanup-{cancelled}" + pause = client.post( + "/approval", + json={ + "runId": "run-pause", + "threadId": thread_id, + "messages": [{"id": "user-1", "role": "user", "content": "Run guarded action"}], + }, + ) + assert pause.status_code == 200 + finished = [event for event in _decode_sse_events(pause) if event.get("type") == "RUN_FINISHED"] + interrupts = _run_finished_interrupts(finished[-1]) + assert len(interrupts) == 1 + approval_id = interrupts[0]["id"] + + snapshot = await snapshot_store.get(scope="tenant-a", thread_id=thread_id) + assert snapshot is not None + assert snapshot.session_state is not None + security_state = snapshot.session_state[security.source_id] + pending = security_state["pending_policy_approvals"] + assert set(pending) == {approval_id} + pending["unrelated-occurrence"] = json.loads(json.dumps(pending[approval_id])) + await snapshot_store.save(scope="tenant-a", thread_id=thread_id, snapshot=snapshot) + + unauthorized = client.post( + "/approval", + json={ + "runId": "run-unauthorized", + "threadId": thread_id, + "messages": [], + "resume": [{"interruptId": "unknown-approval", "status": "cancelled"}], + }, + ) + unauthorized_errors = [event for event in _decode_sse_events(unauthorized) if event.get("type") == "RUN_ERROR"] + assert [error["code"] for error in unauthorized_errors] == ["APPROVAL_RESUME_NOT_FOUND"] + unchanged = await snapshot_store.get(scope="tenant-a", thread_id=thread_id) + assert unchanged is not None and unchanged.session_state is not None + assert set(unchanged.session_state[security.source_id]["pending_policy_approvals"]) == { + approval_id, + "unrelated-occurrence", + } + + resume_entry: dict[str, Any] = {"interruptId": approval_id} + if cancelled: + resume_entry["status"] = "cancelled" + else: + resume_entry.update({"status": "resolved", "payload": {"approved": False}}) + completed = client.post( + "/approval", + json={ + "runId": "run-non-grant", + "threadId": thread_id, + "messages": [], + "resume": [resume_entry], + }, + ) + assert completed.status_code == 200 + assert not [event for event in _decode_sse_events(completed) if event.get("type") == "RUN_ERROR"] + + cleaned = await snapshot_store.get(scope="tenant-a", thread_id=thread_id) + assert cleaned is not None and cleaned.session_state is not None + assert set(cleaned.session_state[security.source_id]["pending_policy_approvals"]) == {"unrelated-occurrence"} + assert executed == [] + + +async def test_endpoint_fides_mixed_cancel_and_reject_clean_each_authenticated_occurrence( + streaming_chat_client_stub: Any, +) -> None: + """A mixed AG-UI resume notifies FIDES for both cancelled and rejected occurrences.""" + client, snapshot_store, security, executed, _ = _build_fides_policy_approval_endpoint( + streaming_chat_client_stub, + multiple_guarded=True, + ) + thread_id = "thread-fides-mixed-non-grants" + pause = client.post( + "/approval", + json={ + "runId": "run-pause", + "threadId": thread_id, + "messages": [{"id": "user-1", "role": "user", "content": "Run guarded actions"}], + }, + ) + interrupts = _run_finished_interrupts( + [event for event in _decode_sse_events(pause) if event.get("type") == "RUN_FINISHED"][-1] + ) + assert len(interrupts) == 2 + snapshot = await snapshot_store.get(scope="tenant-a", thread_id=thread_id) + assert snapshot is not None and snapshot.session_state is not None + pending = snapshot.session_state[security.source_id]["pending_policy_approvals"] + interrupt_ids = [interrupt["id"] for interrupt in interrupts] + assert set(pending) == set(interrupt_ids) + pending["unrelated-occurrence"] = json.loads(json.dumps(pending[interrupt_ids[0]])) + await snapshot_store.save(scope="tenant-a", thread_id=thread_id, snapshot=snapshot) + + mixed = client.post( + "/approval", + json={ + "runId": "run-mixed", + "threadId": thread_id, + "messages": [], + "resume": [ + {"interruptId": interrupt_ids[0], "status": "cancelled"}, + { + "interruptId": interrupt_ids[1], + "status": "resolved", + "payload": {"approved": False}, + }, + ], + }, + ) + + assert mixed.status_code == 200 + assert not [event for event in _decode_sse_events(mixed) if event.get("type") == "RUN_ERROR"] + assert executed == [] + cleaned = await snapshot_store.get(scope="tenant-a", thread_id=thread_id) + assert cleaned is not None and cleaned.session_state is not None + assert set(cleaned.session_state[security.source_id]["pending_policy_approvals"]) == {"unrelated-occurrence"} + + +async def test_endpoint_fides_malformed_legacy_non_grant_does_not_clean_policy_authority( + streaming_chat_client_stub: Any, +) -> None: + """A non-boolean legacy control cannot trigger FIDES cleanup.""" + client, snapshot_store, security, executed, wrapped_agent = _build_fides_policy_approval_endpoint( + streaming_chat_client_stub + ) + thread_id = "thread-fides-malformed-legacy" + pause = client.post( + "/approval", + json={ + "runId": "run-pause", + "threadId": thread_id, + "messages": [{"id": "user-1", "role": "user", "content": "Run guarded action"}], + }, + ) + interrupt = _run_finished_interrupts( + [event for event in _decode_sse_events(pause) if event.get("type") == "RUN_FINISHED"][-1] + )[0] + approval_id = interrupt["id"] + + malformed = client.post( + "/approval", + json={ + "runId": "run-malformed", + "threadId": thread_id, + "messages": [ + { + "id": "malformed-response", + "role": "user", + "function_approvals": [ + { + "id": approval_id, + "call_id": interrupt["toolCallId"], + "name": "guarded_action", + "approved": "false", + "arguments": {}, + } + ], + } + ], + }, + ) + + assert malformed.status_code == 200 + assert executed == [] + snapshot = await snapshot_store.get(scope="tenant-a", thread_id=thread_id) + assert snapshot is not None and snapshot.session_state is not None + assert set(snapshot.session_state[security.source_id]["pending_policy_approvals"]) == {approval_id} + + +async def test_endpoint_fides_replacement_rotates_lifecycle_generation( + streaming_chat_client_stub: Any, +) -> None: + """An expired FIDES grant becomes a visible fresh AG-UI interrupt before execution.""" + client, snapshot_store, security, executed, wrapped_agent = _build_fides_policy_approval_endpoint( + streaming_chat_client_stub + ) + thread_id = "thread-fides-replacement-generation" + pause = client.post( + "/approval", + json={ + "runId": "run-pause", + "threadId": thread_id, + "messages": [{"id": "user-1", "role": "user", "content": "Run guarded action"}], + }, + ) + original_interrupt = _run_finished_interrupts( + [event for event in _decode_sse_events(pause) if event.get("type") == "RUN_FINISHED"][-1] + )[0] + original_id = original_interrupt["id"] + snapshot = await snapshot_store.get(scope="tenant-a", thread_id=thread_id) + assert snapshot is not None and snapshot.session_state is not None + snapshot.session_state[security.source_id]["pending_policy_approvals"][original_id]["created_at"] = 0.0 + assert snapshot.interrupt is not None + snapshot.interrupt.append( + { + "id": "other-input", + "reason": "input_required", + "message": "Choose another value.", + } + ) + await snapshot_store.save(scope="tenant-a", thread_id=thread_id, snapshot=snapshot) + + stale = client.post( + "/approval", + json={ + "runId": "run-stale", + "threadId": thread_id, + "messages": [], + "resume": [{"interruptId": original_id, "status": "resolved", "payload": {"approved": True}}], + }, + ) + + assert stale.status_code == 200 + replacement_interrupts = _run_finished_interrupts( + [event for event in _decode_sse_events(stale) if event.get("type") == "RUN_FINISHED"][-1] + ) + assert len(replacement_interrupts) == 1 + replacement_id = replacement_interrupts[0]["id"] + assert replacement_id != original_id + assert replacement_interrupts[0]["toolCallId"] == "call-guarded" + assert executed == [] + replaced_snapshot = await snapshot_store.get(scope="tenant-a", thread_id=thread_id) + assert replaced_snapshot is not None + assert replaced_snapshot.interrupt is not None + assert {interrupt["id"] for interrupt in replaced_snapshot.interrupt} == {replacement_id, "other-input"} + approval_state = wrapped_agent._approval_state_store.get_tool_approval_state( + approval_state_thread_id(scope="tenant-a", thread_id=thread_id) + ) + assert approval_state is not None + pending_snapshots = approval_state["pending_approval_requests"] + assert [(item["id"], item["function_call"]["id"]) for item in pending_snapshots] == [(replacement_id, original_id)] + + legacy_stale = client.post( + "/approval", + json={ + "runId": "run-legacy-stale", + "threadId": thread_id, + "messages": [ + { + "role": "tool", + "toolCallId": "call-guarded", + "content": json.dumps({"accepted": True}), + } + ], + }, + ) + legacy_errors = [event for event in _decode_sse_events(legacy_stale) if event.get("type") == "RUN_ERROR"] + assert [error["code"] for error in legacy_errors] == ["APPROVAL_RESUME_INVALID"] + assert executed == [] + + stale_again = client.post( + "/approval", + json={ + "runId": "run-stale-again", + "threadId": thread_id, + "messages": [], + "resume": [{"interruptId": original_id, "status": "resolved", "payload": {"approved": True}}], + }, + ) + stale_errors = [event for event in _decode_sse_events(stale_again) if event.get("type") == "RUN_ERROR"] + assert [error["code"] for error in stale_errors] == ["APPROVAL_RESUME_NOT_FOUND"] + assert executed == [] + + approved = client.post( + "/approval", + json={ + "runId": "run-fresh", + "threadId": thread_id, + "messages": [], + "resume": [{"interruptId": replacement_id, "status": "resolved", "payload": {"approved": True}}], + }, + ) + results = [event for event in _decode_sse_events(approved) if event.get("type") == "TOOL_CALL_RESULT"] + assert [(event["toolCallId"], event["content"]) for event in results] == [("call-guarded", "guarded result")] + assert executed == ["guarded"] + + +async def test_endpoint_fides_approval_uses_lifecycle_bound_request_generation( + streaming_chat_client_stub: Any, +) -> None: + """AG-UI preserves trusted request generation while rebinding to occurrence identity.""" + client, snapshot_store, security, executed, wrapped_agent = _build_fides_policy_approval_endpoint( + streaming_chat_client_stub + ) + thread_id = "thread-fides-approved-generation" + pause = client.post( + "/approval", + json={ + "runId": "run-pause", + "threadId": thread_id, + "messages": [{"id": "user-1", "role": "user", "content": "Run guarded action"}], + }, + ) + finished = [event for event in _decode_sse_events(pause) if event.get("type") == "RUN_FINISHED"] + approval_id = _run_finished_interrupts(finished[-1])[0]["id"] + + approved = client.post( + "/approval", + json={ + "runId": "run-approved", + "threadId": thread_id, + "messages": [], + "resume": [{"interruptId": approval_id, "status": "resolved", "payload": {"approved": True}}], + }, + ) + + assert approved.status_code == 200 + results = [event for event in _decode_sse_events(approved) if event.get("type") == "TOOL_CALL_RESULT"] + assert [(event["toolCallId"], event["content"]) for event in results] == [("call-guarded", "guarded result")] + assert executed == ["guarded"] + snapshot = await snapshot_store.get(scope="tenant-a", thread_id=thread_id) + assert snapshot is not None and snapshot.session_state is not None + assert snapshot.session_state[security.source_id]["pending_policy_approvals"] == {} + + +async def test_endpoint_agent_approval_resume_persists_replayable_tool_results(streaming_chat_client_stub): + """Approved batches should hydrate with real results under original tool call ids.""" + client, executed, messages_received, state = _build_mixed_approval_batch_endpoint( + streaming_chat_client_stub, + snapshot_store=InMemoryAGUIThreadSnapshotStore(), + ) + + pause_response = client.post( + "/approval", + json={ + "runId": "run-pause", + "threadId": "thread-mixed-replay", + "messages": [{"id": "user-1", "role": "user", "content": "Run both tools"}], + }, + ) + assert pause_response.status_code == 200 + pause_finished = [event for event in _decode_sse_events(pause_response) if event.get("type") == "RUN_FINISHED"] + interrupts = _run_finished_interrupts(pause_finished[-1]) + assert len(interrupts) == 1 + approval_id = interrupts[0]["id"] + assert approval_id.startswith("af-call-") + assert interrupts[0]["toolCallId"] == "call_sensitive" + + state["phase"] = "resume" + resume_response = client.post( + "/approval", + json={ + "runId": "run-resume", + "threadId": "thread-mixed-replay", + "messages": [], + "resume": [{"interruptId": approval_id, "status": "resolved", "payload": {"accepted": True}}], + }, + ) + + assert resume_response.status_code == 200 + resume_events = _decode_sse_events(resume_response) + live_results = [ + (event["toolCallId"], event["content"]) for event in resume_events if event.get("type") == "TOOL_CALL_RESULT" + ] + assert live_results == [ + ("call_sensitive", "Sensitive action in Seattle"), + ("call_weather", "Weather in Seattle"), + ] + assert Counter(call_id for call_id, _ in live_results) == {"call_sensitive": 1, "call_weather": 1} + assert executed == ["sensitive:Seattle", "weather:Seattle"] + + hydrate_response = client.post( + "/approval", + json={"runId": "run-hydrate", "threadId": "thread-mixed-replay", "messages": []}, + ) + + assert hydrate_response.status_code == 200 + hydrated_messages = _latest_messages_snapshot(hydrate_response) + tool_messages = [message for message in hydrated_messages if message.get("role") == "tool"] + replayed_results = [ + (message.get("toolCallId"), message.get("content")) + for message in tool_messages + if message.get("toolCallId") in {"call_sensitive", "call_weather"} + ] + assert replayed_results == [ + ("call_sensitive", "Sensitive action in Seattle"), + ("call_weather", "Weather in Seattle"), + ] + assert Counter(call_id for call_id, _ in replayed_results) == {"call_sensitive": 1, "call_weather": 1} + assert not any(message.get("function_approvals") for message in hydrated_messages) + assert not any("Tool execution skipped" in str(message.get("content")) for message in hydrated_messages) + + state["phase"] = "next" + next_response = client.post( + "/approval", + json={ + "runId": "run-next", + "threadId": "thread-mixed-replay", + "messages": [{"id": "user-2", "role": "user", "content": "Continue"}], + }, + ) + assert next_response.status_code == 200 + provider_results = [ + content for message in messages_received for content in message.contents if content.type == "function_result" + ] + assert [(content.call_id, content.result) for content in provider_results] == [ + ("call_sensitive", "Sensitive action in Seattle"), + ("call_weather", "Weather in Seattle"), + ] + assert Counter(content.call_id for content in provider_results) == {"call_sensitive": 1, "call_weather": 1} + + +async def test_endpoint_agent_approval_resume_surfaces_queued_tool_approval(streaming_chat_client_stub): + """Queued approvals retain decisions until the wrapped Agent's middleware releases the batch.""" + client, executed, messages_received, state, _ = _build_tool_approval_queue_endpoint(streaming_chat_client_stub) + pause_response = client.post( + "/approval", + json={ + "runId": "run-pause", + "threadId": "thread-queued-approval", + "messages": [{"role": "user", "content": "Run both tools"}], + }, + ) + assert pause_response.status_code == 200 + pause_finished = [event for event in _decode_sse_events(pause_response) if event.get("type") == "RUN_FINISHED"] + first_interrupts = _run_finished_interrupts(pause_finished[-1]) + assert len(first_interrupts) == 1 + first_approval_id = first_interrupts[0]["id"] + assert first_approval_id.startswith("af-call-") + assert first_interrupts[0]["toolCallId"] == "call_first" + assert executed == [] + + state["phase"] = "resume" + first_resume = client.post( + "/approval", + json={ + "runId": "run-resume-first", + "threadId": "thread-queued-approval", + "messages": [], + "resume": [{"interruptId": first_approval_id, "status": "resolved", "payload": {"accepted": True}}], + }, + ) + + assert first_resume.status_code == 200 + first_resume_events = _decode_sse_events(first_resume) + tool_results = [event for event in first_resume_events if event.get("type") == "TOOL_CALL_RESULT"] + assert tool_results == [] + first_resume_finished = [event for event in first_resume_events if event.get("type") == "RUN_FINISHED"] + second_interrupts = _run_finished_interrupts(first_resume_finished[-1]) + assert len(second_interrupts) == 1 + second_approval_id = second_interrupts[0]["id"] + assert second_approval_id.startswith("af-call-") + assert second_interrupts[0]["toolCallId"] == "call_second" + assert not [ + event + for event in first_resume_events + if event.get("type") == "TOOL_CALL_END" and event.get("toolCallId") == "call_first" + ] + assert executed == [] + assert messages_received == [] + + final_resume = client.post( + "/approval", + json={ + "runId": "run-resume-second", + "threadId": "thread-queued-approval", + "messages": [], + "resume": [{"interruptId": second_approval_id, "status": "resolved", "payload": {"accepted": True}}], + }, + ) + + assert final_resume.status_code == 200 + final_events = _decode_sse_events(final_resume) + final_tool_results = [event for event in final_events if event.get("type") == "TOOL_CALL_RESULT"] + assert [(event["toolCallId"], event["content"]) for event in final_tool_results] == [ + ("call_first", "first result"), + ("call_second", "second result"), + ] + assert executed == ["first", "second"] + 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_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 - hydrate_response = client.post( - "/approval", - json={"runId": "run-hydrate", "threadId": "thread-mixed-replay", "messages": []}, - ) + 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"] - assert hydrate_response.status_code == 200 - hydrated_messages = _latest_messages_snapshot(hydrate_response) - tool_messages = [message for message in hydrated_messages if message.get("role") == "tool"] - replayed_results = [ - (message.get("toolCallId"), message.get("content")) - for message in tool_messages - if message.get("toolCallId") in {"call_sensitive", "call_weather"} - ] - assert replayed_results == [ - ("call_sensitive", "Sensitive action in Seattle"), - ("call_weather", "Weather in Seattle"), + # 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) ] - assert Counter(call_id for call_id, _ in replayed_results) == {"call_sensitive": 1, "call_weather": 1} - assert not any(message.get("function_approvals") for message in hydrated_messages) - assert not any("Tool execution skipped" in str(message.get("content")) for message in hydrated_messages) - state["phase"] = "next" - next_response = client.post( - "/approval", - json={ - "runId": "run-next", - "threadId": "thread-mixed-replay", - "messages": [{"id": "user-2", "role": "user", "content": "Continue"}], - }, - ) - assert next_response.status_code == 200 - provider_results = [ - content for message in messages_received for content in message.contents if content.type == "function_result" + 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 [(content.call_id, content.result) for content in provider_results] == [ - ("call_sensitive", "Sensitive action in Seattle"), - ("call_weather", "Weather in Seattle"), + assert [entry[:2] for entry in replayed] == [ + ("function_call", "call_reused"), + ("function_result", "call_reused"), + ("function_call", "call_reused"), + ("function_result", "call_reused"), ] - assert Counter(content.call_id for content in provider_results) == {"call_sensitive": 1, "call_weather": 1} + assert replayed[1][2] == "result 1" + assert "Tool execution skipped" in str(replayed[3][2]) -async def test_endpoint_agent_approval_resume_surfaces_queued_tool_approval(streaming_chat_client_stub): - """Queued approvals retain decisions until the wrapped Agent's middleware releases the batch.""" - client, executed, messages_received, state, _ = _build_tool_approval_queue_endpoint(streaming_chat_client_stub) - pause_response = client.post( +@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 +) -> None: + """A collected user grant cannot be recreated after its original authority expires.""" + now = 0.0 + lifecycle = ApprovalLifecycle( + pending_retention_seconds=30, + terminal_retention_seconds=1, + clock=lambda: now, + ) + client, executed, _, state, _ = _build_tool_approval_queue_endpoint(streaming_chat_client_stub, lifecycle=lifecycle) + pause = client.post( "/approval", - json={ - "runId": "run-pause", - "threadId": "thread-queued-approval", - "messages": [{"role": "user", "content": "Run both tools"}], - }, + json={"runId": "pause", "threadId": "queue-expiration", "messages": [{"role": "user", "content": "Act"}]}, ) - assert pause_response.status_code == 200 - pause_finished = [event for event in _decode_sse_events(pause_response) if event.get("type") == "RUN_FINISHED"] - first_interrupts = _run_finished_interrupts(pause_finished[-1]) - assert len(first_interrupts) == 1 - first_approval_id = first_interrupts[0]["id"] - assert first_approval_id.startswith("af-call-") - assert first_interrupts[0]["toolCallId"] == "call_first" - assert executed == [] - + first_finished = [event for event in _decode_sse_events(pause) if event["type"] == "RUN_FINISHED"] + first = _run_finished_interrupts(first_finished[-1])[0]["id"] + now = 20.0 state["phase"] = "resume" - first_resume = client.post( + collected = client.post( "/approval", json={ - "runId": "run-resume-first", - "threadId": "thread-queued-approval", + "runId": "collect", + "threadId": "queue-expiration", "messages": [], - "resume": [{"interruptId": first_approval_id, "status": "resolved", "payload": {"accepted": True}}], + "resume": [{"interruptId": first, "status": "resolved", "payload": {"approved": True}}], }, ) - - assert first_resume.status_code == 200 - first_resume_events = _decode_sse_events(first_resume) - tool_results = [event for event in first_resume_events if event.get("type") == "TOOL_CALL_RESULT"] - assert tool_results == [] - first_resume_finished = [event for event in first_resume_events if event.get("type") == "RUN_FINISHED"] - second_interrupts = _run_finished_interrupts(first_resume_finished[-1]) - assert len(second_interrupts) == 1 - second_approval_id = second_interrupts[0]["id"] - assert second_approval_id.startswith("af-call-") - assert second_interrupts[0]["toolCallId"] == "call_second" - assert not [ - event - for event in first_resume_events - if event.get("type") == "TOOL_CALL_END" and event.get("toolCallId") == "call_first" - ] + second_finished = [event for event in _decode_sse_events(collected) if event["type"] == "RUN_FINISHED"] + second = _run_finished_interrupts(second_finished[-1])[0]["id"] assert executed == [] - assert messages_received == [] - - final_resume = client.post( + now = retry_time + resumed = client.post( "/approval", json={ - "runId": "run-resume-second", - "threadId": "thread-queued-approval", + "runId": "complete", + "threadId": "queue-expiration", "messages": [], - "resume": [{"interruptId": second_approval_id, "status": "resolved", "payload": {"accepted": True}}], + "resume": [{"interruptId": second, "status": "resolved", "payload": {"approved": True}}], }, ) + errors = [event for event in _decode_sse_events(resumed) if event["type"] == "RUN_ERROR"] + if retry_time == 25.0: + assert errors == [] + assert executed == ["first", "second"] + else: + assert errors + assert executed == [] - assert final_resume.status_code == 200 - final_events = _decode_sse_events(final_resume) - final_tool_results = [event for event in final_events if event.get("type") == "TOOL_CALL_RESULT"] - assert [(event["toolCallId"], event["content"]) for event in final_tool_results] == [ - ("call_first", "first result"), - ("call_second", "second result"), - ] - assert executed == ["first", "second"] - 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("mixed_decisions", [False, True], ids=["deny-only", "mixed-approve-deny"]) +async def test_endpoint_fides_authenticated_rejection_preserves_restored_context( + streaming_chat_client_stub: Any, + mixed_decisions: bool, +) -> None: + """A lifecycle-authenticated denial remains control-plane replay, not external data.""" + phase = "approval" + executed: list[str] = [] + provider_calls = 0 -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 guarded_action() -> str: + executed.append("guarded") + return "guarded result" - 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 guarded_action_two() -> str: + executed.append("guarded-two") + return "guarded result two" - 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 protected_action() -> str: + executed.append("protected") + return "protected result" - def tool_result_call_ids(messages: list[dict[str, Any]]) -> list[str]: - return [ - str(message.get("toolCallId") or message.get("tool_call_id")) + async def stream_fn( + messages: list[Message], + options: dict[str, Any], + **kwargs: Any, + ) -> AsyncIterator[ChatResponseUpdate]: + nonlocal provider_calls + del options, kwargs + provider_calls += 1 + if phase == "approval": + guarded_calls = [ + Content.from_function_call( + call_id="guarded-call", + name="guarded_action", + arguments={}, + ) + ] + if mixed_decisions: + guarded_calls.append( + Content.from_function_call( + call_id="guarded-call-two", + name="guarded_action_two", + arguments={}, + ) + ) + yield ChatResponseUpdate( + role="assistant", + contents=guarded_calls, + ) + return + if phase == "denied": + yield ChatResponseUpdate(role="assistant", contents=[Content.from_text("denied")]) + return + if any( + content.type == "function_result" and content.call_id == "protected-call" for message in messages - if message.get("role") == "tool" - ] + for content in message.contents + ): + yield ChatResponseUpdate(role="assistant", contents=[Content.from_text("done")]) + return + yield ChatResponseUpdate( + role="assistant", + contents=[Content.from_function_call(call_id="protected-call", name="protected_action", arguments={})], + ) - 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"] + security = SecureAgentConfig() + agent = Agent( + client=streaming_chat_client_stub(stream_fn), + tools=[ + FunctionTool( + func=guarded_action, + name="guarded_action", + description="Require approval", + approval_mode="always_require", + ), + FunctionTool( + func=protected_action, + name="protected_action", + description="Run an unrelated protected action", + additional_properties={ + "source_integrity": "trusted", + "max_allowed_confidentiality": "public", + }, + ), + FunctionTool( + func=guarded_action_two, + name="guarded_action_two", + description="Run an already-approved sibling action", + approval_mode="never_require", + ), + ], + context_providers=[security], + ) + store = InMemoryAGUIThreadSnapshotStore() + app = FastAPI() + add_agent_framework_fastapi_endpoint( + app, + agent, + path="/fides-authenticated-rejection", + snapshot_store=store, + snapshot_scope_resolver=lambda _request: "test", + keepalive_seconds=None, + ) + + with TestClient(app) as client: + pause = client.post( + "/fides-authenticated-rejection", + json={ + "runId": "run-rejection-pause", + "threadId": "thread-authenticated-rejection", + "messages": [{"id": "user-1", "role": "user", "content": "Prepare guarded action"}], + }, + ) + finished = [event for event in _decode_sse_events(pause) if event.get("type") == "RUN_FINISHED"] + interrupts = _run_finished_interrupts(finished[-1]) + assert len(interrupts) == 1 + approval_ids = {interrupt["toolCallId"]: interrupt["id"] for interrupt in interrupts} + resume_entries = [ + { + "interruptId": approval_ids["guarded-call"], + "status": "resolved", + "payload": {"approved": False}, + } + ] - 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"] + phase = "denied" + denied = client.post( + "/fides-authenticated-rejection", + json={ + "runId": "run-rejection-denied", + "threadId": "thread-authenticated-rejection", + "messages": [], + "resume": resume_entries, + }, + ) + assert denied.status_code == 200 + assert not [event for event in _decode_sse_events(denied) if event.get("type") == "RUN_ERROR"] + snapshot = await store.get(scope="test", thread_id="thread-authenticated-rejection") + assert snapshot is not None + assert snapshot.interrupt is None + assert snapshot.session_state is not None + label = json.loads(snapshot.session_state[security.source_id]["context_label"]) + assert label["integrity"] == "trusted" + assert label["confidentiality"] == "public" + + phase = "protected" + protected = client.post( + "/fides-authenticated-rejection", + json={ + "runId": "run-rejection-protected", + "threadId": "thread-authenticated-rejection", + "messages": [{"id": "user-2", "role": "user", "content": "Run protected action"}], + }, + ) - 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"] + assert protected.status_code == 200 + assert not [event for event in _decode_sse_events(protected) if event.get("type") == "RUN_ERROR"] + assert executed == (["guarded-two", "protected"] if mixed_decisions else ["protected"]) + assert provider_calls == 4 - 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_fides_rejection_does_not_cover_client_result_with_reused_call_id( + streaming_chat_client_stub: Any, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Only the lifecycle-synthesized rejection result receives backend provenance.""" + import agent_framework_ag_ui._agent_run as agent_run_module + + monkeypatch.setattr( + agent_run_module, + "_filter_local_approval_responses_for_provider", + lambda messages, **_kwargs: list(messages), + ) -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() + phase = "approval" executed: list[str] = [] - messages_received: list[Message] = [] - turn = {"number": 0} - def lookup() -> str: - executed.append(f"turn-{turn['number']}") - return f"result {turn['number']}" + def guarded_action() -> str: + return "guarded result" + + def protected_action() -> str: + executed.append("protected") + return "protected result" async def stream_fn( messages: list[Message], @@ -7630,113 +8841,147 @@ async def stream_fn( **kwargs: Any, ) -> AsyncIterator[ChatResponseUpdate]: del options, kwargs - messages_received[:] = list(messages) - if messages[-1].role == "user": - turn["number"] += 1 + if phase == "approval": yield ChatResponseUpdate( - contents=[Content.from_function_call(call_id="call_reused", name="lookup", arguments="{}")], role="assistant", + contents=[ + Content.from_function_call( + call_id="shared-call", + name="guarded_action", + arguments={}, + ) + ], ) return - yield ChatResponseUpdate(contents=[Content.from_text(text=f"Done {turn['number']}.")], role="assistant") + if any( + content.type == "function_result" and content.call_id == "protected-call" + for message in messages + for content in message.contents + ): + yield ChatResponseUpdate(role="assistant", contents=[Content.from_text("done")]) + return + yield ChatResponseUpdate( + role="assistant", + contents=[Content.from_function_call(call_id="protected-call", name="protected_action", arguments={})], + ) + security = SecureAgentConfig() 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()], + tools=[ + FunctionTool( + func=guarded_action, + name="guarded_action", + description="Require approval", + approval_mode="always_require", + ), + FunctionTool( + func=protected_action, + name="protected_action", + description="Run a trusted public action", + additional_properties={ + "source_integrity": "trusted", + "max_allowed_confidentiality": "public", + }, + ), + ], + context_providers=[security], ) + store = InMemoryAGUIThreadSnapshotStore() app = FastAPI() add_agent_framework_fastapi_endpoint( app, - AgentFrameworkAgent(agent=agent, require_confirmation=False), - path="/approval", + agent, + path="/fides-reused-rejection-result", snapshot_store=store, snapshot_scope_resolver=lambda _request: "test", + keepalive_seconds=None, ) - 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 + with TestClient(app) as client: + pause = client.post( + "/fides-reused-rejection-result", + json={ + "runId": "run-reused-rejection-pause", + "threadId": "thread-reused-rejection-result", + "messages": [{"id": "user-1", "role": "user", "content": "Prepare the guarded action"}], + }, + ) + finished = [event for event in _decode_sse_events(pause) if event.get("type") == "RUN_FINISHED"] + approval_id = _run_finished_interrupts(finished[-1])[0]["id"] + snapshot = await store.get(scope="test", thread_id="thread-reused-rejection-result") + assert snapshot is not None + snapshot.messages = [ + message + for message in snapshot.messages + if not (message.get("role") == "assistant" and (message.get("toolCalls") or message.get("tool_calls"))) + ] + snapshot.messages.append( + { + "id": "assistant-earlier-call", + "role": "assistant", + "content": "", + "toolCalls": [ + { + "id": "shared-call", + "type": "function", + "function": {"name": "client_lookup", "arguments": "{}"}, + } + ], + } + ) + await store.save(scope="test", thread_id="thread-reused-rejection-result", snapshot=snapshot) - saved = await store.get(scope="test", thread_id=thread_id) - assert saved is not None - assert transcript(saved.messages) == expected + phase = "protected" + denied = client.post( + "/fides-reused-rejection-result", + json={ + "runId": "run-reused-rejection-denied", + "threadId": "thread-reused-rejection-result", + "messages": [ + { + "id": "client-owned-result", + "role": "tool", + "toolCallId": "shared-call", + "content": "Error: Tool call invocation was rejected by user.", + } + ], + "resume": [ + { + "interruptId": approval_id, + "status": "resolved", + "payload": {"approved": False}, + } + ], + }, + ) - 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"), - ] + assert denied.status_code == 200 + events = _decode_sse_events(denied) + assert not [event for event in events if event.get("type") == "RUN_ERROR"] + assert executed == [] + assert _is_policy_violation([event["content"] for event in events if event["type"] == "TOOL_CALL_RESULT"]) + snapshot = await store.get(scope="test", thread_id="thread-reused-rejection-result") + assert snapshot is not None + assert snapshot.session_state is not None + label = json.loads(snapshot.session_state[security.source_id]["context_label"]) + assert label["integrity"] == "untrusted" + assert label["confidentiality"] == "private" -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() +async def test_endpoint_fides_rejected_edited_approval_invalidates_snapshot_provenance( + streaming_chat_client_stub: Any, +) -> None: + """Client edits to a stored call remain conservative when lifecycle validation rejects them.""" + phase = "approval" executed: list[str] = [] - messages_received: list[Message] = [] - def lookup() -> str: - executed.append("lookup") - return "result 1" + def guarded_action(value: str) -> str: + return value + + def protected_action() -> str: + executed.append("protected") + return "protected result" async def stream_fn( messages: list[Message], @@ -7744,139 +8989,118 @@ async def stream_fn( **kwargs: Any, ) -> AsyncIterator[ChatResponseUpdate]: del options, kwargs - messages_received[:] = list(messages) - last = messages[-1] - if last.role == "user" and last.text == "First": + if phase == "approval": yield ChatResponseUpdate( - contents=[Content.from_function_call(call_id="call_reused", name="lookup", arguments="{}")], role="assistant", + contents=[ + Content.from_function_call( + call_id="guarded-call", + name="guarded_action", + arguments={"value": "original"}, + ) + ], ) 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", - ) + if any( + content.type == "function_result" and content.call_id == "protected-call" + for message in messages + for content in message.contents + ): + yield ChatResponseUpdate(role="assistant", contents=[Content.from_text("done")]) return - yield ChatResponseUpdate(contents=[Content.from_text(text="Done.")], role="assistant") + yield ChatResponseUpdate( + role="assistant", + contents=[Content.from_function_call(call_id="protected-call", name="protected_action", arguments={})], + ) + security = SecureAgentConfig() 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()], + tools=[ + FunctionTool( + func=guarded_action, + name="guarded_action", + description="Require approval", + approval_mode="always_require", + ), + FunctionTool( + func=protected_action, + name="protected_action", + description="Run a trusted public action", + additional_properties={ + "source_integrity": "trusted", + "max_allowed_confidentiality": "public", + }, + ), + ], + context_providers=[security], ) - wrapped_agent = AgentFrameworkAgent(agent=agent, require_confirmation=False) + store = InMemoryAGUIThreadSnapshotStore() 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) + add_agent_framework_fastapi_endpoint( + app, + agent, + path="/fides-rejected-edited-approval", + snapshot_store=store, + snapshot_scope_resolver=lambda _request: "test", + keepalive_seconds=None, ) - 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]) + with TestClient(app) as client: + pause = client.post( + "/fides-rejected-edited-approval", + json={ + "runId": "run-edited-approval-pause", + "threadId": "thread-edited-approval", + "messages": [{"id": "user-1", "role": "user", "content": "Prepare the guarded action"}], + }, + ) + finished = [event for event in _decode_sse_events(pause) if event.get("type") == "RUN_FINISHED"] + interrupt = _run_finished_interrupts(finished[-1])[0] + stored_snapshot = await store.get(scope="test", thread_id="thread-edited-approval") + assert stored_snapshot is not None + phase = "protected" + response = client.post( + "/fides-rejected-edited-approval", + json={ + "runId": "run-edited-approval-rejected", + "threadId": "thread-edited-approval", + "messages": [ + *stored_snapshot.messages, + { + "id": "edited-approval", + "role": "tool", + "toolCallId": interrupt["toolCallId"], + "content": json.dumps({"accepted": True, "value": "client-edited"}), + }, + { + "id": "user-2", + "role": "user", + "content": "Continue without accepting the edited call.", + }, + ], + "resume": [ + { + "interruptId": interrupt["id"], + "status": "resolved", + "payload": {"approved": False}, + } + ], + }, + ) -@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 -) -> None: - """A collected user grant cannot be recreated after its original authority expires.""" - now = 0.0 - lifecycle = ApprovalLifecycle( - pending_retention_seconds=30, - terminal_retention_seconds=1, - clock=lambda: now, - ) - client, executed, _, state, _ = _build_tool_approval_queue_endpoint(streaming_chat_client_stub, lifecycle=lifecycle) - pause = client.post( - "/approval", - json={"runId": "pause", "threadId": "queue-expiration", "messages": [{"role": "user", "content": "Act"}]}, - ) - first_finished = [event for event in _decode_sse_events(pause) if event["type"] == "RUN_FINISHED"] - first = _run_finished_interrupts(first_finished[-1])[0]["id"] - now = 20.0 - state["phase"] = "resume" - collected = client.post( - "/approval", - json={ - "runId": "collect", - "threadId": "queue-expiration", - "messages": [], - "resume": [{"interruptId": first, "status": "resolved", "payload": {"approved": True}}], - }, - ) - second_finished = [event for event in _decode_sse_events(collected) if event["type"] == "RUN_FINISHED"] - second = _run_finished_interrupts(second_finished[-1])[0]["id"] + assert response.status_code == 200 + events = _decode_sse_events(response) + assert not [event for event in events if event.get("type") == "RUN_ERROR"] assert executed == [] - now = retry_time - resumed = client.post( - "/approval", - json={ - "runId": "complete", - "threadId": "queue-expiration", - "messages": [], - "resume": [{"interruptId": second, "status": "resolved", "payload": {"approved": True}}], - }, - ) - errors = [event for event in _decode_sse_events(resumed) if event["type"] == "RUN_ERROR"] - if retry_time == 25.0: - assert errors == [] - assert executed == ["first", "second"] - else: - assert errors - assert executed == [] + assert _is_policy_violation([event["content"] for event in events if event["type"] == "TOOL_CALL_RESULT"]) + snapshot = await store.get(scope="test", thread_id="thread-edited-approval") + assert snapshot is not None + assert snapshot.session_state is not None + label = json.loads(snapshot.session_state[security.source_id]["context_label"]) + assert label["integrity"] == "untrusted" + assert label["confidentiality"] == "private" async def test_endpoint_agent_approval_cancel_discards_queued_tool_approval(streaming_chat_client_stub): @@ -12293,6 +13517,88 @@ async def stream_fn(messages: Any, options: Any, **kwargs: Any): assert any(message.get("role") == "user" and message.get("content") == "Draft the plan" for message in messages) +async def test_agent_endpoint_confirmation_without_fides_does_not_rebuild_snapshot( + streaming_chat_client_stub: Any, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Confirmation avoids a second full transcript conversion when no FIDES replay folder exists.""" + import agent_framework_ag_ui._agent_run as agent_run_module + + normalization_calls = 0 + original_normalize = agent_run_module.normalize_agui_input_messages + + def track_normalization(*args: Any, **kwargs: Any) -> Any: + nonlocal normalization_calls + normalization_calls += 1 + return original_normalize(*args, **kwargs) + + monkeypatch.setattr(agent_run_module, "normalize_agui_input_messages", track_normalization) + + async def stream_fn( + messages: list[Message], + options: dict[str, Any], + **kwargs: Any, + ) -> AsyncIterator[ChatResponseUpdate]: + del messages, options, kwargs + yield ChatResponseUpdate( + role="assistant", + contents=[ + Content.from_function_call( + name="draft_steps", + call_id="draft-call", + arguments={"steps": [{"description": "Draft outline"}]}, + ) + ], + ) + + agent = Agent(name="test", instructions="Test agent", client=streaming_chat_client_stub(stream_fn)) + store = InMemoryAGUIThreadSnapshotStore() + app = FastAPI() + add_agent_framework_fastapi_endpoint( + app, + agent, + path="/confirm-without-fides", + state_schema={"steps": {"type": "array", "items": {"type": "object"}}}, + predict_state_config={"steps": {"tool": "draft_steps", "tool_argument": "steps"}}, + snapshot_store=store, + snapshot_scope_resolver=lambda _request: "test", + keepalive_seconds=None, + ) + + with TestClient(app) as client: + pause = client.post( + "/confirm-without-fides", + json={ + "runId": "run-confirm-without-fides-pause", + "threadId": "thread-confirm-without-fides", + "messages": [{"id": "user-1", "role": "user", "content": "Draft the plan"}], + }, + ) + finished = [event for event in _decode_sse_events(pause) if event.get("type") == "RUN_FINISHED"] + confirm_id = _run_finished_interrupts(finished[-1])[0]["id"] + normalization_calls = 0 + + confirmed = client.post( + "/confirm-without-fides", + json={ + "runId": "run-confirm-without-fides-finish", + "threadId": "thread-confirm-without-fides", + "messages": [], + "resume": [ + { + "interruptId": confirm_id, + "status": "resolved", + "payload": json.dumps({"accepted": True, "steps": []}), + } + ], + }, + ) + + assert confirmed.status_code == 200 + assert not [event for event in _decode_sse_events(confirmed) if event.get("type") == "RUN_ERROR"] + assert normalization_calls == 1 + + async def test_agent_endpoint_default_state_does_not_reset_persisted_state(streaming_chat_client_stub): """Endpoint defaults fill missing keys but never override persisted Shared State.""" app = FastAPI() @@ -12749,6 +14055,186 @@ async def updates() -> AsyncIterator[AgentResponseUpdate]: assert "Tool call invocation was rejected by user" in str(snapshot.messages) +async def test_endpoint_fides_rejection_folds_client_result_before_provider_failure( + streaming_chat_client_stub: Any, +) -> None: + """A pre-run rejection snapshot persists fresh client-result restrictions.""" + from agent_framework import AgentResponse, ResponseInvalidatedException, ResponseStream + + phase = "pause" + executed: list[str] = [] + + def guarded_action() -> str: + return "guarded result" + + def protected_action() -> str: + executed.append("protected") + return "protected result" + + async def stream_fn( + messages: list[Message], + options: dict[str, Any], + **kwargs: Any, + ) -> AsyncIterator[ChatResponseUpdate]: + del options, kwargs + if phase == "pause": + yield ChatResponseUpdate( + role="assistant", + contents=[ + Content.from_function_call( + call_id="guarded-call", + name="guarded_action", + arguments={}, + ) + ], + ) + return + if any( + content.type == "function_result" and content.call_id == "protected-call" + for message in messages + for content in message.contents + ): + yield ChatResponseUpdate(role="assistant", contents=[Content.from_text("done")]) + return + yield ChatResponseUpdate( + role="assistant", + contents=[Content.from_function_call(call_id="protected-call", name="protected_action", arguments={})], + ) + + config = SecureAgentConfig() + agent = Agent( + client=streaming_chat_client_stub(stream_fn), + tools=[ + FunctionTool( + func=guarded_action, + name="guarded_action", + description="Require approval", + approval_mode="always_require", + ), + FunctionTool( + func=protected_action, + name="protected_action", + description="Run an unrelated protected action", + additional_properties={ + "source_integrity": "trusted", + "max_allowed_confidentiality": "public", + }, + ), + ], + context_providers=[config], + ) + store = InMemoryAGUIThreadSnapshotStore() + app = FastAPI() + add_agent_framework_fastapi_endpoint( + app, + agent, + path="/fides-rejection-provider-failure", + snapshot_store=store, + snapshot_scope_resolver=lambda _request: "test", + keepalive_seconds=None, + ) + + with TestClient(app) as client: + pause_response = client.post( + "/fides-rejection-provider-failure", + json={ + "runId": "run-fides-rejection-pause", + "threadId": "thread-fides-rejection-provider-failure", + "messages": [{"id": "user-1", "role": "user", "content": "Prepare the guarded action"}], + }, + ) + pause_finished = [event for event in _decode_sse_events(pause_response) if event.get("type") == "RUN_FINISHED"] + approval_id = _run_finished_interrupts(pause_finished[-1])[0]["id"] + snapshot = await store.get(scope="test", thread_id="thread-fides-rejection-provider-failure") + assert snapshot is not None + assert snapshot.session_state is not None + initial_label = json.loads(snapshot.session_state[config.source_id]["context_label"]) + assert initial_label["integrity"] == "trusted" + assert initial_label["confidentiality"] == "public" + snapshot.messages.append( + { + "id": "assistant-client-call", + "role": "assistant", + "content": "", + "toolCalls": [ + { + "id": "client-call", + "type": "function", + "function": {"name": "client_lookup", "arguments": "{}"}, + } + ], + } + ) + await store.save( + scope="test", + thread_id="thread-fides-rejection-provider-failure", + snapshot=snapshot, + ) + + original_run = agent.run + + def failing_run(*args: Any, **kwargs: Any) -> Any: + if not kwargs.get("stream", False): + return original_run(*args, **kwargs) + + async def updates() -> AsyncIterator[AgentResponseUpdate]: + if False: # pragma: no cover + yield AgentResponseUpdate() + raise ResponseInvalidatedException("provider preparation failed") + + return ResponseStream(updates(), finalizer=AgentResponse.from_updates) + + agent.run = failing_run # type: ignore[assignment, method-assign] # ty: ignore[invalid-assignment] + rejected_response = client.post( + "/fides-rejection-provider-failure", + json={ + "runId": "run-fides-rejection-failure", + "threadId": "thread-fides-rejection-provider-failure", + "messages": [ + { + "id": "client-result", + "role": "tool", + "toolCallId": "client-call", + "content": "fresh external client result", + } + ], + "resume": [ + { + "interruptId": approval_id, + "status": "resolved", + "payload": {"approved": False}, + } + ], + }, + ) + rejected_events = _decode_sse_events(rejected_response) + assert [event["code"] for event in rejected_events if event.get("type") == "RUN_ERROR"] == [ + "ResponseInvalidatedException" + ] + failed_snapshot = await store.get(scope="test", thread_id="thread-fides-rejection-provider-failure") + assert failed_snapshot is not None + assert failed_snapshot.session_state is not None + failed_label = json.loads(failed_snapshot.session_state[config.source_id]["context_label"]) + assert failed_label["integrity"] == "untrusted" + assert failed_label["confidentiality"] == "private" + + agent.run = original_run # type: ignore[assignment, method-assign] # ty: ignore[invalid-assignment] + phase = "protected" + retry_response = client.post( + "/fides-rejection-provider-failure", + json={ + "runId": "run-fides-rejection-retry", + "threadId": "thread-fides-rejection-provider-failure", + "messages": [{"id": "user-2", "role": "user", "content": "Run the protected action"}], + }, + ) + + retry_events = _decode_sse_events(retry_response) + assert not [event for event in retry_events if event.get("type") == "RUN_ERROR"] + assert executed == [] + assert _is_policy_violation([event["content"] for event in retry_events if event["type"] == "TOOL_CALL_RESULT"]) + + async def test_agent_endpoint_approval_resume_preserves_other_stored_interrupts(): """A successful approval continuation retains unrelated stored interrupts.""" store = InMemoryAGUIThreadSnapshotStore() diff --git a/python/packages/ag-ui/tests/ag_ui/test_message_hygiene.py b/python/packages/ag-ui/tests/ag_ui/test_message_hygiene.py index fba1b91fb09..d7cdd9b74f6 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_message_hygiene.py +++ b/python/packages/ag-ui/tests/ag_ui/test_message_hygiene.py @@ -145,6 +145,36 @@ def test_convert_approval_results_preserves_other_user_content() -> None: assert messages[2].contents[0].type == "text" +def test_convert_approval_results_marks_only_exact_backend_rejection() -> None: + """Reusable call IDs and rejection text cannot authenticate a client-owned result.""" + from agent_framework_ag_ui._agent_run import _convert_approval_results_to_tool_messages + from agent_framework_ag_ui._message_adapters import _take_backend_snapshot_messages + + client_result = Content.from_function_result( + call_id="shared-call", + result="Error: Tool call invocation was rejected by user.", + ) + synthesized_rejection = Content.from_function_result( + call_id="shared-call", + result="Error: Tool call invocation was rejected by user.", + ) + messages = [ + Message(role="tool", contents=[client_result]), + Message(role="user", contents=[synthesized_rejection]), + ] + + _convert_approval_results_to_tool_messages( + messages, + backend_result_content_ids={id(synthesized_rejection)}, + ) + + covered_messages = _take_backend_snapshot_messages(messages) + assert len(covered_messages) == 1 + assert len(covered_messages[0].contents) == 1 + assert covered_messages[0].contents[0] is synthesized_rejection + assert all(content is not client_result for message in covered_messages for content in message.contents) + + def test_sanitize_tool_history_filters_confirm_changes_keeps_other_tools() -> None: """Test that confirm_changes is filtered but other tools are preserved. diff --git a/python/packages/core/agent_framework/security.py b/python/packages/core/agent_framework/security.py index e4bb70d856c..502cdb8e68b 100644 --- a/python/packages/core/agent_framework/security.py +++ b/python/packages/core/agent_framework/security.py @@ -120,6 +120,8 @@ _SECURITY_MATERIALIZED_AUTHORITY_CONTEXT_KEY = "_security_materialized_authority" _SECURITY_RESOLVED_AUTHORITY_CONTEXT_KEY = "_security_resolved_authority" _MAX_REPLAY_CONTENT_DEPTH = 32 +_REPLAY_STATE_COVERED_KEY = "_agent_framework_fides_replay_state_covered" +_REPLAY_STATE_COVERED_MARKER = object() _APPROVAL_HISTORY_RECONSTRUCTION_KEY = "_agent_framework_fides_reconstruct_history_unstamped" _APPROVAL_REPLAY_AUTHORITY_KEY = "_agent_framework_fides_authenticated_approval_calls" _TOOL_SECURITY_POLICY_KEYS = frozenset({ @@ -1280,16 +1282,17 @@ def _is_bound_approval_control( content: Content, *, scope: _SecurityScope, + transport_authoritative: bool, approval_authority: _ApprovalReplayAuthority | None, ) -> Content | None: - """Return whether session state authenticates this approval control.""" + """Return whether trusted transport or session state authenticates this approval control.""" if content.type not in {"function_approval_request", "function_approval_response"} or content.id is None: return None authoritative_call = approval_authority.function_call_for(content) if approval_authority is not None else None function_call = authoritative_call or content.function_call if function_call is None or function_call.id is None or _function_call_replay_key(function_call) is None: return None - if authoritative_call is not None or _is_pending_approval_control(content, scope=scope): + if authoritative_call is not None or transport_authoritative or _is_pending_approval_control(content, scope=scope): return function_call return None @@ -1328,6 +1331,7 @@ def _replayed_message_label( has_authenticated_hidden_content = False has_role_independent_untrusted_content = False has_bound_approval_control = False + transport_authoritative = _is_replay_state_covered(message) bound_approval_call_keys = { key for content in message.contents @@ -1336,6 +1340,7 @@ def _replayed_message_label( authoritative_call := _is_bound_approval_control( content, scope=scope, + transport_authoritative=transport_authoritative, approval_authority=approval_authority, ) ) @@ -1356,6 +1361,7 @@ def _replayed_message_label( authoritative_approval_call = _is_bound_approval_control( content, scope=scope, + transport_authoritative=transport_authoritative, approval_authority=approval_authority, ) is_bound_approval_control = authoritative_approval_call is not None @@ -1416,8 +1422,20 @@ def _replayed_message_label( return combine_labels(*labels) +def _is_replay_state_covered(message: Message) -> bool: + """Return whether trusted transport code marked this replay as represented by restored state.""" + return message.additional_properties.get(_REPLAY_STATE_COVERED_KEY) is _REPLAY_STATE_COVERED_MARKER + + +def _clear_replay_state_covered_markers(messages: Sequence[Message]) -> None: + """Remove private replay provenance before messages reach a chat provider.""" + for message in messages: + if message.additional_properties.get(_REPLAY_STATE_COVERED_KEY) is _REPLAY_STATE_COVERED_MARKER: + message.additional_properties.pop(_REPLAY_STATE_COVERED_KEY, None) + + class _ReplaySecurityChatMiddleware(ChatMiddleware): - """Raise the session label from replayed history-provider messages.""" + """Raise the session label from replayed caller and history-provider messages.""" def __init__( self, @@ -1446,6 +1464,7 @@ async def process( context: ChatContext, call_next: Callable[[], Awaitable[None]], ) -> None: + _clear_replay_state_covered_markers(context.messages) if not self._processed: self._processed = True history_messages = [message for message in context.messages if self._is_history_message(message)] @@ -4298,6 +4317,55 @@ def _function_middleware_for_approval_resolution(self, session: AgentSession) -> """Return provider-customized middleware bound to one restored session scope.""" return self._middleware_for_scope(self._scope_for_session(session)) + def _mark_replay_state_covered_messages(self, messages: Sequence[Message]) -> None: + """Mark trusted transport-restored messages as already represented by session state.""" + for message in messages: + message.additional_properties[_REPLAY_STATE_COVERED_KEY] = _REPLAY_STATE_COVERED_MARKER + + def _fold_uncovered_replay_messages( + self, + session: AgentSession, + messages: Sequence[Message], + *, + complete_replay_messages: Sequence[Message] | None = None, + replay_history_complete: bool = True, + ) -> None: + """Fold new transport messages before a snapshot-only return persists them.""" + state = session.state.setdefault(self.source_id, {}) + if not isinstance(state, dict): + raise ValueError("Security provider state must be a dictionary.") + restored_context_label = _STATE_CONTEXT_LABEL in state + if not restored_context_label and not replay_history_complete: + _clear_replay_state_covered_markers(messages) + if complete_replay_messages is not None and complete_replay_messages is not messages: + _clear_replay_state_covered_markers(complete_replay_messages) + raise RuntimeError( + "Cannot reconstruct service-managed history without an authoritative FIDES context label." + ) + scope = _SecurityScope(cast(dict[str, Any], state), scope_id=session.session_id) + replay_messages = ( + messages if restored_context_label or complete_replay_messages is None else complete_replay_messages + ) + approval_authority = self._authenticated_core_approval_calls( + session, + replay_messages, + ) + replay_labels = [ + _replayed_message_label( + message, + reconstruct_unstamped=True, + scope=scope, + approval_authority=approval_authority, + ) + for message in replay_messages + if not restored_context_label or not _is_replay_state_covered(message) + ] + if replay_labels: + scope.context_label = combine_labels(scope.context_label, *replay_labels) + _clear_replay_state_covered_markers(messages) + if replay_messages is not messages: + _clear_replay_state_covered_markers(replay_messages) + def _approval_history_metadata_key(self) -> str: """Return the run-local SessionContext metadata key for this provider.""" return f"{_APPROVAL_HISTORY_RECONSTRUCTION_KEY}:{self.source_id}" @@ -4460,7 +4528,7 @@ async def before_run( input_labels = [ _replayed_message_label( message, - reconstruct_unstamped=True, + reconstruct_unstamped=not restored_context_label or not _is_replay_state_covered(message), scope=scope, approval_authority=approval_authority, ) diff --git a/python/packages/core/tests/test_security.py b/python/packages/core/tests/test_security.py index 7313e919519..6a28e5bcae7 100644 --- a/python/packages/core/tests/test_security.py +++ b/python/packages/core/tests/test_security.py @@ -3699,6 +3699,64 @@ async def execute(_: FunctionInvocationContext) -> list[Content]: assert executed is True + async def test_transport_marked_replay_uses_restored_state_and_clears_marker(self) -> None: + """Trusted transport replay keeps restored confidentiality without leaking its marker.""" + config = SecureAgentConfig() + session = AgentSession(session_id="transport-covered-history") + tracker, _ = await _get_session_security_middleware(config, session) + tracker._update_context_label(ContentLabel(integrity=IntegrityLabel.UNTRUSTED)) + restored_message = Message(role="assistant", contents=["Previously covered history"]) + config._mark_replay_state_covered_messages([restored_message]) + provider_context = SessionContext(session_id=session.session_id, input_messages=[restored_message]) + + await config.before_run( + agent=SimpleNamespace(context_providers=[config]), + session=session, + context=provider_context, + state=session.state.setdefault(config.source_id, {}), + ) + + scoped_tracker = provider_context.get_middleware()[0] + replay_middleware = cast(ChatMiddleware, provider_context.get_middleware()[2]) + assert isinstance(scoped_tracker, LabelTrackingFunctionMiddleware) + label = scoped_tracker.get_context_label() + assert label.integrity == IntegrityLabel.UNTRUSTED + assert label.confidentiality == ConfidentialityLabel.PUBLIC + chat_context = ChatContext( + client=cast(Any, SimpleNamespace()), + messages=[restored_message], + options={}, + session=session, + ) + + async def call_next() -> None: + assert "_agent_framework_fides_replay_state_covered" not in restored_message.additional_properties + + await replay_middleware.process(chat_context, call_next) + + def test_transport_fold_rejects_missing_label_for_incomplete_service_history(self) -> None: + """Snapshot replay cannot stand in for unavailable service-owned history.""" + config = SecureAgentConfig() + session = AgentSession(session_id="incomplete-service-history", service_session_id="remote-conversation") + session.state[config.source_id] = {"history_reconstruction_complete": True} + replayed = Message( + role="tool", + contents=[Content.from_function_result(call_id="external-call", result="snapshot-visible result")], + ) + + with pytest.raises( + RuntimeError, + match="Cannot reconstruct service-managed history without an authoritative FIDES context label", + ): + config._fold_uncovered_replay_messages( + session, + [replayed], + complete_replay_messages=[replayed], + replay_history_complete=False, + ) + + assert "context_label" not in session.state[config.source_id] + async def test_hidden_history_placeholder_does_not_taint_restored_context(self) -> None: """A framework-owned hidden placeholder stays integrity-clean on a later run.""" executed: list[str] = []