Skip to content
Open
Show file tree
Hide file tree
Changes from 5 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
54 changes: 45 additions & 9 deletions python/packages/core/agent_framework/security.py

Copy link
Copy Markdown
Contributor

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_approval now defines only call_id, but passes approval_id at security.py:1971, so every violating call with approval_on_violation=True raises NameError before MiddlewareTermination can return the approval request.

Copy link
Copy Markdown
Contributor Author

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_id used for storing, matching, and consuming the pending approval.

Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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.

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

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Should the bounded map stay keyed by approval_id rather than provider call_id? Framework calls can carry a distinct function_call_occurrence_id; _matches_pending_approval looks up that occurrence-aware ID at security.py:1868 and _consume_pending_approval removes it at security.py:1908, but this path inserts by call_id. Once the crash above is fixed, valid occurrence-aware approvals will still miss the stored record and be re-requested instead of executing.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The 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 _get_approval_id(context) instead of the provider call_id, while keeping call_id as part of the existing approval-binding validation.

I also added a regression test covering the function_call_occurrence_id case to verify the full request → store → approve → match → consume lifecycle.

additional_properties: dict[str, Any] = {
"policy_violation": True,
"violation_type": primary["violation_type"],
Expand Down
Loading
Loading