diff --git a/python/packages/core/agent_framework/_middleware.py b/python/packages/core/agent_framework/_middleware.py index 77b452214e..4c32608127 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, @@ -1050,6 +1051,21 @@ 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_rejected_approvals(self, responses: Sequence[Content]) -> None: + """Let middleware observe rejected approval decisions without executing tools. + + The approval resolver converts rejected decisions into synthetic function results and + does not re-enter :meth:`execute`. Middleware that retains pending approval state + (for example policy-enforcement bindings) can implement + ``discard_rejected_policy_approvals`` to clear that state here. + """ + if not responses: + return + for middleware in self._middleware: + discard = getattr(middleware, "discard_rejected_policy_approvals", None) + if callable(discard): + discard(responses) + 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 f2f04cf8c2..1d29029070 100644 --- a/python/packages/core/agent_framework/_tools.py +++ b/python/packages/core/agent_framework/_tools.py @@ -2851,6 +2851,7 @@ async def _resolve_approval_responses( max_errors: int, execute_function_calls: _FunctionCallExecutor, invocation_session: AgentSession | None = None, + middleware_pipeline: FunctionMiddlewarePipeline | None = None, ) -> _FunctionProcessingResult: """Resolve inbound approval responses before the next model call.""" from ._types import Message @@ -2877,9 +2878,16 @@ async def _resolve_approval_responses( return _FunctionProcessingResult(errors_in_a_row=errors_in_a_row) # 3. Execute approved decisions once. Rejected decisions are converted to results during normalization below. + # Notify middleware of rejections separately so pending approval state (e.g. policy bindings) + # can be cleared without re-entering tool execution. responses_to_execute = [ response for response in pending_approval_responses.values() if _is_approval_granted(response.approved) ] + rejected_responses = [ + response for response in pending_approval_responses.values() if not _is_approval_granted(response.approved) + ] + if middleware_pipeline is not None and rejected_responses: + middleware_pipeline.notify_rejected_approvals(rejected_responses) execution_result_groups: list[list[Content]] = [] should_terminate = False reached_error_limit = False @@ -3064,6 +3072,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 ._types import ChatResponse, add_usage_details @@ -3086,6 +3095,7 @@ async def _get_response_with_function_invocation( max_errors=max_errors, execute_function_calls=execute_function_calls, invocation_session=invocation_session, + middleware_pipeline=middleware_pipeline, ) function_call_messages.extend(approval_processing.response_messages) errors_in_a_row = approval_processing.errors_in_a_row @@ -3197,6 +3207,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.""" errors_in_a_row = 0 @@ -3215,6 +3226,7 @@ async def _stream_response_with_function_invocation( max_errors=max_errors, execute_function_calls=execute_function_calls, invocation_session=invocation_session, + middleware_pipeline=middleware_pipeline, ) errors_in_a_row = approval_processing.errors_in_a_row total_function_calls = _record_function_calls( @@ -3490,6 +3502,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") @@ -3505,6 +3518,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 397f9af62f..61a6e29073 100644 --- a/python/packages/core/agent_framework/security.py +++ b/python/packages/core/agent_framework/security.py @@ -20,10 +20,12 @@ import logging import re import threading +import time import uuid -from collections.abc import Awaitable, Callable, MutableMapping +from collections import OrderedDict +from collections.abc import Awaitable, Callable, MutableMapping, Sequence from copy import deepcopy -from datetime import datetime +from datetime import datetime, timedelta from enum import Enum from typing import TYPE_CHECKING, Annotated, Any, NamedTuple, cast @@ -1631,12 +1633,19 @@ class _PendingPolicyApproval(NamedTuple): there is no separate user identity here); ``disclosed_violations`` the canonical set of violation types disclosed in the approval request, so an approval granted for one set of risks cannot wave a different (e.g. larger) set that a replay computes after the tool's policy metadata changes. + ``created_at`` is a ``time.monotonic()`` timestamp used for TTL eviction of abandoned approvals. """ body_signature: str label_key: str session_key: str disclosed_violations: tuple[str, ...] + created_at: float + + +# Bounds for abandoned policy-approval state on long-lived middleware instances (#7890). +_DEFAULT_MAX_PENDING_POLICY_APPROVALS = 256 +_DEFAULT_PENDING_POLICY_APPROVAL_TTL = timedelta(hours=1) @experimental(feature_id=ExperimentalFeature.FIDES) @@ -1677,6 +1686,9 @@ def __init__( block_on_violation: bool = True, enable_audit_log: bool = True, approval_on_violation: bool = False, + *, + max_pending_policy_approvals: int = _DEFAULT_MAX_PENDING_POLICY_APPROVALS, + pending_policy_approval_ttl: timedelta | None = _DEFAULT_PENDING_POLICY_APPROVAL_TTL, ) -> None: """Initialize PolicyEnforcementFunctionMiddleware. @@ -1689,19 +1701,33 @@ def __init__( 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_policy_approvals: Maximum number of unconsumed policy-approval + bindings retained on this instance. Oldest entries are evicted when the + limit is exceeded so abandoned approvals cannot grow without bound. + pending_policy_approval_ttl: How long an unconsumed pending approval may live + before it is discarded. ``None`` disables time-based expiration. """ + if max_pending_policy_approvals < 1: + raise ValueError("max_pending_policy_approvals must be >= 1.") + if pending_policy_approval_ttl is not None and pending_policy_approval_ttl.total_seconds() <= 0: + raise ValueError("pending_policy_approval_ttl must be positive when set.") self.allow_untrusted_tools = allow_untrusted_tools or set() self.approval_on_violation = approval_on_violation # If approval_on_violation is True, we don't block - we request approval instead self.block_on_violation = block_on_violation if not approval_on_violation else False self.enable_audit_log = enable_audit_log self.audit_log: list[dict[str, Any]] = [] + self._max_pending_policy_approvals = max_pending_policy_approvals + self._pending_policy_approval_ttl = pending_policy_approval_ttl # Track call_ids awaiting approval, each mapped to a binding record capturing the exact # invocation the approval was requested for: the function name + arguments, the security # label (integrity/confidentiality) shown for review, and the session. Combined with the # call_id key and consume-on-use, an approval cannot re-authorize a repeated call, a # different function, changed arguments, a different security label, or a different session. - self._pending_policy_approvals: dict[str, _PendingPolicyApproval] = {} + # OrderedDict preserves insertion order so abandoned entries can be evicted FIFO (#7890). + self._pending_policy_approvals: OrderedDict[str, _PendingPolicyApproval] = OrderedDict() def _get_call_id(self, context: FunctionInvocationContext) -> str: """Get the tool call id for this invocation context.""" @@ -1790,8 +1816,41 @@ def _pending_record( label_key=self._context_label_key(context), session_key=self._session_key(context), disclosed_violations=self._violation_set_key(violations), + created_at=time.monotonic(), ) + def _prune_pending_policy_approvals(self) -> None: + """Drop expired pending approvals and enforce the max-size bound.""" + ttl = self._pending_policy_approval_ttl + if ttl is not None: + ttl_seconds = ttl.total_seconds() + now = time.monotonic() + expired = [ + call_id + for call_id, pending in self._pending_policy_approvals.items() + if now - pending.created_at > ttl_seconds + ] + for call_id in expired: + self._pending_policy_approvals.pop(call_id, None) + while len(self._pending_policy_approvals) > self._max_pending_policy_approvals: + self._pending_policy_approvals.popitem(last=False) + + def _store_pending_policy_approval(self, call_id: str, pending: _PendingPolicyApproval) -> None: + """Record a pending approval, refreshing TTL/size bounds first.""" + self._prune_pending_policy_approvals() + # Re-insert so a re-request for the same call_id moves to the newest end. + self._pending_policy_approvals.pop(call_id, None) + self._pending_policy_approvals[call_id] = pending + while len(self._pending_policy_approvals) > self._max_pending_policy_approvals: + self._pending_policy_approvals.popitem(last=False) + + def _pending_approval_is_alive(self, pending: _PendingPolicyApproval) -> bool: + """Return whether a pending approval is still within its TTL.""" + ttl = self._pending_policy_approval_ttl + if ttl is None: + return True + return time.monotonic() - pending.created_at <= ttl.total_seconds() + def _signature_from_function_call(self, function_call: Any) -> str | None: """Compute the body signature for a ``function_call`` Content, or None if it is not one.""" if not (isinstance(function_call, Content) and function_call.type == "function_call"): @@ -1840,9 +1899,13 @@ def _matches_pending_approval( call_id = self._get_call_id(context) if not call_id: return False + self._prune_pending_policy_approvals() pending = self._pending_policy_approvals.get(call_id) if pending is None: return False + if not self._pending_approval_is_alive(pending): + self._pending_policy_approvals.pop(call_id, None) + return False approval_response = context.metadata.get("approval_response") if not ( isinstance(approval_response, Content) @@ -1861,6 +1924,36 @@ def _matches_pending_approval( and self._violation_set_key(current_violations) == pending.disclosed_violations ) + def discard_rejected_policy_approvals(self, responses: Sequence[Content]) -> None: + """Clear pending approvals for explicitly rejected decisions. + + The normal agent approval path converts rejections into synthetic function results + without re-entering :meth:`process`. The approval resolver notifies this middleware + through the function middleware pipeline so abandoned bindings are released promptly + instead of waiting for TTL or max-size eviction. + """ + for response in responses: + if not ( + isinstance(response, Content) + and response.type == "function_approval_response" + and response.approved is False + ): + continue + function_call = response.function_call + call_id = function_call.call_id if function_call is not None else None + if not call_id: + continue + pending = self._pending_policy_approvals.get(call_id) + if pending is None: + continue + if self._response_matches_pending(response, call_id, pending.body_signature): + self._pending_policy_approvals.pop(call_id, None) + logger.info( + "Cleared pending policy approval for rejected call_id=%s (function=%s).", + call_id, + function_call.name if function_call is not None else None, + ) + def _consume_pending_approval(self, context: FunctionInvocationContext) -> None: """Remove the pending approval for this call so it authorizes exactly one invocation. @@ -1900,7 +1993,7 @@ def _request_policy_violation_approval( ) call_id = self._get_call_id(context) if call_id: - self._pending_policy_approvals[call_id] = self._pending_record(context, violations) + self._store_pending_policy_approval(call_id, self._pending_record(context, violations)) additional_properties: dict[str, Any] = { "policy_violation": True, "violation_type": primary["violation_type"], diff --git a/python/packages/core/tests/test_security.py b/python/packages/core/tests/test_security.py index 3b9932e910..0cb00ffbb9 100644 --- a/python/packages/core/tests/test_security.py +++ b/python/packages/core/tests/test_security.py @@ -2,7 +2,9 @@ """Unit tests for prompt injection defense system.""" +import asyncio import json +from datetime import timedelta from types import SimpleNamespace import pytest @@ -10,8 +12,13 @@ from agent_framework import AgentSession, ExperimentalFeature, FunctionInvocationContext, FunctionMiddleware from agent_framework._middleware import FunctionMiddlewarePipeline, MiddlewareTermination -from agent_framework._tools import FunctionTool, _auto_invoke_function, normalize_function_invocation_configuration -from agent_framework._types import Content +from agent_framework._tools import ( + FunctionTool, + _auto_invoke_function, + _resolve_approval_responses, + normalize_function_invocation_configuration, +) +from agent_framework._types import Content, Message from agent_framework.security import ( ConfidentialityLabel, ContentLabel, @@ -697,6 +704,136 @@ 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_abandoned_policy_approvals_do_not_grow_without_bound(self, mock_function): + """Regression for #7890: unconsumed approvals must not accumulate without bound.""" + max_pending = 8 + middleware = PolicyEnforcementFunctionMiddleware( + approval_on_violation=True, + max_pending_policy_approvals=max_pending, + pending_policy_approval_ttl=None, + ) + + async def stop_before_execute() -> None: + pytest.fail("Tool execution should not continue before approval") + + for i in range(max_pending + 25): + context = FunctionInvocationContext( + function=mock_function, + arguments=mock_function.args_schema(arg="test"), + ) + context.metadata["context_label"] = ContentLabel(integrity=IntegrityLabel.UNTRUSTED) + context.metadata["call_id"] = f"call-abandoned-{i}" + with pytest.raises(MiddlewareTermination): + await middleware.process(context, stop_before_execute) + + assert len(middleware._pending_policy_approvals) == max_pending + # FIFO eviction: the oldest abandoned entries are gone; the newest remain. + assert "call-abandoned-0" not in middleware._pending_policy_approvals + assert f"call-abandoned-{max_pending + 24}" in middleware._pending_policy_approvals + + async def test_expired_policy_approvals_are_discarded(self, mock_function): + """Pending approvals older than the configured TTL must not authorize a replay.""" + middleware = PolicyEnforcementFunctionMiddleware( + approval_on_violation=True, + pending_policy_approval_ttl=timedelta(milliseconds=1), + ) + request_context = FunctionInvocationContext( + function=mock_function, + arguments=mock_function.args_schema(arg="test"), + ) + request_context.metadata["context_label"] = ContentLabel(integrity=IntegrityLabel.UNTRUSTED) + request_context.metadata["call_id"] = "call-expired" + + async def stop_before_execute() -> None: + pytest.fail("Tool execution should not continue before approval") + + with pytest.raises(MiddlewareTermination): + await middleware.process(request_context, stop_before_execute) + + approval_request = request_context.result + assert isinstance(approval_request, Content) + assert "call-expired" in middleware._pending_policy_approvals + + await asyncio.sleep(0.02) + + replay_context = FunctionInvocationContext( + function=mock_function, + arguments=mock_function.args_schema(arg="test"), + ) + replay_context.metadata["context_label"] = ContentLabel(integrity=IntegrityLabel.UNTRUSTED) + replay_context.metadata["call_id"] = "call-expired" + replay_context.metadata["approval_response"] = approval_request.to_function_approval_response(True) + + async def next_fn() -> None: + pytest.fail("Expired approvals must not authorize execution") + + with pytest.raises(MiddlewareTermination): + await middleware.process(replay_context, next_fn) + + # Expired grant must not execute; a fresh approval request may be stored again. + assert isinstance(replay_context.result, Content) + assert replay_context.result.type == "function_approval_request" + assert replay_context.metadata.get("user_approved_violation") is not True + pending = middleware._pending_policy_approvals.get("call-expired") + assert pending is not None + assert middleware._pending_approval_is_alive(pending) + + async def test_rejected_policy_approval_clears_pending_state_via_resolver(self, mock_function): + """Rejection cleanup must run through the approval resolver, not process(). + + ``_resolve_approval_responses`` converts rejections into synthetic results without + re-entering function middleware. The resolver must notify policy middleware so the + pending binding is released immediately (regression for #7890 / Copilot review). + """ + middleware = PolicyEnforcementFunctionMiddleware(approval_on_violation=True) + request_context = FunctionInvocationContext( + function=mock_function, + arguments=mock_function.args_schema(arg="test"), + ) + request_context.metadata["context_label"] = ContentLabel(integrity=IntegrityLabel.UNTRUSTED) + request_context.metadata["call_id"] = "call-rejected" + + async def stop_before_execute() -> None: + pytest.fail("Tool execution should not continue before approval") + + with pytest.raises(MiddlewareTermination): + await middleware.process(request_context, stop_before_execute) + + approval_request = request_context.result + assert isinstance(approval_request, Content) + assert approval_request.type == "function_approval_request" + assert "call-rejected" in middleware._pending_policy_approvals + + rejection = approval_request.to_function_approval_response(False) + function_call = approval_request.function_call + assert function_call is not None + + async def execute_function_calls(**_kwargs: object) -> object: + pytest.fail("Rejected approvals must not execute tools") + + result = await _resolve_approval_responses( + prepared_messages=[ + Message(role="assistant", contents=[function_call, approval_request]), + Message(role="user", contents=[rejection]), + ], + options={"tools": [mock_function]}, + errors_in_a_row=0, + max_errors=3, + execute_function_calls=execute_function_calls, # type: ignore[arg-type] + middleware_pipeline=FunctionMiddlewarePipeline(middleware), + ) + + assert "call-rejected" not in middleware._pending_policy_approvals + rejection_results = [ + content + for message in result.response_messages + for content in message.contents + if content.type == "function_result" + ] + assert len(rejection_results) == 1 + assert rejection_results[0].call_id == "call-rejected" + assert "rejected" in str(rejection_results[0].result).lower() + 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] = {}