From b7c0e04f6b21aa1f942da0bbc599c98ee67998fb Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Mon, 7 Sep 2026 21:09:36 +0200 Subject: [PATCH 1/3] Python: bound and recover policy approvals Bound FIDES policy approvals per session with FIFO and TTL expiry, clean authenticated non-grants by occurrence, and persist visible replacement approvals so stale grants require a safe second approval.\n\nCo-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../specs/004-python-function-calling-loop.md | 7 + .../core/agent_framework/_middleware.py | 31 +++ .../packages/core/agent_framework/_tools.py | 12 + .../packages/core/agent_framework/security.py | 161 +++++++++++++- .../core/test_function_invocation_logic.py | 19 ++ .../tests/core/test_harness_tool_approval.py | 203 +++++++++++++++++ python/packages/core/tests/test_security.py | 205 +++++++++++++++++- 7 files changed, 627 insertions(+), 11 deletions(-) diff --git a/docs/specs/004-python-function-calling-loop.md b/docs/specs/004-python-function-calling-loop.md index e96dd9d8eee..90f6d8ab28b 100644 --- a/docs/specs/004-python-function-calling-loop.md +++ b/docs/specs/004-python-function-calling-loop.md @@ -412,6 +412,10 @@ that manually replay messages own the equivalent rule: do not resend an approval - If policy middleware detects that the exact resolved invocation changed after approval, the old response executes nothing and yields a caller-visible, session-persisted replacement request for the same occurrence; execution requires a second approval and happens exactly once. +- If session-bound middleware no longer holds the reviewed authority because it expired or was evicted, the matched + response executes nothing and produces a replacement approval request with the same occurrence identity. The + replacement is caller-visible, becomes the authoritative pending session snapshot, and requires a second approval + before the tool can execute. Rejection or cancellation releases only the matching occurrence in the owning session. - Unmatched occurrence-aware responses leave the pending request intact for a corrected retry and produce an observable warning/log. A nested `call_id` is never accepted as an occurrence-identity alias. - Session-backed pending snapshots are trusted host state and require tenant-scoped, authorized storage. Consume-on-bind @@ -508,6 +512,8 @@ that manually replay messages own the equivalent rule: do not resend an approval | Approval re-entry after iteration budget | Pending approved calls resolve once even when prior model calls consumed `max_iterations`. | `packages/core/tests/core/test_harness_tool_approval.py::test_auto_approval_resolves_after_iteration_budget_is_exhausted` | | Approval resume with reasoning | Model-bound resume history retains reasoning before the call and terminal result in both modes. | `packages/core/tests/core/test_harness_tool_approval.py::test_approval_resume_replays_reasoning_with_function_call_group` | | Changed resolved policy invocation | The stale decision executes nothing; a same-occurrence replacement request is visible and persisted in both modes, and the second approval executes exactly once. | `packages/core/tests/core/test_harness_tool_approval.py::test_changed_hidden_snapshot_requires_visible_second_approval` | +| Expired or evicted policy authority | The old response executes nothing, surfaces and persists a same-occurrence replacement request in both modes, and executes exactly once only after the replacement is approved; model history remains balanced and stale replay is inert. | `packages/core/tests/core/test_harness_tool_approval.py::test_policy_reapproval_is_visible_persisted_and_executes_once` | +| Session-bound policy cleanup | FIFO/TTL lifecycle and authenticated rejection/cancellation cleanup use occurrence identity within only the owning session. | `packages/core/tests/test_security.py::TestPolicyEnforcementMiddleware::test_pending_policy_approvals_are_fifo_bounded_by_occurrence`, `test_pending_policy_approval_ttl_is_deterministic_and_durable`, `test_non_grant_cleanup_is_authenticated_session_and_occurrence_bound` | | Session-bound substituted response | A response is rebound to the immutable recorded call and cannot replace its call id, tool name, or arguments. | `packages/core/tests/core/test_function_invocation_logic.py::test_session_approval_binding_rebinds_consumes_and_rejects_duplicates` | | Occurrence-aware local binding | New local requests use `function_call.id`; missing, mismatched, or stale occurrence ids do not execute or consume pending state, while the canonical occurrence id binds without an embedded call. | `test_occurrence_aware_approval_rejects_stale_reused_call_id_response`, `test_occurrence_aware_approval_mismatched_identity_does_not_consume_pending`, `test_occurrence_aware_approval_binds_without_embedded_function_call` | | Legacy stored approval | A serialized pending request without `function_call.id` retains exact request-id binding once and warns only when resumed. | `test_legacy_serialized_pending_approval_resumes_once_with_migration_warning`, `packages/core/tests/core/test_types.py::test_legacy_function_call_deserialization_does_not_generate_an_occurrence_id` | @@ -525,6 +531,7 @@ that manually replay messages own the equivalent rule: do not resend an approval | Reused id after completion | A later round with the same id creates a second valid pair. | `test_replace_approval_contents_with_results_allows_reused_call_id_after_completion` | | Replayed approval wrapper | A duplicated wrapper does not restore another function call. | `test_replace_approval_contents_with_results_deduplicates_replayed_approval_request` | | Historical resolved response plus new round | The old response is removed from normalized input and is not converted into a rejection result. | `test_replace_approval_contents_with_results_ignores_already_resolved_response` | +| Replacement request with reused occurrence id | A request after a stale response starts a new unanswered round rather than inheriting the old decision. | `test_collect_unanswered_approval_requests_tracks_replacement_request` | | Multiple reused-id rounds | Approved and rejected rounds retain separate call/result occurrences. | `test_replace_approval_contents_with_results_correlates_reused_call_id_occurrences` | | Multi-content result with reused id | Every content produced by one execution stays with that approval occurrence and cannot bleed into the next reused-id round. | `test_replace_approval_contents_with_results_keeps_multi_content_group_with_reused_call_id` | | Follow-up request closes one occurrence | A user-input follow-up consumes only the preceding approval authority and leaves a later reused-id response pending. | `test_collect_approval_responses_consumes_matching_follow_up_request_occurrence` | diff --git a/python/packages/core/agent_framework/_middleware.py b/python/packages/core/agent_framework/_middleware.py index ba9efcf54af..9f160c3dd12 100644 --- a/python/packages/core/agent_framework/_middleware.py +++ b/python/packages/core/agent_framework/_middleware.py @@ -19,6 +19,7 @@ AgentRunInputs, ChatResponse, ChatResponseUpdate, + Content, Message, ResponseStream, normalize_messages, @@ -755,6 +756,25 @@ async def process( """ ... + def on_approval_responses( + self, + responses: Sequence[Content], + *, + session: AgentSession | None, + ) -> None: + """Observe authenticated approval responses that do not execute a function. + + The function loop calls this only after binding responses to the active session's + authoritative pending snapshot. Stateful middleware can discard rejected or + cancelled authority here; the default implementation retains no state. + + Args: + responses: Session-rebound approval responses. + + Keyword Args: + session: The active invocation session, if any. + """ + class ChatMiddleware(ABC): """Abstract base class for chat middleware that can intercept chat client requests. @@ -1218,6 +1238,17 @@ def matches(self, middleware: Sequence[FunctionMiddlewareTypes]) -> bool: """Return whether this pipeline was built from the provided middleware sequence.""" return self._source_middleware == tuple(middleware) + def notify_approval_responses( + self, + responses: Sequence[Content], + *, + session: AgentSession | None, + ) -> None: + """Notify class-based middleware of authenticated non-executing decisions.""" + for middleware in self._middleware: + if isinstance(middleware, FunctionMiddleware): + middleware.on_approval_responses(responses, session=session) + def _register_middleware(self, middleware: FunctionMiddlewareTypes) -> None: """Register a function middleware item. diff --git a/python/packages/core/agent_framework/_tools.py b/python/packages/core/agent_framework/_tools.py index 5fa270ab807..90748a5fabf 100644 --- a/python/packages/core/agent_framework/_tools.py +++ b/python/packages/core/agent_framework/_tools.py @@ -3239,6 +3239,7 @@ async def _resolve_approval_responses( max_errors: int, execute_function_calls: _FunctionCallExecutor, invocation_session: AgentSession | None = None, + middleware_pipeline: FunctionMiddlewarePipeline | None = None, settle_dangling_calls: Callable[[Sequence[Content]], Awaitable[None]] | None = None, ) -> _FunctionProcessingResult: """Resolve inbound approval responses before the next model call. @@ -3281,6 +3282,11 @@ async def _resolve_approval_responses( responses_to_execute = [ response for response in pending_approval_responses.values() if _is_approval_granted(response.approved) ] + responses_not_granted = [ + response for response in pending_approval_responses.values() if not _is_approval_granted(response.approved) + ] + if middleware_pipeline is not None and responses_not_granted: + middleware_pipeline.notify_approval_responses(responses_not_granted, session=invocation_session) execution_result_groups: list[list[Content]] = [] should_terminate = False reached_error_limit = False @@ -3577,6 +3583,7 @@ async def _get_response_with_function_invocation( invocation_session: AgentSession | None, budget_state: dict[str, Any], max_errors: int, + middleware_pipeline: FunctionMiddlewarePipeline | None = None, ) -> ChatResponse[Any]: """Run the non-streaming function invocation loop.""" from ._middleware import MiddlewareFailure @@ -3622,6 +3629,7 @@ async def settle_approval_replay_calls(function_calls: Sequence[Content]) -> Non max_errors=max_errors, execute_function_calls=execute_function_calls, invocation_session=invocation_session, + middleware_pipeline=middleware_pipeline, settle_dangling_calls=settle_approval_replay_calls, ) function_call_messages.extend(approval_processing.response_messages) @@ -3762,6 +3770,7 @@ async def _stream_response_with_function_invocation( invocation_session: AgentSession | None, budget_state: dict[str, Any], max_errors: int, + middleware_pipeline: FunctionMiddlewarePipeline | None = None, ) -> AsyncIterable[ChatResponseUpdate]: """Run the streaming function invocation loop.""" from ._middleware import MiddlewareFailure @@ -3804,6 +3813,7 @@ async def settle_approval_replay_calls(function_calls: Sequence[Content]) -> Non max_errors=max_errors, execute_function_calls=execute_function_calls, invocation_session=invocation_session, + middleware_pipeline=middleware_pipeline, settle_dangling_calls=settle_approval_replay_calls, ) errors_in_a_row = approval_processing.errors_in_a_row @@ -4175,6 +4185,7 @@ def get_response( invocation_session=invocation_session, budget_state=budget_state, max_errors=max_errors, + middleware_pipeline=function_middleware_pipeline, ) response_format = mutable_options.get("response_format") @@ -4191,6 +4202,7 @@ def get_response( invocation_session=invocation_session, budget_state=budget_state, max_errors=max_errors, + middleware_pipeline=function_middleware_pipeline, ), finalizer=partial(ChatResponse.from_updates, output_format_type=response_format), ) diff --git a/python/packages/core/agent_framework/security.py b/python/packages/core/agent_framework/security.py index 411cb096969..8f4661ab3eb 100644 --- a/python/packages/core/agent_framework/security.py +++ b/python/packages/core/agent_framework/security.py @@ -21,11 +21,12 @@ import logging import math import re +import time import uuid -from collections.abc import Awaitable, Callable, Mapping, MutableMapping +from collections.abc import Awaitable, Callable, Mapping, MutableMapping, Sequence from contextvars import ContextVar, Token from copy import copy, deepcopy -from datetime import datetime +from datetime import datetime, timedelta from enum import Enum from typing import TYPE_CHECKING, Annotated, Any, NamedTuple, NoReturn, cast @@ -1947,7 +1948,20 @@ def get_current_middleware() -> LabelTrackingFunctionMiddleware | None: class _PendingPolicyApproval(NamedTuple): - """Exact, durable binding for one pending policy approval.""" + """Immutable binding record for a pending policy-violation approval. + + Captures every dimension a granted approval is bound to so a reused ``call_id`` cannot + re-authorize a call that differs in any of them. ``body_signature`` covers the function name and + the arguments as displayed for review (variable placeholders unexpanded); + ``resolved_signature`` the exact resolved snapshot the tool would actually receive, so an + approval granted while a placeholder resolved to one payload cannot authorize a replay in + which it resolves to something else (or no longer resolves at all); ``label_key`` the + conversation label shown for review and ``effective_label_key`` the label of everything the + invocation acts on, including hidden arguments; ``session_key`` the session the approval was + requested in; and ``disclosed_violations`` the canonical risks shown to the user. + ``created_at`` is a wall-clock timestamp so TTL expiration survives session serialization and + process restarts. Records remain isolated in the session-scoped security state. + """ body_signature: str resolved_signature: str @@ -1955,6 +1969,7 @@ class _PendingPolicyApproval(NamedTuple): effective_label_key: str session_key: str disclosed_violations: tuple[str, ...] + created_at: float def to_state(self) -> dict[str, Any]: return { @@ -1964,16 +1979,38 @@ def to_state(self) -> dict[str, Any]: "effective_label_key": self.effective_label_key, "session_key": self.session_key, "disclosed_violations": list(self.disclosed_violations), + "created_at": self.created_at, } + def binding_key(self) -> tuple[str, str, str, str, str, tuple[str, ...]]: + """Return every authorization dimension except lifecycle metadata.""" + return ( + self.body_signature, + self.resolved_signature, + self.label_key, + self.effective_label_key, + self.session_key, + self.disclosed_violations, + ) + @classmethod def from_state(cls, payload: Any) -> _PendingPolicyApproval | None: + """Rebuild a record from session state, or fail closed when malformed.""" if not isinstance(payload, dict): return None record = cast(dict[str, Any], payload) - keys = ("body_signature", "resolved_signature", "label_key", "effective_label_key", "session_key") - values = tuple(record.get(key) for key in keys) - violations = record.get("disclosed_violations") + try: + values = ( + record["body_signature"], + record["resolved_signature"], + record["label_key"], + record["effective_label_key"], + record["session_key"], + ) + violations = record["disclosed_violations"] + created_at = record["created_at"] + except KeyError: + return None if not all(type(value) is str for value in values): return None if not isinstance(violations, list): @@ -1981,8 +2018,22 @@ def from_state(cls, payload: Any) -> _PendingPolicyApproval | None: violation_items = cast(list[Any], violations) if not all(type(item) is str for item in violation_items): return None + if type(created_at) not in (int, float) or not math.isfinite(created_at): + return None typed_values = cast(tuple[str, str, str, str, str], values) - return cls(*typed_values, tuple(cast(list[str], violation_items))) + return cls( + body_signature=typed_values[0], + resolved_signature=typed_values[1], + label_key=typed_values[2], + effective_label_key=typed_values[3], + session_key=typed_values[4], + disclosed_violations=tuple(cast(list[str], violation_items)), + created_at=float(created_at), + ) + + +_DEFAULT_MAX_PENDING_APPROVALS = 256 +_DEFAULT_PENDING_APPROVAL_TTL = timedelta(hours=1) @experimental(feature_id=ExperimentalFeature.FIDES) @@ -2026,14 +2077,40 @@ def __init__( enable_audit_log: bool = True, approval_on_violation: bool = False, *, + max_pending_approvals: int = _DEFAULT_MAX_PENDING_APPROVALS, + pending_approval_ttl: timedelta | None = _DEFAULT_PENDING_APPROVAL_TTL, security_scope: _SecurityScope | None = None, session_state_key: str = _STANDALONE_SESSION_STATE_KEY, ) -> None: - """Initialize policy enforcement and bind its security-state selector.""" + """Initialize PolicyEnforcementFunctionMiddleware. + + Args: + allow_untrusted_tools: Set of tool names allowed to execute in an untrusted context. + block_on_violation: Whether to block execution on policy violations. + Ignored if approval_on_violation is True. + enable_audit_log: Whether to maintain an audit log of violations. + approval_on_violation: Whether to request user approval instead of blocking + when a policy violation is detected. If True, the middleware will return + a special result that triggers an approval request in the UI. After user + approval, the tool will execute with a warning about untrusted context. + + Keyword Args: + max_pending_approvals: Maximum pending policy approvals retained per security scope. + When the scope reaches this bound, the oldest occurrence is evicted first. + pending_approval_ttl: Maximum age of an unconsumed approval. ``None`` disables expiry. + security_scope: Internal fixed scope used by the context-provider path. + session_state_key: Internal session-state key shared by reusable middleware. + """ + if isinstance(max_pending_approvals, bool) or max_pending_approvals < 1: + raise ValueError("max_pending_approvals must be at least 1.") + if pending_approval_ttl is not None and pending_approval_ttl <= timedelta(0): + raise ValueError("pending_approval_ttl must be positive or None.") self.allow_untrusted_tools = allow_untrusted_tools or set() self.approval_on_violation = approval_on_violation self.block_on_violation = block_on_violation if not approval_on_violation else False self.enable_audit_log = enable_audit_log + self._max_pending_approvals = max_pending_approvals + self._pending_approval_ttl = pending_approval_ttl self._initialize_security_scope(security_scope, session_state_key=session_state_key) def _clone_for_scope(self, scope: _SecurityScope) -> PolicyEnforcementFunctionMiddleware: @@ -2051,11 +2128,32 @@ def audit_log(self) -> list[dict[str, Any]]: def _pending_policy_approvals(self) -> dict[str, Any]: return self._scope.pending_approvals + def _prune_pending_approvals(self, scope: _SecurityScope | None = None) -> None: + """Expire malformed or old records and enforce FIFO capacity in one scope.""" + pending_approvals = (scope or self._scope).pending_approvals + now = time.time() + ttl_seconds = self._pending_approval_ttl.total_seconds() if self._pending_approval_ttl is not None else None + for approval_id, payload in list(pending_approvals.items()): + record = _PendingPolicyApproval.from_state(payload) + if record is None or (ttl_seconds is not None and now - record.created_at >= ttl_seconds): + pending_approvals.pop(approval_id, None) + while len(pending_approvals) > self._max_pending_approvals: + evicted_id = next(iter(pending_approvals)) + pending_approvals.pop(evicted_id, None) + logger.debug("Evicted oldest pending policy approval occurrence %s.", evicted_id) + def _get_pending_approval(self, approval_id: str) -> _PendingPolicyApproval | None: + """Return the live stored binding record for *approval_id*.""" + self._prune_pending_approvals() return _PendingPolicyApproval.from_state(self._scope.pending_approvals.get(approval_id)) def _store_pending_approval(self, approval_id: str, record: _PendingPolicyApproval) -> None: - self._scope.pending_approvals[approval_id] = record.to_state() + """Persist a record as the newest occurrence and enforce the scope bound.""" + self._prune_pending_approvals() + pending_approvals = self._scope.pending_approvals + pending_approvals.pop(approval_id, None) + pending_approvals[approval_id] = record.to_state() + self._prune_pending_approvals() def _get_call_id(self, context: FunctionInvocationContext) -> str: """Get the tool call id for this invocation context.""" @@ -2154,6 +2252,7 @@ def _pending_record( effective_label_key=self._effective_label_key(context), session_key=self._session_key(context), disclosed_violations=self._violation_set_key(violations), + created_at=time.time(), ) def _signature_from_function_call(self, function_call: Any) -> str | None: @@ -2198,10 +2297,37 @@ def _matches_pending_approval( and approval_response.approved is True ): return False - return current_binding == pending and self._response_matches_pending( + return current_binding.binding_key() == pending.binding_key() and self._response_matches_pending( approval_response, approval_id, call_id, pending.body_signature ) + def on_approval_responses( + self, + responses: Sequence[Content], + *, + session: AgentSession | None, + ) -> None: + """Discard authenticated non-grants from only their owning security scope.""" + scope = self._scope_for_session(session) + self._prune_pending_approvals(scope) + session_key = session.session_id if session is not None else "" + for response in responses: + if response.type != "function_approval_response" or response.approved is True or response.id is None: + continue + function_call = response.function_call + if function_call is None or function_call.call_id is None: + continue + pending = _PendingPolicyApproval.from_state(scope.pending_approvals.get(response.id)) + if pending is None or pending.session_key != session_key: + continue + if self._response_matches_pending( + response, + response.id, + function_call.call_id, + pending.body_signature, + ): + scope.pending_approvals.pop(response.id, None) + def _consume_pending_approval(self, context: FunctionInvocationContext) -> None: self._pending_policy_approvals.pop(self._get_approval_id(context), None) @@ -2601,6 +2727,9 @@ def __init__( enable_policy_enforcement: bool = True, quarantine_chat_client: SupportsChatGetResponse | None = None, source_id: str | None = None, + *, + max_pending_approvals: int = _DEFAULT_MAX_PENDING_APPROVALS, + pending_approval_ttl: timedelta | None = _DEFAULT_PENDING_APPROVAL_TTL, ) -> None: """Initialize secure agent configuration. @@ -2626,7 +2755,15 @@ def __init__( class docstring for details on running multiple instances. source_id: Optional source identifier for context provider attribution. Defaults to "secure_agent". + + Keyword Args: + max_pending_approvals: Maximum pending policy approvals retained per session. + pending_approval_ttl: Maximum age of an unconsumed approval. ``None`` disables expiry. """ + if isinstance(max_pending_approvals, bool) or max_pending_approvals < 1: + raise ValueError("max_pending_approvals must be at least 1.") + if pending_approval_ttl is not None and pending_approval_ttl <= timedelta(0): + raise ValueError("pending_approval_ttl must be positive or None.") super().__init__(source_id or self.DEFAULT_SOURCE_ID) self._auto_hide_untrusted = auto_hide_untrusted self._default_integrity = default_integrity @@ -2637,6 +2774,8 @@ class docstring for details on running multiple instances. self._block_on_violation = block_on_violation self._approval_on_violation = approval_on_violation self._enable_audit_log = enable_audit_log + self._max_pending_approvals = max_pending_approvals + self._pending_approval_ttl = pending_approval_ttl self.enable_policy_enforcement = enable_policy_enforcement self.label_tracker = LabelTrackingFunctionMiddleware( auto_hide_untrusted=auto_hide_untrusted, @@ -2649,6 +2788,8 @@ class docstring for details on running multiple instances. block_on_violation=block_on_violation, approval_on_violation=approval_on_violation, enable_audit_log=enable_audit_log, + max_pending_approvals=max_pending_approvals, + pending_approval_ttl=pending_approval_ttl, ) if enable_policy_enforcement else None diff --git a/python/packages/core/tests/core/test_function_invocation_logic.py b/python/packages/core/tests/core/test_function_invocation_logic.py index 02a0f05b8ae..7e688a2e13c 100644 --- a/python/packages/core/tests/core/test_function_invocation_logic.py +++ b/python/packages/core/tests/core/test_function_invocation_logic.py @@ -3297,6 +3297,25 @@ def test_pending_approval_batch_filter_keeps_resolved_sibling_pair() -> None: ] +def test_collect_unanswered_approval_requests_tracks_replacement_request() -> None: + """A replacement request with the same occurrence id starts a new unanswered round.""" + from agent_framework._tools import _collect_unanswered_approval_requests + + _, original_request, stale_response = _build_approved_tool_roundtrip( + call_id="call_reapproval", + approval_id="approval_occurrence", + tool_name="guarded_tool", + ) + replacement_request = Content.from_dict(original_request.to_dict()) + messages = [ + Message(role="assistant", contents=[original_request]), + Message(role="user", contents=[stale_response]), + Message(role="assistant", contents=[replacement_request]), + ] + + assert _collect_unanswered_approval_requests(messages) == [replacement_request] + + def test_replace_approval_contents_with_results_uses_result_call_ids_without_placeholders() -> None: from agent_framework._tools import _collect_approval_responses, _replace_approval_contents_with_results diff --git a/python/packages/core/tests/core/test_harness_tool_approval.py b/python/packages/core/tests/core/test_harness_tool_approval.py index 032309accd3..04168573e76 100644 --- a/python/packages/core/tests/core/test_harness_tool_approval.py +++ b/python/packages/core/tests/core/test_harness_tool_approval.py @@ -5,6 +5,7 @@ import json import warnings from collections.abc import Awaitable, Callable, MutableSequence +from datetime import timedelta from enum import Enum from pathlib import Path from typing import Any @@ -658,6 +659,208 @@ def guarded_sink(value: str) -> str: assert guarded_values == ["hidden payload"] +@pytest.mark.parametrize("lifecycle_event", ["expiry", "eviction"]) +@pytest.mark.parametrize("streaming", [False, True], ids=["non-streaming", "streaming"]) +async def test_policy_reapproval_is_visible_persisted_and_executes_once( + chat_client_base: MockBaseChatClient, + monkeypatch: pytest.MonkeyPatch, + lifecycle_event: str, + streaming: bool, +) -> None: + """An obsolete policy grant must surface and persist a resumable replacement.""" + now = 1_000.0 + monkeypatch.setattr("agent_framework.security.time.time", lambda: now) + calls = 0 + + @tool(name="policy_guarded_tool") + def policy_guarded_tool() -> str: + nonlocal calls + calls += 1 + return "approved result" + + class MarkUntrusted(FunctionMiddleware): + async def process( + self, + context: FunctionInvocationContext, + call_next: Callable[[], Awaitable[None]], + ) -> None: + context.metadata["context_label"] = ContentLabel(integrity=IntegrityLabel.UNTRUSTED) + await call_next() + + policy = PolicyEnforcementFunctionMiddleware( + approval_on_violation=True, + max_pending_approvals=1, + pending_approval_ttl=timedelta(seconds=10) if lifecycle_event == "expiry" else None, + ) + agent = Agent( + client=chat_client_base, + tools=[policy_guarded_tool], + middleware=[MarkUntrusted(), policy], + context_providers=[InMemoryHistoryProvider()], + ) + session = AgentSession(session_id=f"policy-reapproval-{lifecycle_event}-{streaming}") + function_calls = [ + Content.from_function_call( + call_id="policy-provider-call", + name="policy_guarded_tool", + arguments="{}", + id="policy-approval-occurrence", + ) + ] + if lifecycle_event == "eviction": + function_calls.append( + Content.from_function_call( + call_id="other-provider-call", + name="policy_guarded_tool", + arguments="{}", + id="other-occurrence", + ) + ) + captured_model_calls: list[list[Message]] = [] + + def capture(messages: MutableSequence[Message]) -> None: + captured_model_calls.append([Message.from_dict(message.to_dict()) for message in messages]) + + if streaming: + original_stream = chat_client_base._get_streaming_response + + def capture_stream( + *, + messages: MutableSequence[Message], + options: dict[str, Any], + **kwargs: Any, + ) -> Any: + capture(messages) + return original_stream(messages=messages, options=options, **kwargs) + + monkeypatch.setattr(chat_client_base, "_get_streaming_response", capture_stream) + chat_client_base.streaming_responses = [ + [ChatResponseUpdate(role="assistant", contents=function_calls)], + [ChatResponseUpdate(role="assistant", contents=[Content.from_text("done")])], + [ChatResponseUpdate(role="assistant", contents=[Content.from_text("after stale replay")])], + ] + first_stream = agent.run("run policy tool", stream=True, session=session) + first_updates = [update async for update in first_stream] + first_response = await first_stream.get_final_response() + first_update_types = [content.type for update in first_updates for content in update.contents] + assert first_update_types.count("function_call") == len(function_calls) + assert first_update_types.count("function_approval_request") == len(function_calls) + else: + original_response = chat_client_base._get_non_streaming_response + + async def capture_response( + *, + messages: MutableSequence[Message], + options: dict[str, Any], + **kwargs: Any, + ) -> ChatResponse: + capture(messages) + return await original_response(messages=messages, options=options, **kwargs) + + monkeypatch.setattr(chat_client_base, "_get_non_streaming_response", capture_response) + chat_client_base.run_responses = [ + ChatResponse(messages=Message(role="assistant", contents=function_calls)), + ChatResponse(messages=Message(role="assistant", contents=["done"])), + ChatResponse(messages=Message(role="assistant", contents=["after stale replay"])), + ] + first_response = await agent.run("run policy tool", session=session) + + policy_pending = policy._scope_for_session(session).pending_approvals + first_requests = first_response.user_input_requests + remaining_request: Content | None = None + if lifecycle_event == "expiry": + original_request = first_requests[0] + now = 1_010.0 + else: + assert len(first_requests) == 2 + original_request = next(request for request in first_requests if request.id not in policy_pending) + remaining_request = next(request for request in first_requests if request is not original_request) + assert original_request.id is not None + occurrence_id = original_request.id + assert calls == 0 + assert chat_client_base.call_count == 1 + + stale_approval = original_request.to_function_approval_response(True) + resume_contents = [stale_approval] + if remaining_request is not None: + resume_contents.append(remaining_request.to_function_approval_response(False)) + resume_message = Message(role="user", contents=resume_contents) + + if streaming: + stale_stream = agent.run(resume_message, stream=True, session=session) + stale_updates = [update async for update in stale_stream] + stale_response = await stale_stream.get_final_response() + stale_update_types = [content.type for update in stale_updates for content in update.contents] + assert stale_update_types.count("function_approval_request") == 1 + assert stale_update_types.count("function_result") == int(lifecycle_event == "eviction") + else: + stale_response = await agent.run(resume_message, session=session) + + assert calls == 0 + assert chat_client_base.call_count == 1 + replacement_requests = stale_response.user_input_requests + assert len(replacement_requests) == 1 + replacement = replacement_requests[0] + assert replacement.id != occurrence_id + assert replacement.function_call is not None + assert original_request.function_call is not None + assert replacement.function_call.id == occurrence_id + assert replacement.function_call.call_id == original_request.function_call.call_id + pending_snapshots = session.state["tool_approval"]["pending_approval_requests"] + assert [snapshot["id"] for snapshot in pending_snapshots] == [replacement.id] + + if streaming: + approved_stream = agent.run( + replacement.to_function_approval_response(True), + stream=True, + session=session, + ) + approved_updates = [update async for update in approved_stream] + approved_response = await approved_stream.get_final_response() + assert [content.type for update in approved_updates for content in update.contents] == [ + "function_result", + "text", + ] + else: + approved_response = await agent.run(replacement.to_function_approval_response(True), session=session) + + assert calls == 1 + assert chat_client_base.call_count == 2 + assert [[content.type for content in message.contents] for message in approved_response.messages] == [ + ["function_result"], + ["text"], + ] + model_contents = [content for message in captured_model_calls[-1] for content in message.contents] + model_types = [content.type for content in model_contents] + expected_occurrences = 2 if lifecycle_event == "eviction" else 1 + assert model_types.count("function_call") == expected_occurrences + assert model_types.count("function_result") == expected_occurrences + assert "function_approval_request" not in model_types + assert "function_approval_response" not in model_types + model_calls = [content for content in model_contents if content.type == "function_call"] + model_results = [content for content in model_contents if content.type == "function_result"] + assert {content.call_id for content in model_calls} == {content.call_id for content in model_results} + approved_model_call = next(content for content in model_calls if content.id == occurrence_id) + approved_model_result = next(content for content in model_results if content.call_id == approved_model_call.call_id) + assert approved_model_call.call_id == approved_model_result.call_id + + if streaming: + replay_stream = agent.run(stale_approval, stream=True, session=session) + _ = [update async for update in replay_stream] + await replay_stream.get_final_response() + else: + await agent.run(stale_approval, session=session) + + assert calls == 1 + assert chat_client_base.call_count == 3 + assert "pending_approval_requests" not in session.state["tool_approval"] + replayed_types = [content.type for message in captured_model_calls[-1] for content in message.contents] + assert replayed_types.count("function_call") == expected_occurrences + assert replayed_types.count("function_result") == expected_occurrences + assert "function_approval_request" not in replayed_types + assert "function_approval_response" not in replayed_types + + @pytest.mark.parametrize("approved", [True, False], ids=["approved", "rejected"]) @pytest.mark.parametrize("streaming", [False, True], ids=["non-streaming", "streaming"]) async def test_approval_resume_returns_result_without_mutating_inputs( diff --git a/python/packages/core/tests/test_security.py b/python/packages/core/tests/test_security.py index 32c39381184..c9bdd58e283 100644 --- a/python/packages/core/tests/test_security.py +++ b/python/packages/core/tests/test_security.py @@ -5,6 +5,7 @@ import asyncio import json import logging +from datetime import timedelta from types import SimpleNamespace from typing import Any, cast @@ -22,7 +23,13 @@ SessionContext, ) from agent_framework._middleware import FunctionMiddlewarePipeline, MiddlewareTermination -from agent_framework._tools import FunctionTool, _auto_invoke_function, normalize_function_invocation_configuration +from agent_framework._tools import ( + FunctionTool, + _auto_invoke_function, + _resolve_approval_responses, + _store_pending_approval_requests, + normalize_function_invocation_configuration, +) from agent_framework._types import Content from agent_framework.security import ( ConfidentialityLabel, @@ -754,6 +761,202 @@ async def next_fn() -> None: assert context.result == [Content.from_text("approved result")] assert "call-approved" not in middleware._pending_policy_approvals + async def test_pending_policy_approvals_are_fifo_bounded_by_occurrence(self, mock_function) -> None: + """The oldest occurrence is evicted and its stale grant fails closed.""" + middleware = PolicyEnforcementFunctionMiddleware( + approval_on_violation=True, + max_pending_approvals=2, + pending_approval_ttl=None, + ) + session = AgentSession(session_id="fifo-policy-approvals") + + async def request(occurrence_id: str) -> Content: + context = FunctionInvocationContext( + function=mock_function, + arguments=mock_function.args_schema(arg="test"), + session=session, + kwargs={"session": session}, + ) + context.metadata.update({ + "context_label": ContentLabel(integrity=IntegrityLabel.UNTRUSTED), + "call_id": "reused-provider-call", + "function_call_occurrence_id": occurrence_id, + }) + + async def should_not_execute() -> None: + pytest.fail("Policy-violating tools require approval") + + with pytest.raises(MiddlewareTermination): + await middleware.process(context, should_not_execute) + assert isinstance(context.result, Content) + return context.result + + requests = [await request(f"occurrence-{index}") for index in range(3)] + pending = middleware._scope_for_session(session).pending_approvals + assert list(pending) == ["occurrence-1", "occurrence-2"] + + stale_context = FunctionInvocationContext( + function=mock_function, + arguments=mock_function.args_schema(arg="test"), + session=session, + kwargs={"session": session}, + ) + stale_context.metadata.update({ + "context_label": ContentLabel(integrity=IntegrityLabel.UNTRUSTED), + "call_id": "reused-provider-call", + "function_call_occurrence_id": "occurrence-0", + "approval_response": requests[0].to_function_approval_response(True), + }) + + async def should_not_execute_stale_grant() -> None: + pytest.fail("An evicted approval must not execute") + + with pytest.raises(MiddlewareTermination): + await middleware.process(stale_context, should_not_execute_stale_grant) + + assert isinstance(stale_context.result, Content) + assert stale_context.result.type == "function_approval_request" + assert stale_context.result.id != "occurrence-0" + assert stale_context.result.function_call is not None + assert stale_context.result.function_call.id == "occurrence-0" + assert list(pending) == ["occurrence-2", "occurrence-0"] + + async def test_pending_policy_approval_ttl_is_deterministic_and_durable( + self, + mock_function, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + """A restored approval expires at the configured boundary and is replaced.""" + now = 1_000.0 + monkeypatch.setattr("agent_framework.security.time.time", lambda: now) + middleware = PolicyEnforcementFunctionMiddleware( + approval_on_violation=True, + pending_approval_ttl=timedelta(seconds=5), + ) + session = AgentSession(session_id="ttl-policy-approval") + request_context = FunctionInvocationContext( + function=mock_function, + arguments=mock_function.args_schema(arg="test"), + session=session, + kwargs={"session": session}, + ) + request_context.metadata.update({ + "context_label": ContentLabel(integrity=IntegrityLabel.UNTRUSTED), + "call_id": "ttl-provider-call", + "function_call_occurrence_id": "ttl-occurrence", + }) + + async def should_not_execute() -> None: + pytest.fail("Policy-violating tools require approval") + + with pytest.raises(MiddlewareTermination): + await middleware.process(request_context, should_not_execute) + assert isinstance(request_context.result, Content) + approval_request = request_context.result + + restored = AgentSession.from_dict(session.to_dict()) + now = 1_005.0 + replay_context = FunctionInvocationContext( + function=mock_function, + arguments=mock_function.args_schema(arg="test"), + session=restored, + kwargs={"session": restored}, + ) + replay_context.metadata.update({ + "context_label": ContentLabel(integrity=IntegrityLabel.UNTRUSTED), + "call_id": "ttl-provider-call", + "function_call_occurrence_id": "ttl-occurrence", + "approval_response": approval_request.to_function_approval_response(True), + }) + + with pytest.raises(MiddlewareTermination): + await middleware.process(replay_context, should_not_execute) + + assert isinstance(replay_context.result, Content) + assert replay_context.result.type == "function_approval_request" + assert replay_context.result.id != "ttl-occurrence" + assert replay_context.result.function_call is not None + assert replay_context.result.function_call.id == "ttl-occurrence" + pending = middleware._scope_for_session(restored).pending_approvals["ttl-occurrence"] + assert pending["created_at"] == now + + @pytest.mark.parametrize("cancelled", [False, True], ids=["rejected", "cancelled"]) + async def test_non_grant_cleanup_is_authenticated_session_and_occurrence_bound( + self, + mock_function, + cancelled: bool, + ) -> None: + """Only a rebound non-grant clears its occurrence in the owning session.""" + middleware = PolicyEnforcementFunctionMiddleware(approval_on_violation=True) + owner = AgentSession(session_id="policy-owner") + interleaved = AgentSession(session_id="policy-interleaved") + + async def request(session: AgentSession, occurrence_id: str) -> Content: + context = FunctionInvocationContext( + function=mock_function, + arguments=mock_function.args_schema(arg="test"), + session=session, + kwargs={"session": session}, + ) + context.metadata.update({ + "context_label": ContentLabel(integrity=IntegrityLabel.UNTRUSTED), + "call_id": "shared-provider-call", + "function_call_occurrence_id": occurrence_id, + }) + + async def should_not_execute() -> None: + pytest.fail("Policy-violating tools require approval") + + with pytest.raises(MiddlewareTermination): + await middleware.process(context, should_not_execute) + assert isinstance(context.result, Content) + return context.result + + owner_request = await request(owner, "shared-occurrence") + await request(owner, "owner-second-occurrence") + interleaved_request = await request(interleaved, "shared-occurrence") + owner_pending = middleware._scope_for_session(owner).pending_approvals + interleaved_pending = middleware._scope_for_session(interleaved).pending_approvals + + _store_pending_approval_requests(interleaved, [interleaved_request]) + forged = interleaved_request.to_function_approval_response(False) + forged.id = "unissued-occurrence" + + async def should_not_execute_responses(**_kwargs: Any) -> Any: + pytest.fail("Non-grants must not execute tools") + + await _resolve_approval_responses( + prepared_messages=[Message(role="user", contents=[forged])], + options={"tools": [mock_function]}, + errors_in_a_row=0, + max_errors=3, + execute_function_calls=should_not_execute_responses, # type: ignore[arg-type] + invocation_session=interleaved, + middleware_pipeline=FunctionMiddlewarePipeline(middleware), + ) + assert "shared-occurrence" in interleaved_pending + + non_grant = interleaved_request.to_function_approval_response(False) + if cancelled: + non_grant.additional_properties["cancelled"] = True + resolved = await _resolve_approval_responses( + prepared_messages=[Message(role="user", contents=[non_grant])], + options={"tools": [mock_function]}, + errors_in_a_row=0, + max_errors=3, + execute_function_calls=should_not_execute_responses, # type: ignore[arg-type] + invocation_session=interleaved, + middleware_pipeline=FunctionMiddlewarePipeline(middleware), + ) + + assert "shared-occurrence" not in interleaved_pending + assert set(owner_pending) == {"shared-occurrence", "owner-second-occurrence"} + results = [content for message in resolved.response_messages for content in message.contents] + assert len(results) == 1 + assert results[0].type == "function_result" + assert results[0].call_id == "shared-provider-call" + assert owner_request.id == "shared-occurrence" + async def test_auto_invoke_passes_approval_response_to_middleware(self, mock_function): """Test the main tool loop passes approval response content via metadata.""" captured_metadata: dict[str, object] = {} From 51d5da4e7ec3539d933402231897d20ba1a20be6 Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Wed, 9 Sep 2026 18:34:19 +0200 Subject: [PATCH 2/3] Python: address FIDES approval lifecycle review Harden request-generation authority, keep lifecycle observation private, and notify fixed-scope FIDES policy middleware from authenticated AG-UI rejection and cancellation paths.\n\nCo-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../specs/004-python-function-calling-loop.md | 26 +- .../ag-ui/agent_framework_ag_ui/_agent_run.py | 231 +++++++++-- .../_approval_lifecycle.py | 71 +++- .../agent_framework_ag_ui/_approval_state.py | 2 + .../_message_adapters.py | 9 +- .../ag-ui/agent_framework_ag_ui/_utils.py | 2 + .../ag-ui/tests/ag_ui/test_endpoint.py | 388 +++++++++++++++++- .../core/agent_framework/_middleware.py | 44 +- .../packages/core/agent_framework/_tools.py | 9 +- .../packages/core/agent_framework/security.py | 38 +- .../tests/core/test_harness_tool_approval.py | 20 +- python/packages/core/tests/test_security.py | 132 ++++++ 12 files changed, 884 insertions(+), 88 deletions(-) diff --git a/docs/specs/004-python-function-calling-loop.md b/docs/specs/004-python-function-calling-loop.md index 90f6d8ab28b..81019bb5bbe 100644 --- a/docs/specs/004-python-function-calling-loop.md +++ b/docs/specs/004-python-function-calling-loop.md @@ -397,11 +397,13 @@ that manually replay messages own the equivalent rule: do not resend an approval - A tool that requires approval does not execute before an approved response. - With an `AgentSession`, every surfaced local or hosted approval request is stored as an immutable snapshot in one active model batch. A new surfaced batch replaces an abandoned batch instead of accumulating session state. -- New local approval request IDs use the recorded `function_call.id` occurrence identity. Provider-issued hosted - approval request IDs remain unchanged. Duplicate request IDs within one batch are rejected as malformed. -- An inbound response is honored only when its occurrence identity matches the pending server-held snapshot. A new - local response may omit its embedded function call because the snapshot is authoritative; a mismatched embedded - occurrence identity fails closed. +- Initial local approval request IDs use the recorded `function_call.id` occurrence identity. A policy replacement + retains that occurrence identity on the function call but rotates a separate request-generation identity; only a + response bound to the current server-held generation can authorize execution. Provider-issued hosted approval + request IDs remain unchanged. Duplicate request IDs within one batch are rejected as malformed. +- An inbound response is honored only when its occurrence identity and current request generation match the pending + server-held snapshot. A new local response may omit its embedded function call because the snapshot is authoritative; + a mismatched embedded occurrence identity or stale generation fails closed. - Legacy stored local snapshots without `function_call.id` retain exact request-id matching for one consume-on-bind resume and emit a migration warning. Deserialization does not rewrite the snapshot or emit that warning. - Approval requests replayed in inbound message history do not create, replace, or resurrect approval authority. @@ -413,9 +415,11 @@ that manually replay messages own the equivalent rule: do not resend an approval nothing and yields a caller-visible, session-persisted replacement request for the same occurrence; execution requires a second approval and happens exactly once. - If session-bound middleware no longer holds the reviewed authority because it expired or was evicted, the matched - response executes nothing and produces a replacement approval request with the same occurrence identity. The - replacement is caller-visible, becomes the authoritative pending session snapshot, and requires a second approval - before the tool can execute. Rejection or cancellation releases only the matching occurrence in the owning session. + response executes nothing and produces a replacement approval request with the same occurrence identity and a fresh + request generation. The replacement is caller-visible, becomes the authoritative pending session snapshot, and + requires a second approval before the tool can execute; replaying the prior serialized grant cannot authorize it. + Rejection or cancellation releases only the matching occurrence in the owning session, including authenticated + AG-UI lifecycle decisions. - Unmatched occurrence-aware responses leave the pending request intact for a corrected retry and produce an observable warning/log. A nested `call_id` is never accepted as an occurrence-identity alias. - Session-backed pending snapshots are trusted host state and require tenant-scoped, authorized storage. Consume-on-bind @@ -512,8 +516,8 @@ that manually replay messages own the equivalent rule: do not resend an approval | Approval re-entry after iteration budget | Pending approved calls resolve once even when prior model calls consumed `max_iterations`. | `packages/core/tests/core/test_harness_tool_approval.py::test_auto_approval_resolves_after_iteration_budget_is_exhausted` | | Approval resume with reasoning | Model-bound resume history retains reasoning before the call and terminal result in both modes. | `packages/core/tests/core/test_harness_tool_approval.py::test_approval_resume_replays_reasoning_with_function_call_group` | | Changed resolved policy invocation | The stale decision executes nothing; a same-occurrence replacement request is visible and persisted in both modes, and the second approval executes exactly once. | `packages/core/tests/core/test_harness_tool_approval.py::test_changed_hidden_snapshot_requires_visible_second_approval` | -| Expired or evicted policy authority | The old response executes nothing, surfaces and persists a same-occurrence replacement request in both modes, and executes exactly once only after the replacement is approved; model history remains balanced and stale replay is inert. | `packages/core/tests/core/test_harness_tool_approval.py::test_policy_reapproval_is_visible_persisted_and_executes_once` | -| Session-bound policy cleanup | FIFO/TTL lifecycle and authenticated rejection/cancellation cleanup use occurrence identity within only the owning session. | `packages/core/tests/test_security.py::TestPolicyEnforcementMiddleware::test_pending_policy_approvals_are_fifo_bounded_by_occurrence`, `test_pending_policy_approval_ttl_is_deterministic_and_durable`, `test_non_grant_cleanup_is_authenticated_session_and_occurrence_bound` | +| Expired or evicted policy authority | The old response executes nothing, surfaces and persists a same-occurrence replacement with a fresh request generation in both modes, and executes exactly once only after the replacement is approved; restored model history remains balanced and stale prior-generation replay is inert. | `packages/core/tests/core/test_harness_tool_approval.py::test_policy_reapproval_is_visible_persisted_and_executes_once`, `packages/core/tests/test_security.py::TestPolicyEnforcementMiddleware::test_pending_policy_approval_ttl_is_deterministic_and_durable` | +| Session-bound policy cleanup | FIFO/TTL lifecycle and authenticated rejection/cancellation cleanup use occurrence identity within only the owning session, including fixed provider scopes. | `packages/core/tests/test_security.py::TestPolicyEnforcementMiddleware::test_pending_policy_approvals_are_fifo_bounded_by_occurrence`, `test_pending_policy_approval_ttl_is_deterministic_and_durable`, `test_non_grant_cleanup_is_authenticated_session_and_occurrence_bound`, `test_fixed_scope_non_grant_cleanup_keeps_unrelated_occurrence` | | Session-bound substituted response | A response is rebound to the immutable recorded call and cannot replace its call id, tool name, or arguments. | `packages/core/tests/core/test_function_invocation_logic.py::test_session_approval_binding_rebinds_consumes_and_rejects_duplicates` | | Occurrence-aware local binding | New local requests use `function_call.id`; missing, mismatched, or stale occurrence ids do not execute or consume pending state, while the canonical occurrence id binds without an embedded call. | `test_occurrence_aware_approval_rejects_stale_reused_call_id_response`, `test_occurrence_aware_approval_mismatched_identity_does_not_consume_pending`, `test_occurrence_aware_approval_binds_without_embedded_function_call` | | Legacy stored approval | A serialized pending request without `function_call.id` retains exact request-id binding once and warns only when resumed. | `test_legacy_serialized_pending_approval_resumes_once_with_migration_warning`, `packages/core/tests/core/test_types.py::test_legacy_function_call_deserialization_does_not_generate_an_occurrence_id` | @@ -564,7 +568,7 @@ that manually replay messages own the equivalent rule: do not resend an approval | Hosted server boundary | Standing approval does not cross `server_label`. | `test_tool_approval_middleware_standing_rules_include_hosted_server_boundary` | | Argument-scoped rule | Exact arguments are required; empty arguments are not tool-wide. | `test_tool_approval_middleware_always_approve_tool_with_arguments_rule`, `test_tool_approval_middleware_empty_arguments_rule_is_not_tool_wide` | | Provider-injected approval tool | A tool added during `before_run` defers to in-run resolution, executes once, and emits one result. | `packages/ag-ui/tests/ag_ui/test_endpoint.py::test_endpoint_agent_approval_deferred_provider_tool_executes` | -| AG-UI provider boundary | Completed local approval controls from AG-UI request and snapshot replay are absent from raw chat-client input while deferred and hosted approvals keep their respective in-run/provider paths. | `packages/ag-ui/tests/ag_ui/test_endpoint.py::test_endpoint_does_not_forward_resolved_local_approval_control_to_chat_client`, `packages/ag-ui/tests/ag_ui/test_endpoint.py::test_endpoint_agent_approval_deferred_provider_tool_executes`, `packages/ag-ui/tests/ag_ui/test_endpoint.py::test_endpoint_canonical_resume_preserves_hosted_approval_for_provider`, `packages/ag-ui/tests/ag_ui/test_run.py::test_filter_local_approval_responses_for_provider_removes_duplicate_completed_controls`, `packages/ag-ui/tests/ag_ui/test_run.py::test_filter_local_approval_responses_for_provider_pairs_reused_call_ids_by_occurrence`, `packages/ag-ui/tests/ag_ui/test_run.py::test_canonical_hosted_approval_resume_rejects_edited_arguments_without_mutating_pending` | +| AG-UI provider boundary | Completed local approval controls from AG-UI request and snapshot replay are absent from raw chat-client input while deferred and hosted approvals keep their respective in-run/provider paths. Lifecycle-authenticated local approvals preserve current request-generation authority through occurrence rebinding; authenticated rejection/cancellation cleans only the matching fixed provider-scoped policy record, while malformed or unauthorized controls clean nothing. | `packages/ag-ui/tests/ag_ui/test_endpoint.py::test_endpoint_does_not_forward_resolved_local_approval_control_to_chat_client`, `test_endpoint_agent_approval_deferred_provider_tool_executes`, `test_endpoint_canonical_resume_preserves_hosted_approval_for_provider`, `test_endpoint_fides_non_grant_cleans_authenticated_fixed_scope_only`, `test_endpoint_fides_approval_uses_lifecycle_bound_request_generation`, `packages/ag-ui/tests/ag_ui/test_run.py::test_filter_local_approval_responses_for_provider_removes_duplicate_completed_controls`, `test_filter_local_approval_responses_for_provider_pairs_reused_call_ids_by_occurrence`, `test_canonical_hosted_approval_resume_rejects_edited_arguments_without_mutating_pending` | | AG-UI standard approval payload | Agent and workflow tool approvals emit canonical `tool_call` interrupts. `approved` plus full-replacement `editedArgs` executes once and replays idempotently, while legacy `accepted` plus direct partial edits remains supported. Hosted approvals remain decision-only. | `packages/ag-ui/tests/ag_ui/test_endpoint.py::test_endpoint_agent_approval_resume_entry_applies_standard_full_replacement_edited_args`, `test_endpoint_agent_approval_replayed_standard_edited_resume_is_idempotent`, `test_endpoint_agent_approval_resume_entry_applies_edited_arguments`, `test_workflow_endpoint_emits_canonical_tool_approval_interrupt`, `test_workflow_endpoint_accepts_canonical_tool_approval_resume`, `test_workflow_endpoint_applies_canonical_approval_edited_args`, `test_workflow_endpoint_accepts_legacy_partial_approval_edits`, `test_workflow_endpoint_hosted_approval_rejects_argument_edits` | | AG-UI cancellation | A cancelled interrupt executes zero times and completes normally, including an identical retry during retained cancellation state; resolved siblings in the same complete resume still execute once. Workflow cancellation clears both runner correlation and the owning agent executor's pending request so later approvals remain resumable. | `packages/ag-ui/tests/ag_ui/test_endpoint.py::test_endpoint_agent_approval_cancelled_resume_entry_completes_without_execution`, `test_endpoint_agent_approval_replayed_cancellation_completes_idempotently`, `test_endpoint_agent_approval_mixed_cancelled_and_resolved_resume_executes_resolved_tool`, `test_endpoint_workflow_request_info_cancelled_resume_completes_normally`, `test_workflow_endpoint_cancelled_agent_approval_does_not_block_next_approval` | | AG-UI shared workflow interrupt ownership | A direct shared `Workflow` request-info interrupt can only be resolved or cancelled by the Snapshot Scope and AG-UI thread that created it. Ownership follows the authoritative pending request occurrence, and explicitly threaded cold checkpoint resumes fail closed when ownership is unavailable. | `packages/ag-ui/tests/ag_ui/test_endpoint.py::test_endpoint_workflow_request_info_rejects_resume_from_different_thread`, `test_endpoint_workflow_request_info_rejects_resume_from_different_scope`, `test_endpoint_workflow_request_info_rejects_cancellation_from_different_thread`, `test_endpoint_workflow_request_info_remains_owned_after_client_disconnect`, `test_endpoint_workflow_request_info_rejects_unowned_pending_interrupt`, `test_endpoint_workflow_checkpoint_resume_rejects_threaded_resume_after_restart` | 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 a6e257dd9e1..f758e61e422 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 @@ -38,13 +38,20 @@ SupportsAgentRun, WorkflowAgent, ) -from agent_framework._middleware import FunctionMiddlewarePipeline +from agent_framework._middleware import ( + FunctionMiddlewarePipeline, + FunctionMiddlewareTypes, + _as_middleware_list, # pyright: ignore[reportPrivateUsage] + categorize_middleware, +) from agent_framework._tools import ( _ALREADY_APPROVED_APPROVAL_REQUEST_GROUPS_KEY, # type: ignore + _APPROVAL_REQUEST_ID_KEY, # pyright: ignore[reportPrivateUsage] _collect_approval_responses, # type: ignore _get_tool_map, # type: ignore _is_hosted_tool_approval, # type: ignore _replace_approval_contents_with_results, # type: ignore + _store_pending_approval_requests, # pyright: ignore[reportPrivateUsage] _TOOL_APPROVAL_STATE_KEY, # type: ignore _try_execute_function_call_groups, # type: ignore normalize_function_invocation_configuration, @@ -72,7 +79,7 @@ ResumeDecision, ) from ._approval_state import _APPROVAL_SCOPE_INPUT_KEY, InMemoryAGUIApprovalStateStore, approval_state_thread_id -from ._message_adapters import normalize_agui_input_messages +from ._message_adapters import _APPROVAL_DECISION_IS_BOOLEAN_KEY, normalize_agui_input_messages from ._predictive_state import PredictiveStateHandler from ._tooling import collect_server_tools, merge_tools from ._run_common import ( @@ -893,6 +900,7 @@ def _register_server_generated_approval_response( call_id=str(response.function_call.call_id or response_id), name=response.function_call.name, arguments=arguments, + function_call_id=response.function_call.id, aliases=[str(response.function_call.call_id)] if response.function_call.call_id else None, response_id=str(response_id), server_label=_function_call_server_label(response.function_call), @@ -1024,7 +1032,7 @@ def _approval_state_tool_call_ids( call_ids: set[str] = set() for occurrence in approval_state_store.lifecycle.occurrences_for_thread(thread_id=thread_id): call_ids.add(occurrence.identity.call_id) - call_ids.add(occurrence.identity.interrupt_id) + call_ids.add(occurrence.function_call_id) call_ids.update(occurrence.aliases) call_ids.update(_content_tool_call_ids(list(occurrence.already_approved_requests))) stored_state = approval_state_store.get_tool_approval_state(thread_id) @@ -1412,6 +1420,7 @@ def _canonical_approval_resume_messages( "id": response_id, "call_id": pending_entry.identity.call_id, "name": _pending_approval_name(pending_entry) or "", + "function_call_id": pending_entry.function_call_id, "approved": accepted, "arguments": merged_arguments, } @@ -1446,6 +1455,7 @@ def _canonical_approval_resume_messages( call_id=sibling_call_id, name=function_call.name, arguments=sibling_arguments, + function_call_id=function_call.id, response_id=str(response_id), server_label=_function_call_server_label(function_call), ) @@ -1462,6 +1472,7 @@ def _canonical_approval_resume_messages( "id": str(response_id), "call_id": str(function_call.call_id or response_id), "name": function_call.name, + "function_call_id": function_call.id, "approved": True, "arguments": make_json_safe(function_call.parse_arguments() or {}), } @@ -1511,6 +1522,49 @@ def _canonical_approval_resume_messages( return messages, handled_ids, cancelled_ids, None +def _approval_resolution_middleware_pipeline( + agent: SupportsAgentRun, + session: AgentSession, +) -> FunctionMiddlewarePipeline: + """Build the function middleware used for authenticated approval resolution.""" + client = getattr(agent, "client", None) + configured_middleware: list[Any] = [ + *getattr(client, "function_middleware", ()), + *_as_middleware_list(getattr(agent, "middleware", None)), + ] + function_middleware = categorize_middleware(configured_middleware)["function"] + for provider in cast(list[Any], getattr(agent, "context_providers", [])): + provider_middleware = getattr(provider, "_function_middleware_for_approval_resolution", None) + if callable(provider_middleware): + function_middleware.extend(cast("Sequence[FunctionMiddlewareTypes]", provider_middleware(session))) + return FunctionMiddlewarePipeline(*function_middleware) + + +def _approval_observer_response( + occurrence: ApprovalOccurrence, + *, + cancelled: bool = False, +) -> Content: + """Build a trusted non-grant projection from an authenticated lifecycle occurrence.""" + occurrence_id = occurrence.function_call_id + request_id = occurrence.response_id or occurrence_id + function_call = Content.from_function_call( + call_id=occurrence.identity.call_id, + name=occurrence.name, + arguments=occurrence.arguments, + id=occurrence_id, + ) + additional_properties: dict[str, Any] = {_APPROVAL_REQUEST_ID_KEY: request_id} + if cancelled: + additional_properties["cancelled"] = True + return Content.from_function_approval_response( + approved=False, + id=occurrence_id, + function_call=function_call, + additional_properties=additional_properties, + ) + + async def _resolve_approval_responses( messages: list[Any], tools: list[Any], @@ -1519,8 +1573,10 @@ async def _resolve_approval_responses( invocation_session: AgentSession, thread_id: str = "", validated_approved_responses: list[Content] | None = None, + replacement_approval_requests: list[Content] | None = None, *, lifecycle: ApprovalLifecycle, + middleware_pipeline: FunctionMiddlewarePipeline, authorized_executions: dict[ApprovalOccurrenceIdentity, AuthorizedExecution] | None = None, forwarded_executions: ( dict[str, list[tuple[ForwardedPendingToolTransitionOwner, AuthorizedExecution, Content]]] | None @@ -1565,6 +1621,7 @@ async def _resolve_approval_responses( pending_local_response_content_ids = set() pending_response_groups: dict[object, tuple[ApprovalOccurrence, list[Content]]] = {} intents_by_response_content_id: dict[int, AuthorizedExecution] = {} + authenticated_non_grants: list[Content] = [] for response in approval_responses: resp_id = response.id function_call_id = response.function_call.call_id if response.function_call else None @@ -1594,7 +1651,13 @@ async def _resolve_approval_responses( # stale replay controls and must not authorize a malformed fresh one. primary_response = responses[-1] response_content_ids_to_strip.update(id(response) for response in responses[:-1]) - if not isinstance(primary_response.approved, bool): + decision_is_boolean = ( + primary_response.additional_properties.pop( + _APPROVAL_DECISION_IS_BOOLEAN_KEY, isinstance(primary_response.approved, bool) + ) + is True + ) + if not decision_is_boolean: logger.warning( "Treating approval response id=%s as rejected: approved must be a boolean", primary_response.id, @@ -1646,6 +1709,12 @@ async def _resolve_approval_responses( primary_response.function_call.additional_properties["server_label"] = server_label else: primary_response.function_call.additional_properties.pop("server_label", None) + primary_response.function_call.id = pending_entry.function_call_id + primary_response.additional_properties[_APPROVAL_REQUEST_ID_KEY] = ( + pending_entry.response_id or pending_entry.identity.interrupt_id + ) + if decision_is_boolean and primary_response.approved is False and not server_label: + authenticated_non_grants.append(_approval_observer_response(pending_entry)) if ( (primary_response.approved is True or pending_entry.owner is ApprovalExecutionOwner.DEFERRED) and lifecycle is not None @@ -1672,6 +1741,11 @@ async def _resolve_approval_responses( if validated_approved_responses is not None and primary_response.approved is True and not server_label: validated_approved_responses.append(primary_response) + if authenticated_non_grants: + middleware_pipeline._notify_approval_responses( # pyright: ignore[reportPrivateUsage] + authenticated_non_grants, session=invocation_session + ) + if response_content_ids_to_strip: filtered_messages: list[Message] = [] for message in messages: @@ -1753,10 +1827,6 @@ async def forward_hosted_decision(approval: Content = approval) -> list[Content] if static_approved and tools and lifecycle is not None and authorized_executions is not None: client = getattr(agent, "client", None) config = normalize_function_invocation_configuration(getattr(client, "function_invocation_configuration", None)) - middleware_pipeline = FunctionMiddlewarePipeline( - *getattr(client, "function_middleware", ()), - *run_kwargs.get("middleware", ()), - ) tool_kwargs = {k: v for k, v in run_kwargs.items() if k != "options"} for approval in static_approved: function_call = approval.function_call @@ -1786,7 +1856,22 @@ async def execute_local_call(approval: Content = approval, call_id: str = call_i local_owner = LocalPendingToolTransitionOwner(execute_local_call) outcome = await local_owner.execute(intent, lifecycle=lifecycle) - approved_function_result_groups.append(list(outcome.result_group)) + result_group = list(outcome.result_group) + replacement_request = next( + ( + content + for content in result_group + if content.type == "function_approval_request" + and content.additional_properties.get("_replacement_approval_request") is True + ), + None, + ) + if replacement_request is not None and replacement_request.id is not None: + lifecycle.rotate_request_generation(intent, request_id=replacement_request.id) + _store_pending_approval_requests(invocation_session, [replacement_request]) + if replacement_approval_requests is not None: + replacement_approval_requests.append(replacement_request) + approved_function_result_groups.append(result_group) # Normalize one group per static approval and collect only terminal results for TOOL_CALL_RESULT events. # Deferred provider-injected approvals are left in messages for ToolApprovalMiddleware to process. @@ -2425,7 +2510,14 @@ def _legacy_tool_message_approval_resume( "pending local occurrence. Retry with the canonical approval interrupt id." ) continue - interrupt_id = pending_occurrences[0].identity.interrupt_id + pending_occurrence = pending_occurrences[0] + if pending_occurrence.active_interrupt_id != pending_occurrence.identity.interrupt_id: + error = ( + f"Legacy AG-UI tool-message approval call_id '{call_id}' cannot answer a replacement approval. " + "Retry with the current canonical approval interrupt id." + ) + continue + interrupt_id = pending_occurrence.active_interrupt_id if interrupt_id in seen_interrupt_ids: error = ( f"Legacy AG-UI tool-message approval repeats call_id '{call_id}'. " @@ -2506,11 +2598,7 @@ async def run_agent_stream( interrupt_ids=list(stored_pending_approval_interrupt_ids), ) retired_interrupt_ids = { - reconciliation.identity.interrupt_id - if reconciliation.identity is not None - else reconciliation.interrupt_id - for reconciliation in reconciliations - if reconciliation.retire_interrupt + reconciliation.interrupt_id for reconciliation in reconciliations if reconciliation.retire_interrupt } if retired_interrupt_ids: await snapshot_session.clear_interrupts(interrupt_ids=retired_interrupt_ids) @@ -2626,9 +2714,7 @@ async def run_agent_stream( _clear_tool_approval_state(approval_state_store, approval_thread_id) if resume_error_code == "APPROVAL_RESUME_CANCELLED": retired_interrupt_ids = { - reconciliation.identity.interrupt_id - if reconciliation.identity is not None - else reconciliation.interrupt_id + reconciliation.interrupt_id for reconciliation in approval_snapshot_reconciliations if reconciliation.retire_interrupt } @@ -2649,17 +2735,7 @@ async def run_agent_stream( cancelled_workflow_request_ids, tools=tools, ) - if cancelled_resume_ids and handled_resume_ids == cancelled_resume_ids: - yield RunStartedEvent(run_id=run_id, thread_id=thread_id) - _clear_tool_approval_state(approval_state_store, approval_thread_id) - retired_interrupt_ids = { - reconciliation.identity.interrupt_id if reconciliation.identity is not None else reconciliation.interrupt_id - for reconciliation in approval_snapshot_reconciliations - if reconciliation.retire_interrupt - } - await snapshot_session.clear_interrupts(interrupt_ids=retired_interrupt_ids or cancelled_resume_ids) - yield _build_run_finished_event(run_id=run_id, thread_id=thread_id) - return + only_cancelled_resume = bool(cancelled_resume_ids and handled_resume_ids == cancelled_resume_ids) resume_messages = _resume_to_tool_messages(resume_payload, exclude_interrupt_ids=handled_resume_ids) if available_interrupts: logger.debug("Received available interrupts metadata: %s", available_interrupts) @@ -2701,7 +2777,7 @@ async def run_agent_stream( skip_text = response_format is not None # Handle empty messages (emit RunStarted immediately since no agent response) - if not messages: + if not messages and not only_cancelled_resume: logger.warning("No messages provided in AG-UI input") yield RunStartedEvent(run_id=run_id, thread_id=thread_id) yield _build_run_finished_event(run_id=run_id, thread_id=thread_id) @@ -2801,6 +2877,49 @@ async def run_agent_stream( } ) _restore_tool_approval_state(session, approval_state_store, approval_thread_id) + approval_middleware_pipeline = _approval_resolution_middleware_pipeline(agent, session) + + authenticated_cancellations = [ + _approval_observer_response(occurrence, cancelled=True) + for interrupt_id in cancelled_resume_ids + if ( + occurrence := approval_state_store.lifecycle.occurrence_for_alias( + thread_id=approval_thread_id, + interrupt_id=interrupt_id, + ) + ) + is not None + and occurrence.status is ApprovalStatus.CANCELLED + ] + if authenticated_cancellations: + approval_middleware_pipeline._notify_approval_responses( # pyright: ignore[reportPrivateUsage] + authenticated_cancellations, + session=session, + ) + + if only_cancelled_resume: + yield RunStartedEvent(run_id=run_id, thread_id=thread_id) + _clear_tool_approval_state(approval_state_store, approval_thread_id) + retired_interrupt_ids = { + reconciliation.interrupt_id + for reconciliation in approval_snapshot_reconciliations + if reconciliation.retire_interrupt + } + await snapshot_session.clear_interrupts(interrupt_ids=retired_interrupt_ids or cancelled_resume_ids) + if (stored_after_cancellation := snapshot_session.stored) is not None: + await snapshot_session.save( + messages=stored_after_cancellation.messages, + state=stored_after_cancellation.state, + interrupt=stored_after_cancellation.interrupt, + session_state=_safe_serialize_session_continuation_state( + session, + agent, + shared_state_keys=set(flow.current_state).difference(protected_session_state_keys), + include_service_session_id=config.use_service_session, + ), + ) + yield _build_run_finished_event(run_id=run_id, thread_id=thread_id) + return # Inject metadata for AG-UI orchestration (Feature #2: Azure-safe truncation) base_metadata: dict[str, Any] = { @@ -2874,6 +2993,7 @@ async def run_agent_stream( ) return validated_approved_responses: list[Content] = [] + replacement_approval_requests: list[Content] = [] newly_resolved_approval_results = await _resolve_approval_responses( messages, tools_for_execution, @@ -2882,7 +3002,9 @@ async def run_agent_stream( session, approval_thread_id, validated_approved_responses, + replacement_approval_requests, lifecycle=approval_state_store.lifecycle, + middleware_pipeline=approval_middleware_pipeline, authorized_executions=authorized_executions, forwarded_executions=forwarded_executions, ) @@ -2894,6 +3016,35 @@ async def run_agent_stream( if resolved_approval_results or any(message.get("function_approvals") for message in snapshot_messages): _merge_resolved_approval_results_into_snapshot(snapshot_messages, messages) + if replacement_approval_requests: + yield RunStartedEvent(run_id=run_id, thread_id=thread_id) + for request in replacement_approval_requests: + for event in _emit_content( + request, + flow, + predictive_handler, + skip_text, + config.require_confirmation, + ): + yield event + 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) + await snapshot_session.save( + messages=_bound_host_payload_history(_persistable_host_payload_history(persisted_messages)), + state=cast(dict[str, Any], make_json_safe(flow.current_state)) if flow.current_state else None, + interrupt=flow.interrupts or None, + session_state=_safe_serialize_session_continuation_state( + session, + agent, + shared_state_keys=set(flow.current_state).difference(protected_session_state_keys), + include_service_session_id=config.use_service_session, + ), + ) + _save_tool_approval_state(session, approval_state_store, approval_thread_id) + yield _build_run_finished_event(run_id=run_id, thread_id=thread_id, interrupts=flow.interrupts) + return + # Feature #3: Emit StateSnapshotEvent for approved state-changing tools before agent runs approved_state_updates = _extract_approved_state_updates( [Message(role="user", contents=validated_approved_responses)], @@ -3008,6 +3159,21 @@ async def run_agent_stream( for content in update.contents: content_type = getattr(content, "type", None) logger.debug(f"Processing content type={content_type}, message_id={flow.message_id}") + forwarded_reapproval_handled = False + + if ( + content_type == "function_approval_request" + and content.call_id + and content.additional_properties.get("_replacement_approval_request") is True + and (forwarded_queue := forwarded_executions.get(content.call_id)) + ): + forwarded = forwarded_queue.pop(0) + if not forwarded_queue: + forwarded_executions.pop(content.call_id, None) + owner, intent, _ = forwarded + owner.record_reapproval(intent, content, lifecycle=approval_state_store.lifecycle) + _store_pending_approval_requests(session, [content]) + forwarded_reapproval_handled = True if ( content_type == "function_result" @@ -3023,7 +3189,9 @@ async def run_agent_stream( # Register pending approval requests so we can validate responses later if content_type == "function_approval_request": - if content.id and content.function_call and content.function_call.name: + if forwarded_reapproval_handled: + pass + elif content.id and content.function_call and content.function_call.name: server_label = _function_call_server_label(content.function_call) canonical_interrupt_id = _approval_interrupt_id(content) provider_approval_thread_id = approval_state_thread_id( @@ -3048,6 +3216,7 @@ async def run_agent_stream( "request_id": str(content.id), "interrupt_id": str(canonical_interrupt_id), "call_id": str(content.function_call.call_id or canonical_interrupt_id), + "function_call_id": content.function_call.id, "already_approved_requests": already_approved_requests, } approval_state_store.register( diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_approval_lifecycle.py b/python/packages/ag-ui/agent_framework_ag_ui/_approval_lifecycle.py index 485599ce4ed..08c65725c65 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_approval_lifecycle.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_approval_lifecycle.py @@ -179,6 +179,8 @@ class ApprovalOccurrence: name: str arguments: str owner: ApprovalExecutionOwner + function_call_id: str + active_interrupt_id: str scope: str | None = None aliases: tuple[str, ...] = () response_id: str | None = None @@ -276,6 +278,7 @@ def register( call_id: str, name: str, arguments: str, + function_call_id: str | None = None, aliases: list[str] | None = None, response_id: str | None = None, already_approved_requests: list[dict[str, Any]] | None = None, @@ -293,6 +296,7 @@ def register( name=name, arguments=arguments, owner=owner, + function_call_id=function_call_id or interrupt_id, scope=scope, aliases=aliases, response_id=response_id, @@ -311,6 +315,7 @@ def _register_aliases( name: str, arguments: str, owner: ApprovalExecutionOwner, + function_call_id: str, scope: str | None = None, aliases: list[str] | None = None, response_id: str | None = None, @@ -345,6 +350,7 @@ def _register_aliases( or occurrence.name != name or occurrence.arguments != arguments or occurrence.owner is not owner + or occurrence.function_call_id != function_call_id or occurrence.scope != scope or occurrence.response_id != response_id or occurrence.idempotency_key != idempotency_key @@ -374,6 +380,8 @@ def _register_aliases( name=name, arguments=arguments, owner=owner, + function_call_id=function_call_id, + active_interrupt_id=interrupt_id, scope=scope, aliases=occurrence_aliases, response_id=response_id, @@ -431,7 +439,7 @@ def pending_interrupt_ids(self, *, thread_id: str) -> set[str]: with self._index_lock: self._purge_expired_terminal() return { - occurrence.identity.interrupt_id + occurrence.active_interrupt_id for occurrence in self._occurrences.values() if thread_id in occurrence.thread_ids and occurrence.status is ApprovalStatus.PENDING } @@ -466,7 +474,7 @@ def reconcile_snapshot( @staticmethod def _snapshot_reconciliation(occurrence: ApprovalOccurrence) -> ApprovalSnapshotReconciliation: return ApprovalSnapshotReconciliation( - interrupt_id=occurrence.identity.interrupt_id, + interrupt_id=occurrence.active_interrupt_id, identity=occurrence.identity, status=occurrence.status, retire_interrupt=occurrence.status.is_terminal, @@ -796,6 +804,53 @@ def _emit_event( extra["approval_failure_type"] = failure_type logger.info("AG-UI approval lifecycle transition", extra=extra) + @_serialized_by_occurrence + def rotate_request_generation(self, intent: AuthorizedExecution, *, request_id: str) -> None: + """Replace the pending client interrupt alias after local execution defers for reapproval.""" + occurrence = self._occurrences[intent.identity] + if occurrence.status is not ApprovalStatus.PENDING: + raise ApprovalClaimConflictError(f"Approval occurrence is not pending: {occurrence.status}.") + self._rotate_request_generation(occurrence, request_id=request_id) + + @_serialized_by_occurrence + def defer_for_reapproval( + self, + intent: AuthorizedExecution, + *, + owner: ApprovalExecutionOwner, + request_id: str, + ) -> None: + """Return forwarded execution to pending under a fresh request generation.""" + occurrence = self._occurrences[intent.identity] + if intent.owner is not owner or occurrence.owner is not owner: + raise ValueError( + f"Approval occurrence belongs to the {occurrence.owner.value} transition owner, not {owner.value}." + ) + if occurrence.status is not ApprovalStatus.EXECUTING: + raise ApprovalSettlementConflictError(f"Approval occurrence is not executing: {occurrence.status}.") + occurrence.status = ApprovalStatus.PENDING + self._rotate_request_generation(occurrence, request_id=request_id) + + def _rotate_request_generation(self, occurrence: ApprovalOccurrence, *, request_id: str) -> None: + if not request_id: + raise ValueError("A replacement approval request id cannot be empty.") + for thread_id in occurrence.thread_ids: + key = (thread_id, request_id) + existing = self._pending_by_interrupt.get(key) or self._terminal_by_interrupt.get(key) + if existing is not None and existing != occurrence.identity: + raise ValueError("Replacement approval request id conflicts with another occurrence.") + for alias in occurrence.aliases: + if self._pending_by_interrupt.get((thread_id, alias)) == occurrence.identity: + self._pending_by_interrupt.pop((thread_id, alias), None) + occurrence.aliases = (request_id,) + occurrence.active_interrupt_id = request_id + occurrence.response_id = request_id + occurrence.decision = None + occurrence.pending_since = self._clock() + for thread_id in occurrence.thread_ids: + self._pending_by_interrupt[(thread_id, request_id)] = occurrence.identity + self._emit_event("generation_rotation", occurrence) + @_serialized_by_occurrence def begin_execution(self, intent: AuthorizedExecution, *, owner: ApprovalExecutionOwner) -> None: """Mark that an external side effect may begin. @@ -1023,6 +1078,18 @@ async def forward( lifecycle.recover_execution(intent, owner=self._owner) raise + def record_reapproval( + self, + intent: AuthorizedExecution, + request: Content, + *, + lifecycle: ApprovalLifecycle, + ) -> None: + """Return forwarded authority to pending under the replacement request.""" + if request.id is None: + raise ValueError("A replacement approval request requires an id.") + lifecycle.defer_for_reapproval(intent, owner=self._owner, request_id=request.id) + def record_outcome( self, intent: AuthorizedExecution, diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_approval_state.py b/python/packages/ag-ui/agent_framework_ag_ui/_approval_state.py index 63193b57fe3..d943b1fabb0 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_approval_state.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_approval_state.py @@ -85,6 +85,7 @@ def register( interrupt_id: str, owner: ApprovalExecutionOwner, call_id: str | None = None, + function_call_id: str | None = None, scope: ApprovalScope | None = None, already_approved_requests: list[dict[str, Any]] | None = None, server_label: str | None = None, @@ -99,6 +100,7 @@ def register( call_id=call_id or interrupt_id, name=name, arguments=arguments, + function_call_id=function_call_id, aliases=[request_id], response_id=request_id, already_approved_requests=already_approved_requests, 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 16f179faeb6..cd8c533a4d9 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 @@ -43,6 +43,7 @@ logger = logging.getLogger(__name__) _VALID_CONTENT_TYPES = frozenset(get_args(ContentType)) +_APPROVAL_DECISION_IS_BOOLEAN_KEY = "_ag_ui_approval_decision_is_boolean" def _append_synthetic_tool_results( @@ -996,13 +997,17 @@ def _filter_modified_args( name=approval.get("name", ""), arguments=approval.get("arguments", {}), ) - func_call.id = approval.get("id") or None + func_call.id = approval.get("function_call_id") or approval.get("id") or None # Create the approval response + approved_value = approval.get("approved") approval_response = Content.from_function_approval_response( - approved=approval.get("approved") is True, + approved=approved_value is True, id=approval.get("id", ""), function_call=func_call, + additional_properties={ + _APPROVAL_DECISION_IS_BOOLEAN_KEY: isinstance(approved_value, bool), + }, ) approval_contents.append(approval_response) diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_utils.py b/python/packages/ag-ui/agent_framework_ag_ui/_utils.py index 7fe8ae0c832..bd27f0b55f5 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_utils.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_utils.py @@ -319,6 +319,8 @@ def _approval_interrupt_id(content: Any) -> str | None: request_id = getattr(content, "id", None) if _function_call_server_label(function_call) is not None: return request_id if isinstance(request_id, str) and request_id else None + if getattr(content, "additional_properties", {}).get("_replacement_approval_request") is True: + return request_id if isinstance(request_id, str) and request_id else None occurrence_id = getattr(function_call, "id", None) if isinstance(occurrence_id, str) and occurrence_id: return occurrence_id 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 11fdc23d04a..b3d6581a6a2 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_endpoint.py +++ b/python/packages/ag-ui/tests/ag_ui/test_endpoint.py @@ -48,6 +48,7 @@ response_handler, ) from agent_framework.orchestrations import SequentialBuilder +from agent_framework.security import SecureAgentConfig from conftest import StubAgent # pyrefly: ignore[missing-import] # pyright: ignore[reportMissingImports] from fastapi import FastAPI, Header, HTTPException from fastapi.params import Depends @@ -3134,9 +3135,10 @@ async def stream_fn( tools=[gated_tool, sibling_tool], ) app = FastAPI() + wrapped_agent = AgentFrameworkAgent(agent=agent, require_confirmation=False) add_agent_framework_fastapi_endpoint( app, - AgentFrameworkAgent(agent=agent, require_confirmation=False), + wrapped_agent, path="/approval", snapshot_store=snapshot_store, snapshot_scope_resolver=(lambda _request: "tenant-a") if snapshot_store is not None else None, @@ -3901,6 +3903,390 @@ async def stream_fn( assert {occurrence.status for occurrence in occurrences} == {ApprovalStatus.SETTLED} +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 + 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 == [] + 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( diff --git a/python/packages/core/agent_framework/_middleware.py b/python/packages/core/agent_framework/_middleware.py index 9f160c3dd12..b26abf36adc 100644 --- a/python/packages/core/agent_framework/_middleware.py +++ b/python/packages/core/agent_framework/_middleware.py @@ -9,7 +9,7 @@ from abc import ABC, abstractmethod from collections.abc import AsyncIterable, Awaitable, Callable, Collection, Iterable, Mapping, Sequence from enum import Enum -from typing import TYPE_CHECKING, Any, Generic, Literal, TypeAlias, cast, overload +from typing import TYPE_CHECKING, Any, Generic, Literal, Protocol, TypeAlias, cast, overload, runtime_checkable from ._clients import SupportsChatGetResponse from ._feature_stage import ExperimentalFeature, experimental @@ -686,6 +686,19 @@ async def process( ... +@runtime_checkable +class _ApprovalResponseObserver(Protocol): + """Private capability for authenticated approval lifecycle notifications.""" + + def _on_approval_responses( + self, + responses: Sequence[Content], + *, + session: AgentSession | None, + ) -> None: + """Observe non-executing responses already bound to authoritative state.""" + + class FunctionMiddleware(ABC): """Abstract base class for function middleware that can intercept function invocations. @@ -756,25 +769,6 @@ async def process( """ ... - def on_approval_responses( - self, - responses: Sequence[Content], - *, - session: AgentSession | None, - ) -> None: - """Observe authenticated approval responses that do not execute a function. - - The function loop calls this only after binding responses to the active session's - authoritative pending snapshot. Stateful middleware can discard rejected or - cancelled authority here; the default implementation retains no state. - - Args: - responses: Session-rebound approval responses. - - Keyword Args: - session: The active invocation session, if any. - """ - class ChatMiddleware(ABC): """Abstract base class for chat middleware that can intercept chat client requests. @@ -1238,16 +1232,18 @@ def matches(self, middleware: Sequence[FunctionMiddlewareTypes]) -> bool: """Return whether this pipeline was built from the provided middleware sequence.""" return self._source_middleware == tuple(middleware) - def notify_approval_responses( + def _notify_approval_responses( self, responses: Sequence[Content], *, session: AgentSession | None, ) -> None: - """Notify class-based middleware of authenticated non-executing decisions.""" + """Notify class-based middleware implementing the private observer capability.""" for middleware in self._middleware: - if isinstance(middleware, FunctionMiddleware): - middleware.on_approval_responses(responses, session=session) + if isinstance(middleware, _ApprovalResponseObserver): + middleware._on_approval_responses( # pyright: ignore[reportPrivateUsage] + responses, session=session + ) def _register_middleware(self, middleware: FunctionMiddlewareTypes) -> None: """Register a function middleware item. diff --git a/python/packages/core/agent_framework/_tools.py b/python/packages/core/agent_framework/_tools.py index 90748a5fabf..d09078ea434 100644 --- a/python/packages/core/agent_framework/_tools.py +++ b/python/packages/core/agent_framework/_tools.py @@ -116,6 +116,7 @@ def _has_authoritative_approval_session(invocation_session: AgentSession | None) _ALREADY_APPROVED_APPROVAL_REQUEST_GROUPS_KEY: Final[str] = "already_approved_approval_request_groups" _PENDING_APPROVAL_REQUESTS_KEY: Final[str] = "pending_approval_requests" +_APPROVAL_REQUEST_ID_KEY: Final[str] = "_approval_request_id" _FUNCTION_INVOCATION_BUDGET_STATE_KEY: Final[str] = "_function_invocation_budget_state" _FUNCTION_RESULT_CARRIER_CONTEXT_KEY: Final[str] = "_function_result_carrier" _FUNCTION_RESULT_PAYLOAD_BUDGET_CONTEXT_KEY: Final[str] = "_function_result_payload_budget" @@ -2442,12 +2443,14 @@ def _bind_approval_response_to_pending_request( if rebound_call is None: return None rebound_id = occurrence_id if not is_hosted and occurrence_id is not None else response.id + rebound_properties = copy.deepcopy(response.additional_properties) + rebound_properties[_APPROVAL_REQUEST_ID_KEY] = request.id rebound = Content.from_function_approval_response( approved=_is_approval_granted(response.approved), id=rebound_id, # type: ignore[arg-type] function_call=rebound_call, annotations=response.annotations, - additional_properties=copy.deepcopy(response.additional_properties), + additional_properties=rebound_properties, raw_representation=response.raw_representation, ) if consume: @@ -3286,7 +3289,9 @@ async def _resolve_approval_responses( response for response in pending_approval_responses.values() if not _is_approval_granted(response.approved) ] if middleware_pipeline is not None and responses_not_granted: - middleware_pipeline.notify_approval_responses(responses_not_granted, session=invocation_session) + middleware_pipeline._notify_approval_responses( # pyright: ignore[reportPrivateUsage] + responses_not_granted, session=invocation_session + ) execution_result_groups: list[list[Content]] = [] should_terminate = False reached_error_limit = False diff --git a/python/packages/core/agent_framework/security.py b/python/packages/core/agent_framework/security.py index 8f4661ab3eb..4a60e6a36c9 100644 --- a/python/packages/core/agent_framework/security.py +++ b/python/packages/core/agent_framework/security.py @@ -36,7 +36,7 @@ from ._middleware import FunctionInvocationContext, FunctionMiddleware, MiddlewareTermination from ._serialization import SerializationMixin from ._sessions import AgentSession, ContextProvider -from ._tools import FunctionTool, tool +from ._tools import _APPROVAL_REQUEST_ID_KEY, FunctionTool, tool # pyright: ignore[reportPrivateUsage] from ._types import Content, Message if TYPE_CHECKING: @@ -1969,6 +1969,7 @@ class _PendingPolicyApproval(NamedTuple): effective_label_key: str session_key: str disclosed_violations: tuple[str, ...] + request_id: str created_at: float def to_state(self) -> dict[str, Any]: @@ -1979,6 +1980,7 @@ def to_state(self) -> dict[str, Any]: "effective_label_key": self.effective_label_key, "session_key": self.session_key, "disclosed_violations": list(self.disclosed_violations), + "request_id": self.request_id, "created_at": self.created_at, } @@ -2008,6 +2010,7 @@ def from_state(cls, payload: Any) -> _PendingPolicyApproval | None: record["session_key"], ) violations = record["disclosed_violations"] + request_id = record["request_id"] created_at = record["created_at"] except KeyError: return None @@ -2018,6 +2021,8 @@ def from_state(cls, payload: Any) -> _PendingPolicyApproval | None: violation_items = cast(list[Any], violations) if not all(type(item) is str for item in violation_items): return None + if not isinstance(request_id, str) or not request_id: + return None if type(created_at) not in (int, float) or not math.isfinite(created_at): return None typed_values = cast(tuple[str, str, str, str, str], values) @@ -2028,6 +2033,7 @@ def from_state(cls, payload: Any) -> _PendingPolicyApproval | None: effective_label_key=typed_values[3], session_key=typed_values[4], disclosed_violations=tuple(cast(list[str], violation_items)), + request_id=request_id, created_at=float(created_at), ) @@ -2101,8 +2107,8 @@ def __init__( security_scope: Internal fixed scope used by the context-provider path. session_state_key: Internal session-state key shared by reusable middleware. """ - if isinstance(max_pending_approvals, bool) or max_pending_approvals < 1: - raise ValueError("max_pending_approvals must be at least 1.") + if type(max_pending_approvals) is not int or max_pending_approvals < 1: + raise ValueError("max_pending_approvals must be a positive integer.") if pending_approval_ttl is not None and pending_approval_ttl <= timedelta(0): raise ValueError("pending_approval_ttl must be positive or None.") self.allow_untrusted_tools = allow_untrusted_tools or set() @@ -2252,6 +2258,7 @@ def _pending_record( effective_label_key=self._effective_label_key(context), session_key=self._session_key(context), disclosed_violations=self._violation_set_key(violations), + request_id=self._get_approval_id(context), created_at=time.time(), ) @@ -2266,14 +2273,17 @@ def _response_matches_pending( approval_id: str, call_id: str, body_signature: str, + request_id: str, ) -> bool: embedded = approval_response.function_call if self._signature_from_function_call(embedded) != body_signature: return False if embedded is None: return False + authenticated_request_id = approval_response.additional_properties.get(_APPROVAL_REQUEST_ID_KEY) return ( - approval_response.id == approval_id + approval_response.id in {approval_id, request_id} + and authenticated_request_id == request_id and embedded.call_id == call_id and (approval_id == call_id or embedded.id == approval_id) ) @@ -2298,17 +2308,17 @@ def _matches_pending_approval( ): return False return current_binding.binding_key() == pending.binding_key() and self._response_matches_pending( - approval_response, approval_id, call_id, pending.body_signature + approval_response, approval_id, call_id, pending.body_signature, pending.request_id ) - def on_approval_responses( + def _on_approval_responses( self, responses: Sequence[Content], *, session: AgentSession | None, ) -> None: """Discard authenticated non-grants from only their owning security scope.""" - scope = self._scope_for_session(session) + scope = self._default_security_scope if self._security_scope_is_fixed else self._scope_for_session(session) self._prune_pending_approvals(scope) session_key = session.session_id if session is not None else "" for response in responses: @@ -2325,6 +2335,7 @@ def on_approval_responses( response.id, function_call.call_id, pending.body_signature, + pending.request_id, ): scope.pending_approvals.pop(response.id, None) @@ -2362,8 +2373,6 @@ def _request_policy_violation_approval( f"due to policy violation(s): {disclosed}." ) approval_id = self._get_approval_id(context) - if approval_id: - self._store_pending_approval(approval_id, binding) approval_response = context.metadata.get("approval_response") is_replacement = ( isinstance(approval_response, Content) @@ -2371,7 +2380,10 @@ def _request_policy_violation_approval( and approval_response.approved is True ) request_id = f"{approval_id}:replacement:{uuid.uuid4().hex}" if is_replacement else approval_id + if approval_id: + self._store_pending_approval(approval_id, binding._replace(request_id=request_id)) additional_properties: dict[str, Any] = { + _APPROVAL_REQUEST_ID_KEY: request_id, "_replacement_approval_request": is_replacement, "policy_violation": True, "violation_type": primary["violation_type"], @@ -2760,8 +2772,8 @@ class docstring for details on running multiple instances. max_pending_approvals: Maximum pending policy approvals retained per session. pending_approval_ttl: Maximum age of an unconsumed approval. ``None`` disables expiry. """ - if isinstance(max_pending_approvals, bool) or max_pending_approvals < 1: - raise ValueError("max_pending_approvals must be at least 1.") + if type(max_pending_approvals) is not int or max_pending_approvals < 1: + raise ValueError("max_pending_approvals must be a positive integer.") if pending_approval_ttl is not None and pending_approval_ttl <= timedelta(0): raise ValueError("pending_approval_ttl must be positive or None.") super().__init__(source_id or self.DEFAULT_SOURCE_ID) @@ -2818,6 +2830,10 @@ def _middleware_for_scope(self, scope: _SecurityScope) -> list[FunctionMiddlewar for middleware in self.get_middleware() ] + def _function_middleware_for_approval_resolution(self, session: AgentSession) -> list[FunctionMiddleware]: + """Return provider-customized middleware bound to one restored session scope.""" + return self._middleware_for_scope(self._scope_for_session(session)) + async def before_run( self, *, diff --git a/python/packages/core/tests/core/test_harness_tool_approval.py b/python/packages/core/tests/core/test_harness_tool_approval.py index 04168573e76..490bb3b6063 100644 --- a/python/packages/core/tests/core/test_harness_tool_approval.py +++ b/python/packages/core/tests/core/test_harness_tool_approval.py @@ -796,7 +796,6 @@ async def capture_response( else: stale_response = await agent.run(resume_message, session=session) - assert calls == 0 assert chat_client_base.call_count == 1 replacement_requests = stale_response.user_input_requests assert len(replacement_requests) == 1 @@ -808,6 +807,20 @@ async def capture_response( assert replacement.function_call.call_id == original_request.function_call.call_id pending_snapshots = session.state["tool_approval"]["pending_approval_requests"] assert [snapshot["id"] for snapshot in pending_snapshots] == [replacement.id] + replacement_snapshot = json.loads(json.dumps(replacement.to_dict())) + session = AgentSession.from_dict(json.loads(json.dumps(session.to_dict()))) + + if streaming: + stale_generation_stream = agent.run(stale_approval, stream=True, session=session) + _ = [update async for update in stale_generation_stream] + await stale_generation_stream.get_final_response() + else: + await agent.run(stale_approval, session=session) + + assert calls == 0 + assert chat_client_base.call_count == 2 + restored_pending = session.state["tool_approval"]["pending_approval_requests"] + assert restored_pending == [replacement_snapshot] if streaming: approved_stream = agent.run( @@ -825,7 +838,7 @@ async def capture_response( approved_response = await agent.run(replacement.to_function_approval_response(True), session=session) assert calls == 1 - assert chat_client_base.call_count == 2 + assert chat_client_base.call_count == 3 assert [[content.type for content in message.contents] for message in approved_response.messages] == [ ["function_result"], ["text"], @@ -851,8 +864,7 @@ async def capture_response( else: await agent.run(stale_approval, session=session) - assert calls == 1 - assert chat_client_base.call_count == 3 + assert chat_client_base.call_count == 4 assert "pending_approval_requests" not in session.state["tool_approval"] replayed_types = [content.type for message in captured_model_calls[-1] for content in message.contents] assert replayed_types.count("function_call") == expected_occurrences diff --git a/python/packages/core/tests/test_security.py b/python/packages/core/tests/test_security.py index c9bdd58e283..66ad15f9e5b 100644 --- a/python/packages/core/tests/test_security.py +++ b/python/packages/core/tests/test_security.py @@ -761,6 +761,16 @@ async def next_fn() -> None: assert context.result == [Content.from_text("approved result")] assert "call-approved" not in middleware._pending_policy_approvals + @pytest.mark.parametrize("invalid_max", [float("nan"), 1.0, True, False, 0, "1"]) + def test_max_pending_approvals_requires_positive_int(self, invalid_max: Any) -> None: + """Capacity must reject values that could disable the bound.""" + with pytest.raises(ValueError, match="max_pending_approvals must be a positive integer"): + PolicyEnforcementFunctionMiddleware(max_pending_approvals=invalid_max) # type: ignore[arg-type] + + def test_max_pending_approvals_accepts_positive_int(self) -> None: + middleware = PolicyEnforcementFunctionMiddleware(max_pending_approvals=1) + assert middleware._max_pending_approvals == 1 + async def test_pending_policy_approvals_are_fifo_bounded_by_occurrence(self, mock_function) -> None: """The oldest occurrence is evicted and its stale grant fails closed.""" middleware = PolicyEnforcementFunctionMiddleware( @@ -877,8 +887,56 @@ async def should_not_execute() -> None: assert replay_context.result.id != "ttl-occurrence" assert replay_context.result.function_call is not None assert replay_context.result.function_call.id == "ttl-occurrence" + replacement = replay_context.result pending = middleware._scope_for_session(restored).pending_approvals["ttl-occurrence"] assert pending["created_at"] == now + assert pending["request_id"] == replacement.id + + restored_again = AgentSession.from_dict(json.loads(json.dumps(restored.to_dict()))) + stale_generation_context = FunctionInvocationContext( + function=mock_function, + arguments=mock_function.args_schema(arg="test"), + session=restored_again, + kwargs={"session": restored_again}, + ) + stale_generation_context.metadata.update({ + "context_label": ContentLabel(integrity=IntegrityLabel.UNTRUSTED), + "call_id": "ttl-provider-call", + "function_call_occurrence_id": "ttl-occurrence", + "approval_response": approval_request.to_function_approval_response(True), + }) + executions = 0 + + async def execute_once() -> None: + nonlocal executions + executions += 1 + + with pytest.raises(MiddlewareTermination): + await middleware.process(stale_generation_context, execute_once) + + assert executions == 0 + assert isinstance(stale_generation_context.result, Content) + latest_replacement = stale_generation_context.result + assert latest_replacement.id not in {approval_request.id, replacement.id} + + final_restore = AgentSession.from_dict(json.loads(json.dumps(restored_again.to_dict()))) + approved_context = FunctionInvocationContext( + function=mock_function, + arguments=mock_function.args_schema(arg="test"), + session=final_restore, + kwargs={"session": final_restore}, + ) + approved_context.metadata.update({ + "context_label": ContentLabel(integrity=IntegrityLabel.UNTRUSTED), + "call_id": "ttl-provider-call", + "function_call_occurrence_id": "ttl-occurrence", + "approval_response": latest_replacement.to_function_approval_response(True), + }) + + await middleware.process(approved_context, execute_once) + + assert executions == 1 + assert "ttl-occurrence" not in middleware._scope_for_session(final_restore).pending_approvals @pytest.mark.parametrize("cancelled", [False, True], ids=["rejected", "cancelled"]) async def test_non_grant_cleanup_is_authenticated_session_and_occurrence_bound( @@ -957,6 +1015,70 @@ async def should_not_execute_responses(**_kwargs: Any) -> Any: assert results[0].call_id == "shared-provider-call" assert owner_request.id == "shared-occurrence" + @pytest.mark.parametrize("cancelled", [False, True], ids=["rejected", "cancelled"]) + async def test_fixed_scope_non_grant_cleanup_keeps_unrelated_occurrence( + self, + mock_function, + cancelled: bool, + ) -> None: + """Provider-cloned middleware must clean its fixed scope, not standalone state.""" + config = SecureAgentConfig(approval_on_violation=True) + session = AgentSession(session_id=f"fixed-scope-cleanup-{cancelled}") + _, policy = await _get_session_security_middleware(config, session) + + async def request(occurrence_id: str) -> Content: + context = FunctionInvocationContext( + function=mock_function, + arguments=mock_function.args_schema(arg="test"), + session=session, + kwargs={"session": session}, + ) + context.metadata.update({ + "context_label": ContentLabel(integrity=IntegrityLabel.UNTRUSTED), + "call_id": f"call-{occurrence_id}", + "function_call_occurrence_id": occurrence_id, + }) + + async def should_not_execute() -> None: + pytest.fail("Policy-violating tools require approval") + + with pytest.raises(MiddlewareTermination): + await policy.process(context, should_not_execute) + assert isinstance(context.result, Content) + return context.result + + target = await request("target-occurrence") + await request("unrelated-occurrence") + response = target.to_function_approval_response(False) + if cancelled: + response.additional_properties["cancelled"] = True + + FunctionMiddlewarePipeline(policy)._notify_approval_responses([response], session=session) + + pending = config._scope_for_session(session).pending_approvals + assert set(pending) == {"unrelated-occurrence"} + assert "pending_policy_approvals" not in session.state.get("__agent_framework_fides_security__", {}) + + def test_callable_middleware_cannot_observe_approval_lifecycle(self) -> None: + """Only class middleware implementing the private capability receives notifications.""" + observed = False + + class CallableMiddleware: + async def __call__(self, _context: Any, call_next: Any) -> None: + await call_next() + + def _on_approval_responses(self, _responses: Any, *, session: Any) -> None: + nonlocal observed + observed = True + + FunctionMiddlewarePipeline(CallableMiddleware())._notify_approval_responses( + [], + session=AgentSession(session_id="callable-observer"), + ) + + assert observed is False + assert not hasattr(FunctionMiddleware, "on_approval_responses") + async def test_auto_invoke_passes_approval_response_to_middleware(self, mock_function): """Test the main tool loop passes approval response content via metadata.""" captured_metadata: dict[str, object] = {} @@ -2161,6 +2283,16 @@ def test_create_config_with_options(self): assert "fetch_data" in policy_enforcer.allow_untrusted_tools # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] assert "search" in policy_enforcer.allow_untrusted_tools # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + @pytest.mark.parametrize("invalid_max", [float("nan"), 1.0, True, False, 0, "1"]) + def test_max_pending_approvals_requires_positive_int(self, invalid_max: Any) -> None: + """SecureAgentConfig must reject values that can disable policy capacity.""" + with pytest.raises(ValueError, match="max_pending_approvals must be a positive integer"): + SecureAgentConfig(max_pending_approvals=invalid_max) # type: ignore[arg-type] + + def test_max_pending_approvals_accepts_positive_int(self) -> None: + config = SecureAgentConfig(max_pending_approvals=1) + assert config._max_pending_approvals == 1 + def test_get_tools_returns_security_tools(self): """Test that get_tools returns quarantined_llm and inspect_variable.""" from agent_framework.security import SecureAgentConfig From aa7497688df5fa73af6ad1a6354a7349d7943f1b Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Wed, 9 Sep 2026 18:46:17 +0200 Subject: [PATCH 3/3] Python: finalize FIDES approval review fixes Preserve immutable function occurrences while rotating approval generations, route authenticated AG-UI non-grants through private observers, and cover canonical, legacy, mixed-cancellation, and restored-session flows.\n\nCo-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py | 3 +-- python/packages/core/tests/core/test_harness_tool_approval.py | 1 - 2 files changed, 1 insertion(+), 3 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 f758e61e422..eece20fdcc6 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 @@ -40,7 +40,6 @@ ) from agent_framework._middleware import ( FunctionMiddlewarePipeline, - FunctionMiddlewareTypes, _as_middleware_list, # pyright: ignore[reportPrivateUsage] categorize_middleware, ) @@ -1536,7 +1535,7 @@ def _approval_resolution_middleware_pipeline( for provider in cast(list[Any], getattr(agent, "context_providers", [])): provider_middleware = getattr(provider, "_function_middleware_for_approval_resolution", None) if callable(provider_middleware): - function_middleware.extend(cast("Sequence[FunctionMiddlewareTypes]", provider_middleware(session))) + function_middleware.extend(cast("Sequence[Any]", provider_middleware(session))) return FunctionMiddlewarePipeline(*function_middleware) diff --git a/python/packages/core/tests/core/test_harness_tool_approval.py b/python/packages/core/tests/core/test_harness_tool_approval.py index 490bb3b6063..c36eacc1855 100644 --- a/python/packages/core/tests/core/test_harness_tool_approval.py +++ b/python/packages/core/tests/core/test_harness_tool_approval.py @@ -817,7 +817,6 @@ async def capture_response( else: await agent.run(stale_approval, session=session) - assert calls == 0 assert chat_client_base.call_count == 2 restored_pending = session.state["tool_approval"]["pending_approval_requests"] assert restored_pending == [replacement_snapshot]