Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 16 additions & 0 deletions python/packages/core/agent_framework/_middleware.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@
AgentRunInputs,
ChatResponse,
ChatResponseUpdate,
Content,
Message,
ResponseStream,
normalize_messages,
Expand Down Expand Up @@ -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.

Expand Down
14 changes: 14 additions & 0 deletions python/packages/core/agent_framework/_tools.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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(
Expand Down Expand Up @@ -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")
Expand All @@ -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),
)
Expand Down
101 changes: 97 additions & 4 deletions python/packages/core/agent_framework/security.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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.

Expand All @@ -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."""
Expand Down Expand Up @@ -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"):
Expand Down Expand Up @@ -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)
Expand All @@ -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.

Expand Down Expand Up @@ -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"],
Expand Down
Loading
Loading