From 74b7133182e4fff776dbcf47e62bca6afcd72015 Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Fri, 30 Oct 2026 10:53:11 +0100 Subject: [PATCH] Python: preserve FIDES AG-UI replay provenance --- .../ag-ui/agent_framework_ag_ui/_agent_run.py | 91 +- .../_message_adapters.py | 103 +- .../ag-ui/tests/ag_ui/test_endpoint.py | 1001 +++++++++++++++++ .../packages/core/agent_framework/security.py | 68 +- python/packages/core/tests/test_security.py | 35 + 5 files changed, 1270 insertions(+), 28 deletions(-) diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py b/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py index 436607114a9..fdd5f531ba9 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 @@ -90,7 +90,12 @@ 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, + _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 +1769,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_call_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 +1867,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_call_ids.add(pending_entry.identity.call_id) if ( (primary_response.approved is True or pending_entry.owner is ApprovalExecutionOwner.DEFERRED) and lifecycle is not None @@ -1902,6 +1909,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: @@ -2075,12 +2090,19 @@ async def execute_local_call(approval: Content = approval, call_id: str = call_i # 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, + authenticated_rejection_call_ids=authenticated_rejection_call_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], + *, + authenticated_rejection_call_ids: set[str] | 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 @@ -2095,6 +2117,13 @@ def _convert_approval_results_to_tool_messages(messages: list[Message]) -> None: result: list[Message] = [] for msg in messages: + if authenticated_rejection_call_ids and any( + content.type == "function_result" + and content.call_id in authenticated_rejection_call_ids + and content.result == "Error: Tool call invocation was rejected by user." + for content in msg.contents + ): + _mark_backend_snapshot_origin(msg) if get_role_value(msg) != "user": result.append(msg) continue @@ -2113,7 +2142,8 @@ 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]) + result.append(tool_message) # Then user message with remaining content (if any) if other_contents: @@ -2801,6 +2831,33 @@ 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 _fold_fides_uncovered_replay_messages( + agent: SupportsAgentRun, + session: AgentSession, + messages: Sequence[Message], + *, + complete_replay_messages: Sequence[Message] | None = None, +) -> None: + """Fold uncovered replay into FIDES state before a snapshot-only return.""" + for provider in getattr(agent, "context_providers", ()): + fold = getattr(provider, "_fold_uncovered_replay_messages", None) + if callable(fold): + fold(session, messages, complete_replay_messages=complete_replay_messages) + + def _legacy_tool_message_approval_resume( messages: Sequence[Mapping[str, Any]], *, @@ -3004,6 +3061,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 +3131,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 +3258,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 +3423,23 @@ async def save_thread_snapshot( ), ) + def fold_fides_replay_before_snapshot_save() -> None: + 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, + sanitize_tool_history=not (config.use_service_session and stored_snapshot is not None), + protected_tool_call_ids=protected_tool_call_ids, + ) + _fold_fides_uncovered_replay_messages( + agent, + session, + messages, + complete_replay_messages=complete_replay_messages, + ) + _mark_fides_state_covered_messages(agent, messages) + authenticated_cancellations = [ _approval_observer_response(occurrence, cancelled=True) for interrupt_id in cancelled_resume_ids @@ -3552,6 +3631,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 +3693,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 +3727,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 +3736,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..1a306b3c59b 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,30 @@ _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 _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 +82,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 +91,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 +140,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 +161,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 +189,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 +203,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 +243,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 +276,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 +298,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 +342,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 +360,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 +706,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 +717,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 +749,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 @@ -801,6 +858,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 +913,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 +924,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 @@ -993,7 +1056,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 +1075,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 +1091,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 +1127,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 +1171,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..ebbf7b35d3a 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_endpoint.py +++ b/python/packages/ag-ui/tests/ag_ui/test_endpoint.py @@ -5375,6 +5375,665 @@ def _is_policy_violation(tool_results: list[str]) -> bool: return any("Policy violation" in result for result in tool_results) +async def test_endpoint_fides_snapshot_provenance_survives_reasoning_normalization( + streaming_chat_client_stub: Any, +) -> None: + """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] = [] + + 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={})], + ) + + 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, + ) + + 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 == [] + 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_fides_synthetic_skipped_result_keeps_snapshot_provenance( + streaming_chat_client_stub: Any, +) -> None: + """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] = [] + + 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={})], + ) + + 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=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, + ) + + 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 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" + + +async def test_endpoint_fides_normal_continuation_reconstructs_missing_context_label( + streaming_chat_client_stub: Any, +) -> None: + """A normal continuation reconstructs backend history when FIDES state is absent.""" + thread_id = "thread-fides-normal-missing-state" + executed: list[str] = [] + + 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 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 + 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], + ) + 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, + agent, + path="/fides-normal-missing-state", + snapshot_store=store, + snapshot_scope_resolver=lambda _request: "test", + keepalive_seconds=None, + ) + + 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) + 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_fides_confirm_changes_folds_fresh_client_result( + streaming_chat_client_stub: Any, +) -> None: + """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 protected_action() -> str: + executed.append("protected") + return "protected result" + + 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], + ) + 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, + } + }, + ), + ) + app = FastAPI() + 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, + ) + + with TestClient(app) as client: + confirm_response = client.post( + "/fides-confirm-client-result", + json={ + "runId": "run-fides-confirm-resume", + "threadId": thread_id, + "messages": [ + { + "id": "client-result-1", + "role": "tool", + "toolCallId": "client-call", + "content": "fresh external client result", + }, + { + "id": "confirm-result-1", + "role": "tool", + "toolCallId": "confirm-call", + "content": json.dumps( + { + "accepted": True, + "steps": [], + "function_call_id": "unknown-confirm-target", + } + ), + }, + ], + }, + ) + + 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"}], + }, + ) + + 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 protected_action() -> str: + executed.append("protected") + return "protected result" + + 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={})], + ) + + 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": "external-call", + "type": "function", + "function": {"name": "client_lookup", "arguments": "{}"}, + } + ], + }, + { + "id": "external-result", + "role": "tool", + "toolCallId": "external-call", + "content": "stored external result", + }, + { + "id": "assistant-confirm-call", + "role": "assistant", + "content": "", + "toolCalls": [ + { + "id": "confirm-call", + "type": "function", + "function": { + "name": "confirm_changes", + "arguments": json.dumps({"function_call_id": "unknown-confirm-target"}), + }, + } + ], + }, + ], + 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, + 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 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": f"run-fides-confirm-missing-protected-{confirmation_transport}", + "threadId": thread_id, + "messages": [{"id": "user-2", "role": "user", "content": "Run the protected action"}], + }, + ) + + 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("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"]) @@ -7879,6 +8538,168 @@ async def test_endpoint_collected_approval_preserves_original_expiration( assert executed == [] +@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 + + def guarded_action() -> str: + executed.append("guarded") + return "guarded result" + + def guarded_action_two() -> str: + executed.append("guarded-two") + return "guarded result two" + + 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]: + 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 + 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( + 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}, + } + ] + + 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"}], + }, + ) + + 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 + + async def test_endpoint_agent_approval_cancel_discards_queued_tool_approval(streaming_chat_client_stub): """Cancelling a queued approval batch must not replay stale approval prompts on the next user turn.""" client, executed, messages_received, state, _ = _build_tool_approval_queue_endpoint(streaming_chat_client_stub) @@ -12749,6 +13570,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/core/agent_framework/security.py b/python/packages/core/agent_framework/security.py index e4bb70d856c..83909b1693c 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,47 @@ 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, + ) -> 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 + 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 +4520,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..6c7ec046fe6 100644 --- a/python/packages/core/tests/test_security.py +++ b/python/packages/core/tests/test_security.py @@ -3699,6 +3699,41 @@ 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) + 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] = []