-
Notifications
You must be signed in to change notification settings - Fork 2.3k
Python: fix: bound pending policy approvals #7996
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from 5 commits
dc9d297
7045f81
e34a0f9
2842630
c5d577c
a2cc4a1
2842534
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -21,6 +21,7 @@ | |
| import re | ||
| import threading | ||
| import uuid | ||
| from collections import OrderedDict | ||
| from collections.abc import Awaitable, Callable, MutableMapping | ||
| from copy import deepcopy | ||
| from datetime import datetime | ||
|
|
@@ -1677,6 +1678,7 @@ def __init__( | |
| block_on_violation: bool = True, | ||
| enable_audit_log: bool = True, | ||
| approval_on_violation: bool = False, | ||
| max_pending_approvals: int | None = 1000, | ||
| ) -> None: | ||
| """Initialize PolicyEnforcementFunctionMiddleware. | ||
|
|
||
|
|
@@ -1689,19 +1691,27 @@ 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. | ||
| max_pending_approvals: Maximum number of pending approvals to retain. When exceeded, | ||
| the oldest pending approval is evicted (FIFO). Set to None for no limit. | ||
| Defaults to 1000. | ||
| """ | ||
| if max_pending_approvals is not None and max_pending_approvals <= 0: | ||
| raise ValueError("max_pending_approvals must be None or a positive integer") | ||
|
|
||
| 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]] = [] | ||
| # Track occurrence-aware approval ids, each mapped to a binding record capturing the exact | ||
| # invocation the approval was requested for: the provider call id, function name + arguments, | ||
| # security label shown for review, and session. Combined with 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] = {} | ||
| self._max_pending_approvals = max_pending_approvals | ||
| # 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. | ||
| # OrderedDict preserves insertion order for FIFO eviction when bounded. | ||
| 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.""" | ||
|
|
@@ -1858,6 +1868,19 @@ def _matches_pending_approval( | |
| pending = self._pending_policy_approvals.get(approval_id) | ||
| if pending is None: | ||
| return False | ||
|
|
||
| # Session-mismatch cleanup: if the pending entry is from a different session, | ||
| # it can never be consumed (approvals are session-bound). Remove it to prevent | ||
| # unbounded growth when call_ids are reused across sessions. | ||
| current_session_key = self._session_key(context) | ||
| if pending.session_key != current_session_key: | ||
| del self._pending_policy_approvals[call_id] | ||
| logger.debug( | ||
| f"Removed stale pending approval '{call_id}' from session '{pending.session_key}' " | ||
| f"(current session: '{current_session_key}')" | ||
| ) | ||
| return False | ||
|
|
||
| approval_response = context.metadata.get("approval_response") | ||
| if not ( | ||
| isinstance(approval_response, Content) | ||
|
|
@@ -1913,9 +1936,22 @@ def _request_policy_violation_approval( | |
| f"APPROVAL REQUESTED: Tool '{context.function.name}' requires user approval " | ||
| f"due to policy violation(s): {disclosed}." | ||
| ) | ||
| approval_id = self._get_approval_id(context) | ||
| if approval_id: | ||
| self._pending_policy_approvals[approval_id] = self._pending_record(context, violations) | ||
| call_id = self._get_call_id(context) | ||
| if call_id: | ||
| # If bounded, evict oldest entry when adding a new unique call_id would exceed limit. | ||
| # Do not evict when updating an existing call_id (re-request scenario). | ||
| if ( | ||
| self._max_pending_approvals is not None | ||
| and call_id not in self._pending_policy_approvals | ||
| and len(self._pending_policy_approvals) >= self._max_pending_approvals | ||
| ): | ||
| # Evict oldest (first) entry | ||
| oldest_call_id = next(iter(self._pending_policy_approvals)) | ||
| del self._pending_policy_approvals[oldest_call_id] | ||
| logger.debug( | ||
| f"Evicted oldest pending approval '{oldest_call_id}' to maintain limit of {self._max_pending_approvals}" | ||
| ) | ||
| self._pending_policy_approvals[call_id] = self._pending_record(context, violations) | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Should the bounded map stay keyed by
Contributor
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Thanks for catching this. You're right — the pending-approval map needs to use the same canonical identity as the lookup and consume paths. I've updated the insertion path to use I also added a regression test covering the |
||
| additional_properties: dict[str, Any] = { | ||
| "policy_violation": True, | ||
| "violation_type": primary["violation_type"], | ||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Could we restore
approval_id = self._get_approval_id(context)before building the request?_request_policy_violation_approvalnow defines onlycall_id, but passesapproval_idatsecurity.py:1971, so every violating call withapproval_on_violation=TrueraisesNameErrorbeforeMiddlewareTerminationcan return the approval request.There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Thanks, you're right. I had missed restoring the
approval_id = self._get_approval_id(context)assignment in_request_policy_violation_approval().I've restored it and verified that the approval request is now created with the same
approval_idused for storing, matching, and consuming the pending approval.