diff --git a/docs/specs/004-python-function-calling-loop.md b/docs/specs/004-python-function-calling-loop.md index d7d46058180..e96dd9d8eee 100644 --- a/docs/specs/004-python-function-calling-loop.md +++ b/docs/specs/004-python-function-calling-loop.md @@ -409,6 +409,9 @@ that manually replay messages own the equivalent rule: do not resend an approval request, never from the response payload. - A matched approval response consumes its pending entry once. Unmatched, duplicate, and replayed responses do not reach local execution. +- 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. - 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 @@ -500,10 +503,11 @@ that manually replay messages own the equivalent rule: do not resend an approval | Persisted approval replay | Resume executes with the prior call available. | `test_persisted_approval_messages_replay_correctly` | | Hosted approval pass-through | Hosted requests/responses are bound to the recorded provider request and are not processed as local calls. | `test_hosted_tool_approval_response`, `test_hosted_mcp_approval_response_passthrough`, `test_session_approval_binding_reconstructs_hosted_response`, `test_mixed_local_and_hosted_approval_flow` | | Approval-time user input | Every user-input request from one approved execution returns in order with assistant role and no extra model call; the execution consumes one call-budget unit. | `packages/core/tests/core/test_harness_tool_approval.py::test_approval_resume_returns_all_user_input_requests_without_another_model_call`, `packages/core/tests/core/test_function_invocation_logic.py::test_approval_resume_user_input_counts_toward_function_call_budget` | -| Mixed terminal result and follow-up input | Completed siblings remain tool-role while only follow-up input requests use assistant-role messages/updates. | `packages/core/tests/core/test_function_invocation_logic.py::test_approval_resume_separates_terminal_results_from_follow_up_requests`, `packages/openai/tests/openai/test_openai_chat_completion_client.py::test_mixed_approval_resume_roles_serialize_function_result_as_tool` | +| Mixed terminal result and follow-up input | Completed siblings remain tool-role while only follow-up input requests use assistant-role messages/updates. | `packages/core/tests/core/test_function_invocation_logic.py::test_approval_resume_separates_terminal_results_from_follow_up_requests`, `packages/openai/tests/openai/test_openai_chat_completion_client.py::test_mixed_approval_resume_roles_serialize_function_result_as_tool`, `packages/core/tests/core/test_harness_tool_approval.py::test_dynamic_policy_approval_partitions_safe_sibling_result_roles` | | Approval-time middleware termination | Terminal result returns with no extra model call in either response mode. | `packages/core/tests/core/test_function_invocation_logic.py::test_approval_resume_honors_middleware_termination` | | 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` | | 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` | diff --git a/python/packages/core/agent_framework/_sessions.py b/python/packages/core/agent_framework/_sessions.py index c11a8ae9c28..da8cfbcb318 100644 --- a/python/packages/core/agent_framework/_sessions.py +++ b/python/packages/core/agent_framework/_sessions.py @@ -836,20 +836,34 @@ def _is_approval_placeholder_result(content: Content) -> bool: def _approval_controls_to_keep(messages: Sequence[Message]) -> set[int]: unresolved_requests_by_id: dict[str, Content] = {} + local_request_ids_by_call_id: dict[str, deque[str]] = {} + local_request_ids_by_occurrence: dict[str, str] = {} unresolved_local_responses_by_id: dict[str, Content] = {} - local_response_ids_by_call_id: dict[str, deque[str]] = {} + local_responses_by_call_id: dict[str, deque[tuple[str, str | None]]] = {} for message in messages: for content in message.contents: if content.type == "function_approval_request": function_call = content.function_call if content.id is not None and function_call is not None and function_call.call_id is not None: - unresolved_requests_by_id.setdefault(content.id, content) + if content.id not in unresolved_requests_by_id: + unresolved_requests_by_id[content.id] = content + local_request_ids_by_call_id.setdefault(function_call.call_id, deque()).append(content.id) + if function_call.id is not None: + local_request_ids_by_occurrence[function_call.id] = content.id + # A replacement request supersedes the decision that triggered + # reapproval; that old decision is no longer executable authority. + unresolved_local_responses_by_id.pop(content.id, None) + if ( + content.additional_properties.get("_replacement_approval_request") is True + and function_call.id is not None + ): + unresolved_local_responses_by_id.pop(function_call.id, None) continue if content.type == "function_approval_response": function_call = content.function_call if content.id is not None: - unresolved_requests_by_id.pop(content.id, None) + unresolved_requests_by_id.pop(local_request_ids_by_occurrence.get(content.id, content.id), None) if ( content.id is not None and function_call is not None @@ -858,7 +872,11 @@ def _approval_controls_to_keep(messages: Sequence[Message]) -> set[int]: and content.id not in unresolved_local_responses_by_id ): unresolved_local_responses_by_id[content.id] = content - local_response_ids_by_call_id.setdefault(function_call.call_id, deque()).append(content.id) + request_id = local_request_ids_by_occurrence.get(content.id) + local_responses_by_call_id.setdefault(function_call.call_id, deque()).append(( + content.id, + request_id, + )) continue if content.call_id is None: continue @@ -869,8 +887,21 @@ def _approval_controls_to_keep(messages: Sequence[Message]) -> set[int]: } if not (is_terminal_result or is_follow_up_request): continue - if response_ids := local_response_ids_by_call_id.get(content.call_id): - unresolved_local_responses_by_id.pop(response_ids.popleft(), None) + resolved_response = False + if responses := local_responses_by_call_id.get(content.call_id): + while responses and responses[0][0] not in unresolved_local_responses_by_id: + responses.popleft() + if responses: + response_id, request_id = responses.popleft() + unresolved_local_responses_by_id.pop(response_id, None) + if request_id is not None: + unresolved_requests_by_id.pop(request_id, None) + resolved_response = True + if not resolved_response and (request_ids := local_request_ids_by_call_id.get(content.call_id)): + while request_ids and request_ids[0] not in unresolved_requests_by_id: + request_ids.popleft() + if request_ids: + unresolved_requests_by_id.pop(request_ids.popleft(), None) return { id(content) for content in (*unresolved_requests_by_id.values(), *unresolved_local_responses_by_id.values()) diff --git a/python/packages/core/agent_framework/_tools.py b/python/packages/core/agent_framework/_tools.py index e76c37d6650..5fa270ab807 100644 --- a/python/packages/core/agent_framework/_tools.py +++ b/python/packages/core/agent_framework/_tools.py @@ -103,6 +103,17 @@ def _generate_function_call_occurrence_id() -> str: DEFAULT_MAX_CONSECUTIVE_ERRORS_PER_REQUEST: Final[int] = 3 SHELL_TOOL_KIND_VALUE: Final[str] = "shell" _TOOL_APPROVAL_STATE_KEY: Final[str] = "tool_approval" +_RUN_LOCAL_MIDDLEWARE_SESSION_ATTR: Final[str] = "_run_local_function_middleware_session" + + +def _has_authoritative_approval_session(invocation_session: AgentSession | None) -> bool: + """Return whether approval state belongs to a caller-owned session.""" + return ( + invocation_session is not None + and getattr(invocation_session, _RUN_LOCAL_MIDDLEWARE_SESSION_ATTR, False) is not True + ) + + _ALREADY_APPROVED_APPROVAL_REQUEST_GROUPS_KEY: Final[str] = "already_approved_approval_request_groups" _PENDING_APPROVAL_REQUESTS_KEY: Final[str] = "pending_approval_requests" _FUNCTION_INVOCATION_BUDGET_STATE_KEY: Final[str] = "_function_invocation_budget_state" @@ -1978,7 +1989,7 @@ async def _try_execute_function_call_groups( ): visible_requests.append(approval_request) continue - if invocation_session is None: + if not _has_authoritative_approval_session(invocation_session): visible_requests.append(approval_request) continue already_approved_requests.append(approval_request) @@ -2251,26 +2262,29 @@ def _extract_tools( return options.get("tools") if options else None -def _get_tool_approval_state(invocation_session: AgentSession | None) -> dict[str, Any] | None: +def _get_tool_approval_state(invocation_session: AgentSession | None, *, create: bool = True) -> dict[str, Any] | None: """Return the shared tool-approval state bag for the invocation session.""" - if invocation_session is None: + if not _has_authoritative_approval_session(invocation_session): return None - raw_state = invocation_session.state.get(_TOOL_APPROVAL_STATE_KEY) + authoritative_session = cast("AgentSession", invocation_session) + raw_state = authoritative_session.state.get(_TOOL_APPROVAL_STATE_KEY) if isinstance(raw_state, dict): return cast(dict[str, Any], raw_state) from ._harness._tool_approval import ToolApprovalState if isinstance(raw_state, ToolApprovalState): serialized_state = raw_state.to_dict(exclude={"type"}) - invocation_session.state[_TOOL_APPROVAL_STATE_KEY] = serialized_state + authoritative_session.state[_TOOL_APPROVAL_STATE_KEY] = serialized_state return serialized_state if raw_state is not None: raise TypeError( f"Session state for {_TOOL_APPROVAL_STATE_KEY!r} must be a dict or ToolApprovalState, " f"got {type(raw_state).__name__}." ) + if not create: + return None new_state: dict[str, Any] = {} - invocation_session.state[_TOOL_APPROVAL_STATE_KEY] = new_state + authoritative_session.state[_TOOL_APPROVAL_STATE_KEY] = new_state return new_state @@ -2287,7 +2301,7 @@ def _content_from_state(value: Any) -> Content | None: def _load_pending_approval_requests(invocation_session: AgentSession | None) -> dict[str, Content]: """Load immutable approval-request snapshots keyed by request ID.""" - state = _get_tool_approval_state(invocation_session) + state = _get_tool_approval_state(invocation_session, create=False) if state is None: return {} raw_requests = state.get(_PENDING_APPROVAL_REQUESTS_KEY, []) @@ -2367,7 +2381,7 @@ def _bind_approval_response_to_pending_request( """Bind one approval response to a session-recorded request.""" from ._types import Content - if invocation_session is None: + if not _has_authoritative_approval_session(invocation_session): return response pending = _load_pending_approval_requests(invocation_session) request_key = response.id @@ -2381,6 +2395,7 @@ def _bind_approval_response_to_pending_request( (pending_id, candidate) for pending_id, candidate in pending.items() if not _is_hosted_tool_approval(candidate) + and candidate.additional_properties.get("_replacement_approval_request") is not True and candidate.function_call is not None and candidate.function_call.id == response.id ] @@ -2396,17 +2411,22 @@ def _bind_approval_response_to_pending_request( if not is_hosted and occurrence_id is not None: embedded_call = response.function_call uses_occurrence_id = response.id == occurrence_id - uses_legacy_request_id = response.id == request.id + uses_request_id = response.id == request.id + is_replacement = request.additional_properties.get("_replacement_approval_request") is True if not uses_occurrence_id: - if not (uses_legacy_request_id and embedded_call is not None and embedded_call.id == occurrence_id): + if is_replacement and uses_request_id: + if embedded_call is not None and embedded_call.id != occurrence_id: + return None + elif not (uses_request_id and embedded_call is not None and embedded_call.id == occurrence_id): return None - warnings.warn( - "An occurrence-aware approval used the legacy provider call_id request binding. " - "Return function_call.id as the approval response id; legacy request-id binding will be removed " - "in a future release.", - FutureWarning, - stacklevel=3, - ) + else: + warnings.warn( + "An occurrence-aware approval used the legacy provider call_id request binding. " + "Return function_call.id as the approval response id; legacy request-id binding will be removed " + "in a future release.", + FutureWarning, + stacklevel=3, + ) elif embedded_call is not None and embedded_call.id != occurrence_id: return None elif not is_hosted: @@ -2545,14 +2565,21 @@ def _collect_approval_responses( """ approval_responses: list[Content] = [] pending_by_call_id: dict[str, deque[Content]] = {} + pending_by_approval_id: dict[str, Content] = {} resolved_response_ids: set[int] = set() for message in messages: for content in message.contents: + if content.type == "function_approval_request" and content.id is not None: + if superseded := pending_by_approval_id.pop(content.id, None): + resolved_response_ids.add(id(superseded)) + continue if content.type == "function_approval_response" and not _is_hosted_tool_approval(content): function_call = content.function_call if function_call is None or function_call.call_id is None: continue approval_responses.append(content) + if content.id is not None: + pending_by_approval_id[content.id] = content pending_by_call_id.setdefault(function_call.call_id, deque()).append(content) continue if content.call_id is None: @@ -2565,8 +2592,13 @@ def _collect_approval_responses( if not (is_terminal_result or is_follow_up_request): continue pending_responses = pending_by_call_id.get(content.call_id) + while pending_responses and id(pending_responses[0]) in resolved_response_ids: + pending_responses.popleft() if pending_responses: - resolved_response_ids.add(id(pending_responses.popleft())) + resolved = pending_responses.popleft() + resolved_response_ids.add(id(resolved)) + if resolved.id is not None and pending_by_approval_id.get(resolved.id) is resolved: + pending_by_approval_id.pop(resolved.id, None) return { content.id: content @@ -2576,9 +2608,10 @@ def _collect_approval_responses( def _collect_unanswered_approval_requests(messages: Sequence[Message]) -> list[Content]: - approval_requests_by_id: dict[str, Content] = {} - pending_request_ids_by_call_id: dict[str, deque[str]] = {} - answered_approval_ids: set[str] = set() + unanswered_by_id: dict[str, Content] = {} + requests_by_call_id: dict[str, deque[Content]] = {} + request_ids_by_occurrence: dict[str, str] = {} + answered_request_ids_by_call_id: dict[str, deque[str]] = {} for message in messages: for content in message.contents: @@ -2586,13 +2619,19 @@ def _collect_unanswered_approval_requests(messages: Sequence[Message]) -> list[C function_call = content.function_call if content.id is None or function_call is None or function_call.call_id is None: continue - if content.id not in approval_requests_by_id: - approval_requests_by_id[content.id] = content - pending_request_ids_by_call_id.setdefault(function_call.call_id, deque()).append(content.id) + if content.id not in unanswered_by_id: + unanswered_by_id[content.id] = content + requests_by_call_id.setdefault(function_call.call_id, deque()).append(content) + if function_call.id is not None: + request_ids_by_occurrence[function_call.id] = content.id continue if content.type == "function_approval_response": + function_call = content.function_call if content.id is not None: - answered_approval_ids.add(content.id) + request_id = request_ids_by_occurrence.get(content.id, content.id) + unanswered_by_id.pop(request_id, None) + if function_call is not None and function_call.call_id is not None: + answered_request_ids_by_call_id.setdefault(function_call.call_id, deque()).append(request_id) continue if content.call_id is None: continue @@ -2603,12 +2642,19 @@ def _collect_unanswered_approval_requests(messages: Sequence[Message]) -> list[C } if not (is_terminal_result or is_follow_up_request): continue - if request_ids := pending_request_ids_by_call_id.get(content.call_id): - answered_approval_ids.add(request_ids.popleft()) + answered_requests = answered_request_ids_by_call_id.get(content.call_id) + if answered_requests: + answered_requests.popleft() + continue + requests = requests_by_call_id.get(content.call_id) + while requests and (requests[0].id is None or unanswered_by_id.get(requests[0].id) is not requests[0]): + requests.popleft() + if requests: + resolved = requests.popleft() + if resolved.id is not None: + unanswered_by_id.pop(resolved.id, None) - return [ - request for approval_id, request in approval_requests_by_id.items() if approval_id not in answered_approval_ids - ] + return list(unanswered_by_id.values()) def _remove_unanswered_approval_batches_from_model_input(messages: list[Message]) -> None: @@ -2742,7 +2788,21 @@ def _replace_approval_contents_with_results( result_groups_by_call_id: dict[str, deque[list[Content]]] = {} for result_group in approved_function_result_groups: - call_id = next((result.call_id for result in result_group if result.call_id is not None), None) + call_id = next( + ( + result.function_call.call_id + if result.type == "function_approval_request" and result.function_call is not None + else result.call_id + for result in result_group + if result.call_id is not None + or ( + result.type == "function_approval_request" + and result.function_call is not None + and result.function_call.call_id is not None + ) + ), + None, + ) if call_id is not None: result_groups_by_call_id.setdefault(call_id, deque()).append(result_group) @@ -2844,7 +2904,18 @@ def find_approval_occurrence(approval_id: str) -> _ApprovalCallOccurrence | None else: replacement_groups_by_index[content_idx] = replacements if occurrence is not None: - occurrence.closed = True + replacement_request = next( + ( + replacement + for replacement in replacements + if replacement.type == "function_approval_request" + ), + None, + ) + if replacement_request is not None: + occurrence.approval_id = replacement_request.id + else: + occurrence.closed = True resolved_contents.extend(replacements) elif content.type == "function_result": if content.call_id is None: @@ -3116,11 +3187,17 @@ def _handle_function_call_results( # Only add items that aren't already in the message (e.g. function_approval_request wrappers). # Declaration-only function_call items are already present from the LLM response. new_items = [result for result in execution_results if result.type != "function_call"] - if new_items: - if response.messages and response.messages[0].role == "assistant": - response.messages[0].contents.extend(new_items) - else: - response.messages.append(Message(role="assistant", contents=new_items)) + response_messages, _ = _messages_and_updates_for_terminal_contents(new_items) + if ( + response_messages + and all(message.role == "assistant" for message in response_messages) + and response.messages + and response.messages[0].role == "assistant" + ): + for message in response_messages: + response.messages[0].contents.extend(message.contents) + else: + response.messages.extend(response_messages) streaming_items: list[Content] = [] for result in execution_results: if result.type == "function_call": @@ -3129,11 +3206,12 @@ def _handle_function_call_results( streaming_items.append(metadata_only_result) else: streaming_items.append(result) + _, streaming_updates = _messages_and_updates_for_terminal_contents(streaming_items) return _FunctionProcessingResult( errors_in_a_row=errors_in_a_row, action="return", function_call_count=function_call_count, - streaming_updates=(ChatResponseUpdate(contents=streaming_items, role="assistant"),), + streaming_updates=streaming_updates, ) errors_in_a_row, reached_error_limit = _update_consecutive_error_count( @@ -3173,6 +3251,11 @@ async def _resolve_approval_responses( from ._middleware import MiddlewareFailure from ._types import Message + active_pending_ids = ( + set(_load_pending_approval_requests(invocation_session)) + if _has_authoritative_approval_session(invocation_session) + else None + ) _bind_approval_responses_to_pending_requests(prepared_messages, invocation_session) # 1. Restore safe siblings hidden with a prior mixed approval batch when its visible decision arrives. @@ -3222,14 +3305,41 @@ async def _resolve_approval_responses( max_errors=max_errors, ) - # 4. Replace approval controls/placeholders with terminal contents, correlated by logical call occurrence. + # 4. Snapshot unanswered siblings before normalization removes their wrappers, then merge them with + # replacement requests produced while resolving this response. + produced_replacement_request = any( + content.type == "function_approval_request" + and content.additional_properties.get("_replacement_approval_request") is True + for result_group in execution_result_groups + for content in result_group + ) + pending_before_normalization = ( + [ + request + for request in _collect_unanswered_approval_requests(prepared_messages) + if active_pending_ids is None or request.id in active_pending_ids + ] + if produced_replacement_request + else [] + ) terminal_contents = _replace_approval_contents_with_results( prepared_messages, pending_approval_responses, execution_result_groups, ) - if pending_requests := _collect_unanswered_approval_requests(prepared_messages): - terminal_contents.extend(pending_requests) + pending_by_id = { + request.id: request + for request in (*pending_before_normalization, *_collect_unanswered_approval_requests(prepared_messages)) + if request.id is not None + } + if pending_by_id: + surfaced_request_ids = { + content.id for content in terminal_contents if content.type == "function_approval_request" + } + terminal_contents.extend( + request for request_id, request in pending_by_id.items() if request_id not in surfaced_request_ids + ) + _store_pending_approval_requests(invocation_session, list(pending_by_id.values())) # 5. Return role-correct output and tell the outer loop whether to return, stop tools, or call the model. executed_function_count = len(execution_result_groups) @@ -3979,8 +4089,11 @@ def get_response( request_kwargs.pop("middleware", []), supported_categories=("chat", "function") ) - function_middleware_pipeline = self._get_function_middleware_pipeline( - categorized_runtime_middleware["function"] + runtime_function_middleware = categorized_runtime_middleware["function"] + function_middleware_pipeline = self._get_function_middleware_pipeline(runtime_function_middleware) + requires_session_state = any( + getattr(item, "_requires_session_state", False) is True + for item in (*self.function_middleware, *runtime_function_middleware) ) if categorized_runtime_middleware["chat"]: request_kwargs["middleware"] = categorized_runtime_middleware["chat"] @@ -4011,6 +4124,9 @@ def get_response( raw_session = request_kwargs.get("session") invocation_session = raw_session if isinstance(raw_session, _AgentSession) else None + if invocation_session is None and requires_session_state: + invocation_session = _AgentSession() + setattr(invocation_session, _RUN_LOCAL_MIDDLEWARE_SESSION_ATTR, True) # Bind one executor with the run's custom arguments, middleware, configuration, and session. execute_function_calls = partial( diff --git a/python/packages/core/agent_framework/_types.py b/python/packages/core/agent_framework/_types.py index 93a5679c88a..7b0efa634e8 100644 --- a/python/packages/core/agent_framework/_types.py +++ b/python/packages/core/agent_framework/_types.py @@ -1317,6 +1317,7 @@ def from_function_approval_request( and function_call.id is not None and id != function_call.id and function_call.additional_properties.get("server_label") is None + and not (additional_properties or {}).get("_replacement_approval_request", False) ): warnings.warn( "Creating a local function_approval_request whose id differs from function_call.id uses the legacy " diff --git a/python/packages/core/agent_framework/security.py b/python/packages/core/agent_framework/security.py index f2ccf1f673e..1a46bd36876 100644 --- a/python/packages/core/agent_framework/security.py +++ b/python/packages/core/agent_framework/security.py @@ -16,17 +16,18 @@ import asyncio import contextlib +import hashlib import json import logging import math import re import uuid from collections.abc import Awaitable, Callable, Mapping, MutableMapping -from contextvars import ContextVar +from contextvars import ContextVar, Token from copy import copy, deepcopy from datetime import datetime from enum import Enum -from typing import TYPE_CHECKING, Annotated, Any, NamedTuple, cast +from typing import TYPE_CHECKING, Annotated, Any, NamedTuple, NoReturn, cast from pydantic import BaseModel, Field @@ -68,7 +69,12 @@ logger = logging.getLogger(__name__) -_BRACKETED_VAR_REF_RE = re.compile(r"^\[\s*(var_[0-9a-fA-F]+)\s*\]$") +_BRACKETED_VAR_ID = r"var_[0-9a-fA-F]+" +_BARE_VAR_ID = r"var_[0-9a-fA-F]{8,}" +_EMBEDDED_VAR_REF_RE = re.compile(rf"\[\s*(?P{_BRACKETED_VAR_ID})\s*\]|\b(?P{_BARE_VAR_ID})\b") +_WHOLE_VAR_REF_RE = re.compile(rf"\s*(?:\[\s*(?P{_BRACKETED_VAR_ID})\s*\]|(?P{_BARE_VAR_ID}))\s*") +_BARE_REFERENCE_WARNING = "Expanded a bare variable reference in a tool argument; models should use [var_] instead." +_UNRESOLVED = object() # Tools that consume variable IDs literally (as opaque references) and therefore # must NOT have ``var_xxx`` arguments expanded to stored content before execution. @@ -688,6 +694,103 @@ def from_message(cls, message: dict[str, Any], index: int | None = None) -> Labe _STATE_VARIABLE_METADATA = "variable_metadata" _STATE_AUDIT_LOG = "audit_log" _STATE_PENDING_APPROVALS = "pending_policy_approvals" +_STANDALONE_SESSION_STATE_KEY = "__agent_framework_fides_security__" + + +def _strict_json_value( + value: Any, + *, + path: str, + canonical: bool, + allow_content: bool = False, + active_container_ids: set[int] | None = None, +) -> Any: + """Return a strict JSON-compatible value without string coercion.""" + if active_container_ids is None: + active_container_ids = set() + if isinstance(value, Enum): + return _strict_json_value( + value.value, + path=path, + canonical=canonical, + allow_content=allow_content, + active_container_ids=active_container_ids, + ) + if type(value) in (str, int, bool, type(None)): + return value + if type(value) is float: + if not math.isfinite(value): + raise ValueError(f"{path} contains a non-finite float") + return value + if canonical and isinstance(value, AgentSession): + return ["agent_session", value.session_id] + if allow_content and isinstance(value, Content): + return _strict_json_value( + value.to_dict(), + path=path, + canonical=canonical, + active_container_ids=active_container_ids, + ) + + is_list = isinstance(value, list) + is_tuple = isinstance(value, tuple) + if is_list or is_tuple: + container = cast(list[Any] | tuple[Any, ...], value) + if canonical and type(container) not in (list, tuple): + raise TypeError(f"{path} contains unsupported type {type(container).__name__}") + container_id = id(container) + if container_id in active_container_ids: + raise ValueError(f"{path} contains a circular reference") + active_container_ids.add(container_id) + try: + items = [ + _strict_json_value( + item, + path=f"{path}[{index}]", + canonical=canonical, + allow_content=allow_content, + active_container_ids=active_container_ids, + ) + for index, item in enumerate(container) + ] + finally: + active_container_ids.remove(container_id) + return ["tuple" if is_tuple else "list", items] if canonical else items + + if isinstance(value, Mapping): + mapping = cast(Mapping[Any, Any], cast(object, value)) + value_type = type(cast(object, value)) + if canonical and value_type is not dict: + raise TypeError(f"{path} contains unsupported type {value_type.__name__}") + container_id = id(cast(object, value)) + if container_id in active_container_ids: + raise ValueError(f"{path} contains a circular reference") + active_container_ids.add(container_id) + try: + raw_keys = list(mapping.keys()) + if any(type(key) is not str for key in raw_keys): + raise TypeError(f"{path} contains a non-string mapping key") + keys = cast(list[str], raw_keys) + if canonical: + keys.sort() + items = [ + ( + key, + _strict_json_value( + mapping[key], + path=f"{path}.{key}", + canonical=canonical, + allow_content=allow_content, + active_container_ids=active_container_ids, + ), + ) + for key in keys + ] + finally: + active_container_ids.remove(container_id) + return ["mapping", items] if canonical else dict(items) + + raise TypeError(f"{path} contains unsupported type {type(value).__name__}") def _durable_state_snapshot( @@ -859,6 +962,10 @@ def pending_approvals(self) -> dict[str, Any]: """Return serialized pending approval records.""" return self._mapping(_STATE_PENDING_APPROVALS) + def variable_store(self) -> ContentVariableStore: + """Return an owner-aware store backed by this scope.""" + return _ScopedVariableStore(self) + def _audit_entries(self) -> list[Any]: """Return raw audit entries from scope state.""" entries = self._state.get(_STATE_AUDIT_LOG) @@ -896,6 +1003,39 @@ def clear_audit_log(self) -> None: self._audit_entries().clear() +class _SecurityScopeBinding: + """Select task-local session state for a reusable middleware instance.""" + + def _initialize_security_scope(self, scope: _SecurityScope | None, *, session_state_key: str) -> None: + self._default_security_scope = scope if scope is not None else _SecurityScope() + self._security_scope_is_fixed = scope is not None + self._security_session_state_key = session_state_key + self._active_security_scope: ContextVar[_SecurityScope | None] = ContextVar( + f"agent_framework_security_scope_{id(self)}", default=None + ) + + @property + def _scope(self) -> _SecurityScope: + return self._active_security_scope.get() or self._default_security_scope + + def _scope_for_session(self, session: AgentSession | None) -> _SecurityScope: + if session is None: + return self._scope + stored = session.state.get(self._security_session_state_key) + if stored is None: + stored = {} + session.state[self._security_session_state_key] = stored + if not isinstance(stored, dict): + raise ValueError("Security session state must be a dictionary.") + return _SecurityScope(cast(dict[str, Any], stored), scope_id=session.session_id) + + def _activate_security_scope(self, context: FunctionInvocationContext) -> Token[_SecurityScope | None]: + scope = self._default_security_scope + if not self._security_scope_is_fixed and context.session is not None: + scope = self._scope_for_session(context.session) + return self._active_security_scope.set(scope) + + class _ScopedVariableStore(ContentVariableStore): """Variable store backed by one security scope's serializable state.""" @@ -983,7 +1123,7 @@ def list_variables(self) -> list[str]: @experimental(feature_id=ExperimentalFeature.FIDES) -class LabelTrackingFunctionMiddleware(FunctionMiddleware): +class LabelTrackingFunctionMiddleware(FunctionMiddleware, _SecurityScopeBinding): """Middleware that tracks and propagates security labels through tool invocations. Tiered Label Propagation: @@ -1042,70 +1182,53 @@ async def get_weather(city: str) -> str: response = await agent.run(messages=[{"role": "user", "content": "What's the weather?"}]) """ + _requires_session_state = True + def __init__( self, default_integrity: IntegrityLabel = IntegrityLabel.UNTRUSTED, default_confidentiality: ConfidentialityLabel = ConfidentialityLabel.PUBLIC, auto_hide_untrusted: bool = True, hide_threshold: IntegrityLabel = IntegrityLabel.UNTRUSTED, + *, + security_scope: _SecurityScope | None = None, + session_state_key: str = _STANDALONE_SESSION_STATE_KEY, ) -> None: - """Initialize LabelTrackingFunctionMiddleware. - - Args: - default_integrity: Default integrity label for tools without source_integrity. - Defaults to UNTRUSTED for safety (tools must opt-in to TRUSTED). - default_confidentiality: Default confidentiality label. Defaults to PUBLIC. - auto_hide_untrusted: Whether to automatically hide untrusted results. Defaults to True. - hide_threshold: The integrity level at which to hide content. Defaults to UNTRUSTED. - """ + """Initialize label tracking and bind its security-state selector.""" self.default_integrity = default_integrity self.default_confidentiality = default_confidentiality self.auto_hide_untrusted = auto_hide_untrusted self.hide_threshold = hide_threshold - - self._security_scope = _SecurityScope() - self._variable_store = _ScopedVariableStore(self._security_scope) + self._initialize_security_scope(security_scope, session_state_key=session_state_key) def _clone_for_scope(self, scope: _SecurityScope) -> LabelTrackingFunctionMiddleware: - """Clone current middleware configuration into a session scope.""" + """Clone customized middleware configuration into a fixed session scope.""" scoped = copy(self) - scoped._security_scope = scope - scoped._variable_store = _ScopedVariableStore(scope) + scoped._initialize_security_scope(scope, session_state_key=self._security_session_state_key) return scoped @property def _context_label(self) -> ContentLabel: - """Return the cumulative label from this middleware's fixed scope.""" - return self._security_scope.context_label + return self._scope.context_label @_context_label.setter def _context_label(self, label: ContentLabel) -> None: - self._security_scope.context_label = label + self._scope.context_label = label @property def _variable_metadata(self) -> dict[str, Any]: - """Return variable metadata from this middleware's fixed scope.""" - return self._security_scope.variable_metadata - - def get_context_label(self) -> ContentLabel: - """Get the current context-level security label. - - The context label represents the cumulative security state of the conversation. - It starts as TRUSTED + PUBLIC and gets "tainted" as untrusted or private - content is added to the context. + return self._scope.variable_metadata - Returns: - The current context security label. - """ - return self._context_label + def get_context_label(self, session: AgentSession | None = None) -> ContentLabel: + """Get the cumulative context label for an optional explicit session.""" + return self._scope_for_session(session).context_label - def reset_context_label(self) -> None: - """Reset the context label to initial state (TRUSTED + PUBLIC). - - Call this when starting a new conversation or session. - """ - self._context_label = ContentLabel( - integrity=IntegrityLabel.TRUSTED, confidentiality=ConfidentialityLabel.PUBLIC, metadata={"reset": True} + def reset_context_label(self, session: AgentSession | None = None) -> None: + """Reset the cumulative context label for an optional explicit session.""" + self._scope_for_session(session).context_label = ContentLabel( + integrity=IntegrityLabel.TRUSTED, + confidentiality=ConfidentialityLabel.PUBLIC, + metadata={"reset": True}, ) logger.info("Context label reset to TRUSTED + PUBLIC") @@ -1162,176 +1285,69 @@ def _extract_primary_tool_content(expanded_content: Any) -> Any: return expanded_content - def _expand_variable_reference(self, value: Any) -> Any: - """Expand variable references (e.g., ``[var_abc123]``) to stored content. - - This enables direct tool chaining where models pass variable placeholders as - arguments to subsequent local/MCP tool calls. - - Accepts two reference forms: - - 1. **Bracketed** (canonical): ``[var_abc123]`` — the documented form models - are instructed to emit. - 2. **Bare** (lenient fallback): ``var_abc123`` — some models drop the - brackets when copying a variable id into a tool argument. To prevent - the literal token from leaking into a destination (e.g. a write to a - public README), we also resolve bare ``var_`` tokens that - correspond to a known stored variable. A warning is logged on every - bare expansion so the failure mode remains observable. - - When a stored variable is a dict with a ``response`` key (e.g., from - ``quarantined_llm``), only the ``response`` value is extracted and - returned, ensuring tools receive main content without metadata fields. - """ - if isinstance(value, str): - # Unanchored bracketed pattern (canonical form). - bracketed_pattern = r"\[\s*(var_[0-9a-fA-F]+)\s*\]" - # Word-boundary bare token. Require >=8 hex chars to limit false - # positives on unrelated strings that happen to contain "var_*". - # Variable ids generated by ContentVariableStore are 16 hex chars. - bare_pattern = r"\bvar_[0-9a-fA-F]{8,}\b" - - bracketed_matches = re.findall(bracketed_pattern, value) - bare_matches = re.findall(bare_pattern, value) + def _resolve_variable_references(self, value: Any) -> tuple[Any, list[ContentLabel]]: + """Recursively resolve owned variable references and return their stored labels.""" + labels: list[ContentLabel] = [] + return self._resolve_value(value, labels), labels - if not bracketed_matches and not bare_matches: + def _lookup_variable(self, variable_id: str, labels: list[ContentLabel]) -> Any: + try: + stored_content, stored_label = self.get_variable_store().retrieve(variable_id) + except KeyError: + return _UNRESOLVED + labels.append(stored_label) + return self._extract_primary_tool_content(stored_content) + + def _resolve_string(self, value: str, labels: list[ContentLabel]) -> Any: + if not _EMBEDDED_VAR_REF_RE.search(value): + return value + + whole = _WHOLE_VAR_REF_RE.fullmatch(value) + if whole is not None: + variable_id = whole.group("bracketed") or whole.group("bare") + resolved = self._lookup_variable(variable_id, labels) + if resolved is _UNRESOLVED: return value - - # Whole-string canonical match: ``[var_xxx]`` - if len(bracketed_matches) == 1: - whole_bracketed = _BRACKETED_VAR_REF_RE.match(value) - if whole_bracketed is not None: - variable_id = whole_bracketed.group(1) - try: - expanded_content, _ = self._variable_store.retrieve(variable_id) - extracted = self._extract_primary_tool_content(expanded_content) - if extracted is not expanded_content: - logger.debug( - f"Expanded variable placeholder '{value}' for tool argument " - f"(extracted primary content from stored payload)" - ) - return extracted - logger.debug(f"Expanded variable placeholder '{value}' for tool argument") - return expanded_content - except KeyError: - logger.debug(f"Variable placeholder '{value}' could not be resolved") - return value - - # Whole-string bare match: ``var_xxx`` (no brackets). Only treat as - # a variable reference if the id actually exists in the store; this - # keeps random strings that happen to look like ``var_xxx`` from - # being silently mangled. - whole_bare = re.fullmatch(r"\s*(var_[0-9a-fA-F]{8,})\s*", value) - if whole_bare is not None and not bracketed_matches: - variable_id = whole_bare.group(1) - try: - expanded_content, _ = self._variable_store.retrieve(variable_id) - extracted = self._extract_primary_tool_content(expanded_content) - logger.warning( - "Expanded a bare variable reference for a tool argument. Models should wrap " - "variable references in '[ ]' brackets; accepting the bare form prevents the " - "literal handle from leaking to a destination." - ) - if extracted is not expanded_content: - return extracted - return expanded_content - except KeyError: - # Not a known variable id; leave string untouched. - return value - - # Embedded substitutions. Apply bracketed pass first, then bare pass - # on the result so that ``[var_xxx]`` is never double-handled. - def replace_bracketed(match_obj: Any) -> str: - variable_id = match_obj.group(1) - try: - expanded_content, _ = self._variable_store.retrieve(variable_id) - extracted = self._extract_primary_tool_content(expanded_content) - if extracted is not expanded_content: - logger.debug( - f"Expanded embedded variable placeholder '[{variable_id}]' in tool argument " - f"(extracted primary content from stored payload)" - ) - return str(extracted) - logger.debug(f"Expanded embedded variable placeholder '[{variable_id}]' in tool argument") - return str(expanded_content) - except KeyError: - logger.debug(f"Variable placeholder '[{variable_id}]' could not be resolved") - return match_obj.group(0) - - result = re.sub(bracketed_pattern, replace_bracketed, value) - - def replace_bare(match_obj: Any) -> str: - variable_id = match_obj.group(0) - try: - expanded_content, _ = self._variable_store.retrieve(variable_id) - extracted = self._extract_primary_tool_content(expanded_content) - logger.warning( - "Expanded an embedded bare variable reference in a tool argument. Models should " - "wrap variable references in '[ ]' brackets." - ) - if extracted is not expanded_content: - return str(extracted) - return str(expanded_content) - except KeyError: - # Not a known variable id; leave the token in place. - return match_obj.group(0) - - return re.sub(bare_pattern, replace_bare, result) - + if whole.group("bare"): + logger.warning(_BARE_REFERENCE_WARNING) + return resolved + + def replace(match: re.Match[str]) -> str: + variable_id = match.group("bracketed") or match.group("bare") + resolved = self._lookup_variable(variable_id, labels) + if resolved is _UNRESOLVED: + return match.group(0) + if match.group("bare"): + logger.warning(_BARE_REFERENCE_WARNING) + return str(resolved) + + return _EMBEDDED_VAR_REF_RE.sub(replace, value) + + def _resolve_value(self, value: Any, labels: list[ContentLabel]) -> Any: + if isinstance(value, str): + return self._resolve_string(value, labels) if isinstance(value, BaseModel): - return self._expand_variable_reference(value.model_dump()) - + return self._resolve_value(value.model_dump(), labels) if isinstance(value, dict): value_dict = cast(dict[str, Any], value) - return {k: self._expand_variable_reference(v) for k, v in value_dict.items()} - + return {key: self._resolve_value(item, labels) for key, item in value_dict.items()} if isinstance(value, list): - value_list = cast(list[Any], value) - return [self._expand_variable_reference(item) for item in value_list] - + return [self._resolve_value(item, labels) for item in cast(list[Any], value)] if isinstance(value, tuple): - value_tuple = cast(tuple[Any, ...], value) - return tuple(self._expand_variable_reference(item) for item in value_tuple) - + return tuple(self._resolve_value(item, labels) for item in cast(tuple[Any, ...], value)) return value - def _expand_variable_references_in_context(self, context: FunctionInvocationContext) -> None: - """Resolve bracketed variable placeholders in invocation arguments in-place. - - Expands [var_xxx] placeholders to their stored content before tool execution. - Original unexpanded arguments are preserved in metadata for message reconstruction, - ensuring that function_call Content messages keep placeholders hidden from the LLM. - - Tools in ``_VARIABLE_ID_CONSUMERS`` (e.g. ``inspect_variable``) take variable - IDs as literal references and resolve them internally, so their arguments are - left untouched — expanding them would replace the ID with content and break - the lookup. - """ + def _expand_variable_references_in_context(self, context: FunctionInvocationContext) -> list[ContentLabel]: + """Expand owned references in invocation values and return their stored labels.""" if context.function.name in _VARIABLE_ID_CONSUMERS: - return + return [] + labels: list[ContentLabel] = [] if context.arguments: - args_before = str(context.arguments)[:200] if context.arguments else "" - context.arguments = self._expand_variable_reference(context.arguments) - args_after = str(context.arguments)[:200] if context.arguments else "" - has_var_ref_before = "[var_" in args_before - has_var_ref_after = "[var_" in args_after - if has_var_ref_before or has_var_ref_after: - logger.debug( - "Variable expansion for '%s': had_ref_before=%s, had_ref_after=%s", - context.function.name, - has_var_ref_before, - has_var_ref_after, - ) - if has_var_ref_before and not has_var_ref_after: - logger.debug( - "Expanded variable references from: %s... to: %s...", - args_before[:100], - args_after[:100], - ) - + context.arguments = self._resolve_value(context.arguments, labels) if context.kwargs: - context.kwargs = cast(dict[str, Any], self._expand_variable_reference(context.kwargs)) + context.kwargs = cast(dict[str, Any], self._resolve_value(context.kwargs, labels)) + return labels def _get_input_labels(self, context: FunctionInvocationContext) -> list[ContentLabel]: """Extract security labels from tool input arguments. @@ -1482,64 +1498,31 @@ async def process( context: FunctionInvocationContext, call_next: Callable[[], Awaitable[None]], ) -> None: - """Process function invocation with tiered label propagation. - - Label propagation follows a strict 3-tier priority for determining the - result label of a tool call: - - 1. **Tier 1 (Highest)**: Per-item embedded labels in the tool result - (``additional_properties.security_label``). If present, these labels - are used directly for each item. - 2. **Tier 2**: The tool's ``source_integrity`` declaration. If the tool - explicitly declares ``source_integrity`` in its ``additional_properties``, - that declaration alone determines the fallback label (input argument - labels are NOT combined in). - 3. **Tier 3 (Lowest)**: The join (``combine_labels``) of all input argument - labels. Used only when there are no embedded labels AND no - ``source_integrity`` declaration. - - Two metadata keys are set on the context: - - - ``context.metadata["result_label"]``: The security label of THIS tool - call's result (per-call). Set once after result processing. - - ``context.metadata["context_label"]``: The cumulative conversation - security state (cross-call). Used by ``PolicyEnforcementFunctionMiddleware`` - to validate subsequent tool calls. - - Args: - context: The function invocation context. - call_next: Callback to continue to next middleware or function execution. - """ - # Keep security tools bound to this invocation across asyncio task overlap. + """Resolve hidden arguments, publish their labels, and label the result.""" + scope_token = self._activate_security_scope(context) middleware_token = _current_middleware.set(self) - try: function_name = context.function.name + if "original_arguments_for_messages" not in context.metadata: + context.metadata["original_arguments_for_messages"] = deepcopy(context.arguments) + else: + context.arguments = deepcopy(context.metadata["original_arguments_for_messages"]) + if "security_original_runtime_kwargs" not in context.metadata: + context.metadata["security_original_runtime_kwargs"] = dict(context.kwargs) + else: + context.kwargs = dict(cast(dict[str, Any], context.metadata["security_original_runtime_kwargs"])) - # ========== Tiered Label Propagation ========== - # Step 1: Extract labels from input arguments input_labels = self._get_input_labels(context) - - # Step 2: Get tool's source_integrity declaration (may be None) declared_source_integrity = self._get_source_integrity(context) - - # Get confidentiality from function additional_properties or use default confidentiality = self._get_function_confidentiality(context) - # Step 3: Build tiered fallback_label - # This label is used for result items that have NO embedded labels. - # Priority: source_integrity declaration (tier 2) > input labels join (tier 3) if declared_source_integrity is not None: - # Tier 2: Tool explicitly declared source_integrity — use it alone. - # Input argument labels are NOT combined in; the tool's declaration - # is authoritative for the trust level of its output. fallback_label = ContentLabel( integrity=declared_source_integrity, confidentiality=confidentiality, metadata={"source": "source_integrity", "function_name": function_name}, ) elif input_labels: - # Tier 3: No source_integrity declared — join all input labels. combined = combine_labels(*input_labels) fallback_label = ContentLabel( integrity=combined.integrity, @@ -1547,51 +1530,26 @@ async def process( metadata={"source": "input_labels_join", "function_name": function_name}, ) else: - # Tier 3 fallback: No source_integrity AND no input labels. - # Default to UNTRUSTED for safety. fallback_label = ContentLabel( integrity=self.default_integrity, confidentiality=confidentiality, metadata={"source": "default", "function_name": function_name}, ) - # context_label: cumulative conversation security state (cross-call). - # Used by PolicyEnforcementFunctionMiddleware to validate tool calls. - context.metadata["context_label"] = self._context_label - - logger.info( - f"Tool call '{function_name}' fallback label (tiered): " - f"{fallback_label.integrity.value}, {fallback_label.confidentiality.value} " - f"(inputs: {len(input_labels)}, source_integrity: " - f"{declared_source_integrity.value if declared_source_integrity else 'not declared'})" - ) - logger.info( - f"Current context label: {self._context_label.integrity.value}, " - f"{self._context_label.confidentiality.value}" - ) - - # Store original unexpanded arguments for message reconstruction before expanding - if "original_arguments_for_messages" not in context.metadata: - # Deep copy to preserve original state - context.metadata["original_arguments_for_messages"] = deepcopy(context.arguments) - - # Expand bracketed variable references in arguments BEFORE tool execution - # so that tools receive expanded content, but keep originals for message history - self._expand_variable_references_in_context(context) + resolved_labels = self._expand_variable_references_in_context(context) + argument_label = combine_labels(*resolved_labels) if resolved_labels else ContentLabel() + context_label = self._context_label + context.metadata["context_label"] = context_label + context.metadata["argument_label"] = argument_label + context.metadata["effective_invocation_label"] = combine_labels(context_label, argument_label) - # Execute the function await call_next() - - # If middleware set a function_approval_request (e.g., policy violation approval), - # skip all result processing and let it pass through unchanged if isinstance(context.result, Content) and context.result.type == "function_approval_request": - logger.info(f"Tool '{function_name}' returned function_approval_request - skipping result processing") return - - # Label, hide, and update context label for the tool result self._label_result(context, function_name, fallback_label) finally: _current_middleware.reset(middleware_token) + self._active_security_scope.reset(scope_token) def _label_result( self, @@ -1802,7 +1760,7 @@ def _hide_item( # Store the actual content (serialize Content to its text representation) stored_value: Any = item.text if item.type == "text" and item.text is not None else item.to_dict() - var_id = self._variable_store.store(stored_value, label) + var_id = self.get_variable_store().store(stored_value, label) # Store metadata about this variable self._variable_metadata[var_id] = { @@ -1828,37 +1786,23 @@ def _hide_item( additional_properties={"_variable_reference": True, "security_label": label.to_dict()}, ) - def get_variable_store(self) -> ContentVariableStore: - """Get the variable store for this middleware instance. - - Returns: - The ContentVariableStore instance. - """ - return self._variable_store - - def get_variable_metadata(self, var_id: str) -> dict[str, Any] | None: - """Get metadata for a stored variable. - - Args: - var_id: The variable ID. + def get_variable_store(self, session: AgentSession | None = None) -> ContentVariableStore: + """Get the owner-aware variable store for an optional explicit session.""" + return self._scope_for_session(session).variable_store() - Returns: - Metadata dictionary or None if not found. - """ - if not self._variable_store.exists(var_id): + def get_variable_metadata(self, var_id: str, session: AgentSession | None = None) -> dict[str, Any] | None: + """Get variable metadata for an optional explicit session.""" + scope = self._scope_for_session(session) + if not scope.variable_store().exists(var_id): return None - metadata = self._variable_metadata.get(var_id) + metadata = scope.variable_metadata.get(var_id) if not isinstance(metadata, dict): return None return deepcopy(cast(dict[str, Any], metadata)) - def list_variables(self) -> list[str]: - """Get a list of all stored variable IDs. - - Returns: - List of variable ID strings. - """ - return self._variable_store.list_variables() + def list_variables(self, session: AgentSession | None = None) -> list[str]: + """List variable IDs for an optional explicit session.""" + return self.get_variable_store(session).list_variables() def get_security_tools(self) -> list[FunctionTool]: """Get the list of security tools for agent integration. @@ -1934,58 +1878,46 @@ def get_current_middleware() -> LabelTrackingFunctionMiddleware | None: class _PendingPolicyApproval(NamedTuple): - """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 - arguments; ``label_key`` the security label (integrity/confidentiality) shown for review; - ``session_key`` the session the approval was requested in (the isolation boundary at this layer; - 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. - """ + """Exact, durable binding for one pending policy approval.""" body_signature: str + resolved_signature: str label_key: str + effective_label_key: str session_key: str disclosed_violations: tuple[str, ...] def to_state(self) -> dict[str, Any]: - """Return the JSON-compatible representation stored in session state.""" return { "body_signature": self.body_signature, + "resolved_signature": self.resolved_signature, "label_key": self.label_key, + "effective_label_key": self.effective_label_key, "session_key": self.session_key, "disclosed_violations": list(self.disclosed_violations), } @classmethod def from_state(cls, payload: Any) -> _PendingPolicyApproval | None: - """Restore a pending record, failing closed for malformed state.""" if not isinstance(payload, dict): return None record = cast(dict[str, Any], payload) - body_signature = record.get("body_signature") - label_key = record.get("label_key") - session_key = record.get("session_key") + 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") - if not all(isinstance(value, str) for value in (body_signature, label_key, session_key)): + if not all(type(value) is str for value in values): return None if not isinstance(violations, list): return None violation_items = cast(list[Any], violations) - if not all(isinstance(item, str) for item in violation_items): + if not all(type(item) is str for item in violation_items): return None - return cls( - body_signature=cast(str, body_signature), - label_key=cast(str, label_key), - session_key=cast(str, session_key), - disclosed_violations=tuple(cast(list[str], violation_items)), - ) + typed_values = cast(tuple[str, str, str, str, str], values) + return cls(*typed_values, tuple(cast(list[str], violation_items))) @experimental(feature_id=ExperimentalFeature.FIDES) -class PolicyEnforcementFunctionMiddleware(FunctionMiddleware): +class PolicyEnforcementFunctionMiddleware(FunctionMiddleware, _SecurityScopeBinding): """Middleware that enforces security policies on tool invocations. This middleware: @@ -2016,55 +1948,45 @@ class PolicyEnforcementFunctionMiddleware(FunctionMiddleware): ) """ + _requires_session_state = True + def __init__( self, allow_untrusted_tools: set[str] | None = None, block_on_violation: bool = True, enable_audit_log: bool = True, approval_on_violation: bool = False, + *, + security_scope: _SecurityScope | None = None, + session_state_key: str = _STANDALONE_SESSION_STATE_KEY, ) -> None: - """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. - """ + """Initialize policy enforcement and bind its security-state selector.""" 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._security_scope = _SecurityScope() + self._initialize_security_scope(security_scope, session_state_key=session_state_key) def _clone_for_scope(self, scope: _SecurityScope) -> PolicyEnforcementFunctionMiddleware: - """Clone current middleware configuration into a session scope.""" + """Clone customized middleware configuration into a fixed session scope.""" scoped = copy(self) - scoped._security_scope = scope + scoped._initialize_security_scope(scope, session_state_key=self._security_session_state_key) return scoped @property def audit_log(self) -> list[dict[str, Any]]: - """Return the audit log view for this middleware's fixed scope.""" - return self._security_scope.audit_log + """Return the live audit log for the active scope.""" + return self._scope.audit_log @property def _pending_policy_approvals(self) -> dict[str, Any]: - """Return serialized pending approvals for this middleware's fixed scope.""" - return self._security_scope.pending_approvals + return self._scope.pending_approvals def _get_pending_approval(self, approval_id: str) -> _PendingPolicyApproval | None: - """Restore one pending approval from scope state.""" - return _PendingPolicyApproval.from_state(self._pending_policy_approvals.get(approval_id)) + return _PendingPolicyApproval.from_state(self._scope.pending_approvals.get(approval_id)) def _store_pending_approval(self, approval_id: str, record: _PendingPolicyApproval) -> None: - """Store one pending approval in detached JSON-compatible form.""" - self._pending_policy_approvals[approval_id] = record.to_state() + self._scope.pending_approvals[approval_id] = record.to_state() def _get_call_id(self, context: FunctionInvocationContext) -> str: """Get the tool call id for this invocation context.""" @@ -2102,31 +2024,42 @@ def _build_function_call_content(self, context: FunctionInvocationContext) -> Co ) def _signature_from_parts(self, name: str | None, arguments: dict[str, Any]) -> str: - """Canonicalize a (function name, arguments) pair into a stable comparison signature.""" - try: - arguments_repr = json.dumps(arguments, sort_keys=True, default=str) - except (TypeError, ValueError): - arguments_repr = repr(sorted(arguments.items())) - return f"{name or ''}\x00{arguments_repr}" + """Hash a deterministic, strictly representable invocation snapshot.""" + canonical = _strict_json_value( + {"name": name or "", "arguments": arguments}, + path="approval invocation", + canonical=True, + ) + payload = json.dumps( + canonical, + ensure_ascii=False, + allow_nan=False, + separators=(",", ":"), + ).encode() + return hashlib.sha256(payload).hexdigest() def _call_body_signature(self, context: FunctionInvocationContext) -> str: - """Compute the (function name, arguments) signature for the current invocation. - - This is the part of the binding that the approval *response*'s embedded ``function_call`` - can also reproduce, so it is used both to validate the response body and to check that the - invocation about to execute matches the one that was approved. - """ return self._signature_from_parts(context.function.name, self._current_arguments(context)) + def _resolved_arguments(self, context: FunctionInvocationContext) -> dict[str, Any]: + if isinstance(context.arguments, BaseModel): + return context.arguments.model_dump() + return dict(context.arguments) + + def _resolved_call_signature(self, context: FunctionInvocationContext) -> str: + return self._signature_from_parts( + context.function.name, + {"arguments": self._resolved_arguments(context), "kwargs": dict(context.kwargs)}, + ) + def _context_label_key(self, context: FunctionInvocationContext) -> str: - """Canonicalize the current security label (integrity/confidentiality) for binding. + return self._label_key(context.metadata.get("context_label")) - Accepts a ``ContentLabel`` or its dict form from ``context.metadata['context_label']`` and - reduces it to the security-relevant dimensions used for policy decisions, so an approval - granted under one label cannot authorize the same call under a different (e.g. more - sensitive) label. - """ - label_data = context.metadata.get("context_label") + def _effective_label_key(self, context: FunctionInvocationContext) -> str: + return self._label_key(context.metadata.get("effective_invocation_label")) + + @staticmethod + def _label_key(label_data: Any) -> str: if isinstance(label_data, ContentLabel): return f"{label_data.integrity.value}/{label_data.confidentiality.value}" if isinstance(label_data, dict): @@ -2135,40 +2068,29 @@ def _context_label_key(self, context: FunctionInvocationContext) -> str: return "/" def _session_key(self, context: FunctionInvocationContext) -> str: - """Return the session id for binding, or empty string when there is no session.""" - session = context.session - return session.session_id if session is not None else "" + return context.session.session_id if context.session is not None else "" def _violation_set_key(self, violations: list[dict[str, Any]]) -> tuple[str, ...]: - """Canonicalize the disclosed violations into a stable, order-independent key. - - Each entry pairs the violation type with its canonical policy reason, so a replay that - trips the *same* violation type but a materially *different* risk (e.g. the tool's - ``max_allowed_confidentiality`` destination changed, keeping the type but changing the - reason) no longer matches and must re-request approval. The approval is thus bound to - exactly the risks disclosed to the user, not merely to their category names. - """ - return tuple(sorted(f"{v['violation_type']}\x00{v['audit']['reason']}" for v in violations)) + return tuple(sorted(f"{item['violation_type']}\x00{item['audit']['reason']}" for item in violations)) def _pending_record( self, context: FunctionInvocationContext, violations: list[dict[str, Any]], ) -> _PendingPolicyApproval: - """Build the binding record for the current invocation and disclosed violation set.""" return _PendingPolicyApproval( body_signature=self._call_body_signature(context), + resolved_signature=self._resolved_call_signature(context), label_key=self._context_label_key(context), + effective_label_key=self._effective_label_key(context), session_key=self._session_key(context), disclosed_violations=self._violation_set_key(violations), ) 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"): return None - arguments = function_call.parse_arguments() or {} - return self._signature_from_parts(function_call.name, dict(arguments)) + return self._signature_from_parts(function_call.name, dict(function_call.parse_arguments() or {})) def _response_matches_pending( self, @@ -2177,43 +2099,22 @@ def _response_matches_pending( call_id: str, body_signature: str, ) -> bool: - """Validate that the approval response itself corresponds to the pending request. - - The response must carry the request id shown for review and embed the exact function call - (name + arguments) reconstructed from the authoritative pending snapshot. The embedded - ``call_id`` must match provider correlation; an occurrence-aware response and embedded call - must also carry the Agent Framework approval id. Legacy direct middleware callers continue - to use ``call_id`` for both identities. - """ - embedded = getattr(approval_response, "function_call", None) + embedded = approval_response.function_call if self._signature_from_function_call(embedded) != body_signature: return False - # Both identifiers must be present and name the pending request (no None bypass). - response_id = getattr(approval_response, "id", None) - embedded_call_id = getattr(embedded, "call_id", None) - embedded_occurrence_id = getattr(embedded, "id", None) + if embedded is None: + return False return ( - response_id == approval_id - and embedded_call_id == call_id - and (approval_id == call_id or embedded_occurrence_id == approval_id) + approval_response.id == approval_id + and embedded.call_id == call_id + and (approval_id == call_id or embedded.id == approval_id) ) def _matches_pending_approval( self, context: FunctionInvocationContext, - current_violations: list[dict[str, Any]], + current_binding: _PendingPolicyApproval, ) -> bool: - """Return whether an approved, call-bound approval matches this exact invocation. - - True only when an approved ``function_approval_response`` is present, the call_id is still - awaiting approval, the response itself (its id and embedded ``function_call``) matches the - pending request, and the current invocation matches every bound dimension recorded when - approval was requested: function + arguments, the security label, the session, and the - exact set of violation types disclosed to the user. If the invocation now trips a different - set of violations than was disclosed, this returns False so the caller re-requests approval - for the new set. Does not mutate state; consumption is done separately via - :meth:`_consume_pending_approval` once the approval actually waves the detected violations. - """ call_id = self._get_call_id(context) approval_id = self._get_approval_id(context) if not call_id or not approval_id: @@ -2228,23 +2129,11 @@ def _matches_pending_approval( and approval_response.approved is True ): return False - # The approval response must itself match the pending request (id + embedded function_call), - # and the invocation about to execute must match every recorded binding dimension, including - # the exact set of violations that was disclosed for review. - return ( - self._response_matches_pending(approval_response, approval_id, call_id, pending.body_signature) - and self._call_body_signature(context) == pending.body_signature - and self._context_label_key(context) == pending.label_key - and self._session_key(context) == pending.session_key - and self._violation_set_key(current_violations) == pending.disclosed_violations + return current_binding == pending and self._response_matches_pending( + approval_response, approval_id, call_id, pending.body_signature ) def _consume_pending_approval(self, context: FunctionInvocationContext) -> None: - """Remove the pending approval for this call so it authorizes exactly one invocation. - - Idempotent: safe to call for both the integrity and confidentiality checks of a single - invocation. - """ self._pending_policy_approvals.pop(self._get_approval_id(context), None) def _mark_policy_violation_approved( @@ -2263,6 +2152,7 @@ def _request_policy_violation_approval( *, context_label: ContentLabel, violations: list[dict[str, Any]], + binding: _PendingPolicyApproval, ) -> None: """Create a single policy-violation approval request disclosing every detected violation. @@ -2278,8 +2168,16 @@ def _request_policy_violation_approval( ) approval_id = self._get_approval_id(context) if approval_id: - self._store_pending_approval(approval_id, self._pending_record(context, violations)) + self._store_pending_approval(approval_id, binding) + approval_response = context.metadata.get("approval_response") + is_replacement = ( + isinstance(approval_response, Content) + and approval_response.type == "function_approval_response" + and approval_response.approved is True + ) + request_id = f"{approval_id}:replacement:{uuid.uuid4().hex}" if is_replacement else approval_id additional_properties: dict[str, Any] = { + "_replacement_approval_request": is_replacement, "policy_violation": True, "violation_type": primary["violation_type"], "reason": ( @@ -2295,7 +2193,7 @@ def _request_policy_violation_approval( {"violation_type": v["violation_type"], "reason": v["approval_reason"]} for v in violations ] context.result = Content.from_function_approval_request( - id=approval_id, + id=request_id, function_call=self._build_function_call_content(context), additional_properties=additional_properties, ) @@ -2322,73 +2220,66 @@ def _block_policy_violation( context.result = result raise MiddlewareTermination("Policy violation blocked tool execution") + def _block_unsafe_approval_binding( + self, context: FunctionInvocationContext, *, context_label: ContentLabel + ) -> NoReturn: + context.result = { + "error": "Policy violation: approval cannot be safely bound to this invocation", + "function": context.function.name, + "context_label": context_label.to_dict(), + "violation_type": "unsafe_approval_binding", + } + raise MiddlewareTermination("Unsafe policy approval binding") + async def process( self, context: FunctionInvocationContext, call_next: Callable[[], Awaitable[None]], ) -> None: - """Process function invocation with policy enforcement. - - Policy enforcement uses the context_label (cumulative security state of the - conversation) to validate tool calls. This prevents indirect attacks where - untrusted content from previous tool calls could influence dangerous operations. + """Enforce policy using the scope selected from this invocation.""" + scope_token = self._activate_security_scope(context) + try: + await self._process_in_scope(context, call_next) + finally: + self._active_security_scope.reset(scope_token) - Args: - context: The function invocation context. - call_next: Callback to continue to next middleware or function execution. - """ + async def _process_in_scope( + self, + context: FunctionInvocationContext, + call_next: Callable[[], Awaitable[None]], + ) -> None: function_name = context.function.name - - # Get the context label (cumulative security state of the conversation) - # This is set by LabelTrackingFunctionMiddleware and represents the - # combined security state of all content that has entered the context context_label_data = context.metadata.get("context_label") - if context_label_data is None: logger.warning( - f"No context label found for tool '{function_name}'. " - "Ensure LabelTrackingFunctionMiddleware runs before PolicyEnforcementFunctionMiddleware." + "No context label found for tool '%s'. Ensure LabelTrackingFunctionMiddleware runs first.", + function_name, ) - # Continue execution without policy check await call_next() return - - # Convert context label to ContentLabel if it's a dict if isinstance(context_label_data, dict): context_label = ContentLabel.from_dict(cast(dict[str, Any], context_label_data)) elif isinstance(context_label_data, ContentLabel): context_label = context_label_data else: - logger.error(f"Invalid context label type: {type(context_label_data)}") + logger.error("Invalid context label type: %s", type(cast(object, context_label_data)).__name__) await call_next() return - logger.debug( - f"Policy enforcement for '{function_name}': " - f"context_label={context_label.integrity.value}/{context_label.confidentiality.value}" - ) + argument_label = self._resolve_label(context.metadata.get("argument_label")) + effective_label = combine_labels(context_label, argument_label) function_props = _get_additional_properties(context.function) - - # Detect every applicable policy violation up front so a single approval decision can - # disclose all of them together. Evaluating both the integrity and confidentiality checks - # before acting prevents an approval that was requested (and granted) for one violation - # from silently waving a second, undisclosed violation when the call is replayed. + accepts_untrusted = ( + function_name in self.allow_untrusted_tools or function_props.get("accepts_untrusted") is True + ) violations: list[dict[str, Any]] = [] - # Integrity policy: an UNTRUSTED (tainted) context may not drive a tool that has not - # opted in to untrusted input. - if ( - context_label.integrity == IntegrityLabel.UNTRUSTED - and function_name not in self.allow_untrusted_tools - and not function_props.get("accepts_untrusted", False) - ): + if context_label.integrity == IntegrityLabel.UNTRUSTED and not accepts_untrusted: violations.append({ "violation_type": "untrusted_context", "approval_reason": ( f"Tool '{function_name}' is being called in an UNTRUSTED context. " - "The conversation contains data from untrusted sources which could " - "influence this operation. Approve to proceed anyway (the agent will " - "continue with a warning about untrusted context)." + "Approve to proceed despite possible untrusted influence." ), "block_error": "Policy violation: Tool cannot be called in untrusted context", "block_violation_type": None, @@ -2397,19 +2288,37 @@ async def process( "function": function_name, "context_label": context_label.to_dict(), "turn": context.metadata.get("turn_number", -1), - "reason": "Context is UNTRUSTED and tool is not allowed to execute in an untrusted context", + "reason": "Context is UNTRUSTED and the tool does not accept untrusted input", + }, + }) + + if argument_label.integrity == IntegrityLabel.UNTRUSTED and not accepts_untrusted: + violations.append({ + "violation_type": "untrusted_arguments", + "approval_reason": ( + f"Tool '{function_name}' would receive UNTRUSTED content resolved from a hidden variable. " + "Approve to forward that content." + ), + "block_error": "Policy violation: Tool cannot receive untrusted variable content", + "block_violation_type": "untrusted_arguments", + "audit": { + "type": "untrusted_arguments", + "function": function_name, + "context_label": context_label.to_dict(), + "argument_label": argument_label.to_dict(), + "effective_label": effective_label.to_dict(), + "turn": context.metadata.get("turn_number", -1), + "reason": "Arguments resolve to UNTRUSTED content and the tool does not accept untrusted input", }, }) - # Confidentiality policy: block writing higher-confidentiality data to a lower - # confidentiality destination (data exfiltration). - conf_result = self._check_confidentiality_policy_detailed(context, context_label) + conf_result = self._check_confidentiality_policy_detailed(context, effective_label) if not conf_result["passed"]: violations.append({ "violation_type": conf_result["failure_type"], "approval_reason": ( - f"Tool '{function_name}' violates confidentiality policy: " - f"{conf_result['reason']}. Approve to proceed anyway." + f"Tool '{function_name}' violates confidentiality policy: {conf_result['reason']}. " + "Approve to proceed anyway." ), "block_error": f"Policy violation: {conf_result['reason']}", "block_violation_type": conf_result["failure_type"], @@ -2418,60 +2327,60 @@ async def process( "subtype": conf_result["failure_type"], "function": function_name, "context_label": context_label.to_dict(), + "argument_label": argument_label.to_dict(), + "effective_label": effective_label.to_dict(), "reason": conf_result["reason"], "turn": context.metadata.get("turn_number", -1), }, }) if not violations: - # Policy check passed, continue execution - logger.debug(f"Policy check passed for tool '{function_name}'") await call_next() return - for violation in violations: self._log_violation(violation["audit"]) - # Resolve the approval decision against the exact violation set now detected. A pending - # approval only counts if it was granted for this same set; a replay that trips a - # different set (e.g. after the tool's policy metadata changed) falls through and - # re-requests approval for the new set. - approved = self._matches_pending_approval(context, violations) - - disclosed = ", ".join(v["violation_type"] for v in violations) + binding: _PendingPolicyApproval | None = None + approved = False + if self.approval_on_violation: + try: + binding = self._pending_record(context, violations) + except (TypeError, ValueError, OverflowError): + self._block_unsafe_approval_binding(context, context_label=context_label) + approved = self._matches_pending_approval(context, binding) + disclosed = ", ".join(item["violation_type"] for item in violations) if approved: - # A single approval waves every violation it disclosed for this exact invocation; - # consume it once so it cannot authorize a repeated or different call. self._consume_pending_approval(context) self._mark_policy_violation_approved( context, - warning_message=( - f"APPROVED BY USER: Tool '{function_name}' executing despite policy " - f"violation(s) [{disclosed}]. User acknowledged the security risk and " - "approved execution." - ), + warning_message=f"APPROVED BY USER: '{function_name}' executing despite [{disclosed}].", ) - elif self.approval_on_violation: + elif binding is not None: self._request_policy_violation_approval( context, context_label=context_label, violations=violations, + binding=binding, ) - return elif self.block_on_violation: - logger.warning(f"BLOCKED: Tool '{function_name}' policy violation(s): {disclosed}") - self._block_policy_violation( - context, - context_label=context_label, - violations=violations, - ) - return + self._block_policy_violation(context, context_label=context_label, violations=violations) else: - logger.warning(f"WARNING: Tool '{function_name}' policy violation(s) [{disclosed}] (allowed)") + logger.warning("WARNING: Tool '%s' policy violation(s) [%s] (allowed)", function_name, disclosed) await call_next() + @staticmethod + def _resolve_label(label_data: Any) -> ContentLabel: + if isinstance(label_data, ContentLabel): + return label_data + if isinstance(label_data, dict): + try: + return ContentLabel.from_dict(cast(dict[str, Any], label_data)) + except (TypeError, ValueError): + logger.warning("Ignoring unparseable invocation label.") + return ContentLabel() + def _check_confidentiality_policy( self, context: FunctionInvocationContext, @@ -2545,22 +2454,18 @@ def _log_violation(self, violation: dict[str, Any]) -> None: violation: Dictionary containing violation details. """ if self.enable_audit_log: - self._security_scope.append_audit_entry(violation) + self._scope.append_audit_entry(violation) logger.warning("Policy violation detected") logger.debug("Policy violation details: %s", violation) - def get_audit_log(self) -> list[dict[str, Any]]: - """Get the audit log of policy violations. - - Returns: - List of violation records. - """ - return self._security_scope.get_audit_log() + def get_audit_log(self, session: AgentSession | None = None) -> list[dict[str, Any]]: + """Get a detached audit log for an optional explicit session.""" + return self._scope_for_session(session).get_audit_log() - def clear_audit_log(self) -> None: - """Clear the audit log.""" - self._security_scope.clear_audit_log() + def clear_audit_log(self, session: AgentSession | None = None) -> None: + """Clear the audit log for an optional explicit session.""" + self._scope_for_session(session).clear_audit_log() @experimental(feature_id=ExperimentalFeature.FIDES) @@ -2654,26 +2559,31 @@ class docstring for details on running multiple instances. Defaults to "secure_agent". """ super().__init__(source_id or self.DEFAULT_SOURCE_ID) - + self._auto_hide_untrusted = auto_hide_untrusted + self._default_integrity = default_integrity + self._default_confidentiality = default_confidentiality + self._allow_untrusted_tools = {"quarantined_llm", "inspect_variable"} + if allow_untrusted_tools: + self._allow_untrusted_tools.update(allow_untrusted_tools) + self._block_on_violation = block_on_violation + self._approval_on_violation = approval_on_violation + self._enable_audit_log = enable_audit_log + self.enable_policy_enforcement = enable_policy_enforcement self.label_tracker = LabelTrackingFunctionMiddleware( auto_hide_untrusted=auto_hide_untrusted, default_integrity=default_integrity, default_confidentiality=default_confidentiality, ) - - self.enable_policy_enforcement = enable_policy_enforcement - if enable_policy_enforcement: - tools_allowing_untrusted = {"quarantined_llm", "inspect_variable"} - if allow_untrusted_tools: - tools_allowing_untrusted.update(allow_untrusted_tools) - self.policy_enforcer: PolicyEnforcementFunctionMiddleware | None = PolicyEnforcementFunctionMiddleware( - allow_untrusted_tools=tools_allowing_untrusted, + self.policy_enforcer = ( + PolicyEnforcementFunctionMiddleware( + allow_untrusted_tools=set(self._allow_untrusted_tools), block_on_violation=block_on_violation, approval_on_violation=approval_on_violation, enable_audit_log=enable_audit_log, ) - else: - self.policy_enforcer = None + if enable_policy_enforcement + else None + ) self._provider_state_used = False # Store and configure quarantine client for real LLM calls @@ -2938,15 +2848,20 @@ def get_quarantine_client() -> SupportsChatGetResponse | None: write_file(path="out.txt", content="[var_abc123]") ``` -**INCORRECT** — do NOT pass the bare id, do NOT quote it, do NOT prefix it: +**INCORRECT** — do not rely on the bare-id safety fallback, quote it, or prefix it: ``` -write_file(content="var_abc123") # ❌ bare id, will be written verbatim +write_file(content="var_abc123") # discouraged; use brackets write_file(content="${var_abc123}") # ❌ wrong syntax write_file(content="") # ❌ wrong syntax ``` Always use the exact form ``[var_]``. The id is opaque — do NOT shorten, -truncate, or modify it. +truncate, or modify it. Known bare IDs are expanded only as a safety fallback. Unknown IDs and IDs +owned by another session remain literal. + +Forwarding is allowed only when the destination tool declares `accepts_untrusted=True`; otherwise +the call is blocked, audited, or sent for policy approval. That opt-in does not bypass the tool's +`max_allowed_confidentiality` limit. Hidden content stays out of the model context. ### Best Practices: 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 18615b7c60f..02a0f05b8ae 100644 --- a/python/packages/core/tests/core/test_function_invocation_logic.py +++ b/python/packages/core/tests/core/test_function_invocation_logic.py @@ -59,6 +59,48 @@ def _build_approved_tool_roundtrip( return function_call, approval_request, approval_response +def test_collect_unanswered_replacement_requests_correlates_reused_call_id_by_occurrence() -> None: + """A resolved replacement must not consume a pending reused-call-id sibling.""" + from agent_framework._tools import _collect_unanswered_approval_requests + + first_call = Content.from_function_call( + call_id="reused", + name="guarded", + arguments="{}", + id="occurrence-1", + ) + first_request = Content.from_function_approval_request( + id="replacement-1", + function_call=first_call, + additional_properties={"_replacement_approval_request": True}, + ) + second_call = Content.from_function_call( + call_id="reused", + name="guarded", + arguments="{}", + id="occurrence-2", + ) + second_request = Content.from_function_approval_request( + id="replacement-2", + function_call=second_call, + additional_properties={"_replacement_approval_request": True}, + ) + second_response = Content.from_function_approval_response( + approved=True, + id="occurrence-2", + function_call=second_call, + ) + + unanswered = _collect_unanswered_approval_requests([ + Message(role="assistant", contents=[first_request]), + Message(role="assistant", contents=[second_request]), + Message(role="user", contents=[second_response]), + Message(role="tool", contents=[Content.from_function_result(call_id="reused", result="done")]), + ]) + + assert unanswered == [first_request] + + def test_session_approval_binding_rebinds_consumes_and_rejects_duplicates() -> None: """Session binding must use the recorded call and honor one response once.""" from agent_framework._tools import ( @@ -539,6 +581,15 @@ def fragment(call_id: str, name: str, arguments: str, index: int) -> Content: assert caught == [] +def test_loading_pending_approval_requests_does_not_create_state() -> None: + from agent_framework._tools import _load_pending_approval_requests + + session = AgentSession(session_id="approval-read-only") + + assert _load_pending_approval_requests(session) == {} + assert "tool_approval" not in session.state + + def test_occurrence_aware_approval_rejects_stale_reused_call_id_response(caplog: pytest.LogCaptureFixture) -> None: from agent_framework._tools import ( _bind_approval_responses_to_pending_requests, 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 b637bc751e5..032309accd3 100644 --- a/python/packages/core/tests/core/test_harness_tool_approval.py +++ b/python/packages/core/tests/core/test_harness_tool_approval.py @@ -2,12 +2,16 @@ from __future__ import annotations +import json import warnings -from collections.abc import MutableSequence +from collections.abc import Awaitable, Callable, MutableSequence +from enum import Enum from pathlib import Path from typing import Any +from uuid import UUID import pytest +from pydantic import create_model from agent_framework import ( DEFAULT_TOOL_APPROVAL_SOURCE_ID, @@ -17,6 +21,10 @@ ChatResponseUpdate, Content, FileHistoryProvider, + FunctionInvocationContext, + FunctionMiddleware, + FunctionTool, + InMemoryHistoryProvider, Message, ToolApprovalMiddleware, ToolApprovalState, @@ -25,6 +33,12 @@ tool, ) from agent_framework._feature_stage import ExperimentalWarning +from agent_framework.security import ( + ContentLabel, + IntegrityLabel, + LabelTrackingFunctionMiddleware, + PolicyEnforcementFunctionMiddleware, +) from .conftest import MockBaseChatClient @@ -40,6 +54,610 @@ def _function_call(request: Content) -> Content: return request.function_call +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() + + +async def test_manual_fides_no_session_preserves_standard_tool_approval( + chat_client_base: MockBaseChatClient, +) -> None: + """Run-local FIDES state must not make ordinary no-session approval authoritative.""" + calls: list[str] = [] + + @tool(name="approved_tool", approval_mode="always_require") + def approved_tool() -> str: + calls.append("approved") + return "approved" + + @tool(name="safe_tool") + def safe_tool() -> str: + calls.append("safe") + return "safe" + + agent = Agent( + client=chat_client_base, + tools=[approved_tool, safe_tool], + middleware=[LabelTrackingFunctionMiddleware(), PolicyEnforcementFunctionMiddleware()], + ) + chat_client_base.run_responses = [ + ChatResponse( + messages=Message( + role="assistant", + contents=[ + Content.from_function_call( + call_id="approved-call", + name="approved_tool", + arguments="{}", + id="approved-occurrence", + ), + Content.from_function_call( + call_id="safe-call", + name="safe_tool", + arguments="{}", + id="safe-occurrence", + ), + ], + ) + ), + ChatResponse(messages=Message(role="assistant", contents=["done"])), + ] + + first = await agent.run("request approval") + assert {request.id for request in first.user_input_requests} == { + "approved-occurrence", + "safe-occurrence", + } + resumed = await agent.run( + Message( + role="user", + contents=[request.to_function_approval_response(True) for request in first.user_input_requests], + ) + ) + + assert calls == ["approved", "safe"] + assert [(message.role, [content.type for content in message.contents]) for message in resumed.messages] == [ + ("tool", ["function_result", "function_result"]), + ("assistant", ["text"]), + ] + + +class _StringApprovalValue(str, Enum): + ALPHA = "alpha" + + +class _IntApprovalValue(int, Enum): + ONE = 1 + + +@pytest.mark.parametrize("max_iterations", [3], indirect=True) +async def test_manual_fides_no_session_uses_isolated_run_scope( + chat_client_base: MockBaseChatClient, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Manual FIDES middleware shares one run scope without leaking it to a later run.""" + monkeypatch.setattr( + "agent_framework.security.uuid.uuid4", + lambda: UUID("01234567-89ab-cdef-0123-456789abcdef"), + ) + hidden_id = "var_0123456789abcdef" + received: list[str] = [] + run_sessions: list[AgentSession | None] = [] + + @tool(name="untrusted_source", additional_properties={"source_integrity": "untrusted"}) + def untrusted_source(ctx: FunctionInvocationContext) -> str: + run_sessions.append(ctx.session) + return "hidden payload" + + @tool( + name="forward_sink", + additional_properties={"accepts_untrusted": True, "source_integrity": "trusted"}, + ) + def forward_sink(value: str, ctx: FunctionInvocationContext) -> str: + run_sessions.append(ctx.session) + received.append(value) + return "sent" + + tracker = LabelTrackingFunctionMiddleware() + policy = PolicyEnforcementFunctionMiddleware() + agent = Agent( + client=chat_client_base, + tools=[untrusted_source, forward_sink], + middleware=[tracker, policy], + ) + chat_client_base.run_responses = [ + ChatResponse( + messages=Message( + role="assistant", + contents=[ + Content.from_function_call( + call_id="source-call", + name="untrusted_source", + arguments="{}", + ) + ], + ) + ), + ChatResponse( + messages=Message( + role="assistant", + contents=[ + Content.from_function_call( + call_id="forward-same-run", + name="forward_sink", + arguments={"value": f"[{hidden_id}]"}, + ) + ], + ) + ), + ChatResponse(messages=Message(role="assistant", contents=["first done"])), + ChatResponse( + messages=Message( + role="assistant", + contents=[ + Content.from_function_call( + call_id="forward-next-run", + name="forward_sink", + arguments={"value": f"[{hidden_id}]"}, + ) + ], + ) + ), + ChatResponse(messages=Message(role="assistant", contents=["second done"])), + ] + + first = await agent.run("first run") + second = await agent.run("second run") + + assert first.text == "first done" + assert second.text == "second done" + assert received == ["hidden payload", f"[{hidden_id}]"] + assert run_sessions[0] is not None + assert run_sessions[0] is run_sessions[1] + assert run_sessions[2] is not None + assert run_sessions[2] is not run_sessions[0] + assert tracker.list_variables() == [] + + +@pytest.mark.parametrize( + ("enum_type", "wire_value", "expected"), + [ + pytest.param(_StringApprovalValue, "alpha", _StringApprovalValue.ALPHA, id="string-enum"), + pytest.param(_IntApprovalValue, 1, _IntApprovalValue.ONE, id="int-enum"), + ], +) +@pytest.mark.parametrize("streaming", [False, True], ids=["non-streaming", "streaming"]) +async def test_policy_approval_resume_supports_enum_arguments( + chat_client_base: MockBaseChatClient, + monkeypatch: pytest.MonkeyPatch, + enum_type: type[Enum], + wire_value: str | int, + expected: Enum, + streaming: bool, +) -> None: + """Strict approval binding accepts Pydantic-validated string and integer enums.""" + input_model = create_model("PolicyEnumArguments", value=(enum_type, ...)) + received: list[Enum] = [] + + def guarded_enum(value: Enum) -> str: + received.append(value) + return "approved enum" + + guarded_tool = FunctionTool( + func=guarded_enum, + name="guarded_enum", + description="Use one enum value", + input_model=input_model, + ) + policy = PolicyEnforcementFunctionMiddleware(approval_on_violation=True) + agent = Agent( + client=chat_client_base, + tools=[guarded_tool], + middleware=[_MarkUntrusted(), policy], + context_providers=[InMemoryHistoryProvider()], + ) + session = AgentSession(session_id=f"enum-policy-{enum_type.__name__}-{streaming}") + function_call = Content.from_function_call( + call_id="enum-call", + name="guarded_enum", + arguments={"value": wire_value}, + id="enum-occurrence", + ) + captured_model_calls: list[list[Message]] = [] + + if streaming: + original_stream = chat_client_base._get_streaming_response + + def capture_stream( + *, + messages: MutableSequence[Message], + options: dict[str, Any], + **kwargs: Any, + ) -> Any: + captured_model_calls.append([Message.from_dict(message.to_dict()) for message in 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_call])], + [ChatResponseUpdate(role="assistant", contents=[Content.from_text("done")])], + ] + first_stream = agent.run("run enum", stream=True, session=session) + first_updates = [update async for update in first_stream] + first_response = await first_stream.get_final_response() + assert [content.type for update in first_updates for content in update.contents] == [ + "function_call", + "function_approval_request", + ] + else: + original_response = chat_client_base._get_non_streaming_response + + async def capture_response( + *, + messages: MutableSequence[Message], + options: dict[str, Any], + **kwargs: Any, + ) -> ChatResponse: + captured_model_calls.append([Message.from_dict(message.to_dict()) for message in 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_call])), + ChatResponse(messages=Message(role="assistant", contents=["done"])), + ] + first_response = await agent.run("run enum", session=session) + + assert received == [] + request = first_response.user_input_requests[0] + assert request.id == "enum-occurrence" + + if streaming: + resumed_stream = agent.run(request.to_function_approval_response(True), stream=True, session=session) + resumed_updates = [update async for update in resumed_stream] + resumed = await resumed_stream.get_final_response() + assert [(update.role, [content.type for content in update.contents]) for update in resumed_updates] == [ + ("tool", ["function_result"]), + ("assistant", ["text"]), + ] + else: + resumed = await agent.run(request.to_function_approval_response(True), session=session) + + assert received == [expected] + assert [(message.role, [content.type for content in message.contents]) for message in resumed.messages] == [ + ("tool", ["function_result"]), + ("assistant", ["text"]), + ] + result = resumed.messages[0].contents[0] + assert result.call_id == "enum-call" + model_contents = [content for message in captured_model_calls[-1] for content in message.contents] + 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] == ["enum-call"] + assert [content.call_id for content in model_results] == ["enum-call"] + assert not any(content.type.startswith("function_approval_") for content in model_contents) + + +@pytest.mark.parametrize("streaming", [False, True], ids=["non-streaming", "streaming"]) +async def test_changed_hidden_snapshot_requires_visible_second_approval( + chat_client_base: MockBaseChatClient, + monkeypatch: pytest.MonkeyPatch, + streaming: bool, +) -> None: + """A changed resolved snapshot surfaces a persisted replacement before executing once.""" + received: list[str] = [] + + @tool(name="guarded_sink") + def guarded_sink(value: str) -> str: + received.append(value) + return "approved hidden value" + + tracker = LabelTrackingFunctionMiddleware() + policy = PolicyEnforcementFunctionMiddleware(approval_on_violation=True) + session = AgentSession(session_id=f"changed-hidden-{streaming}") + variable_id = tracker.get_variable_store(session).store( + "original", + ContentLabel(integrity=IntegrityLabel.UNTRUSTED), + ) + agent = Agent( + client=chat_client_base, + tools=[guarded_sink], + middleware=[tracker, policy], + context_providers=[InMemoryHistoryProvider()], + ) + function_call = Content.from_function_call( + call_id="hidden-call", + name="guarded_sink", + arguments={"value": f"[{variable_id}]"}, + id="hidden-occurrence", + ) + captured_model_calls: list[list[Message]] = [] + + if streaming: + original_stream = chat_client_base._get_streaming_response + + def capture_stream( + *, + messages: MutableSequence[Message], + options: dict[str, Any], + **kwargs: Any, + ) -> Any: + captured_model_calls.append([Message.from_dict(message.to_dict()) for message in 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_call])], + [ChatResponseUpdate(role="assistant", contents=[Content.from_text("ignored stale replay")])], + [ChatResponseUpdate(role="assistant", contents=[Content.from_text("done")])], + [ChatResponseUpdate(role="assistant", contents=[Content.from_text("later")])], + ] + first_stream = agent.run("run hidden", stream=True, session=session) + _ = [update async for update in first_stream] + first = await first_stream.get_final_response() + else: + original_response = chat_client_base._get_non_streaming_response + + async def capture_response( + *, + messages: MutableSequence[Message], + options: dict[str, Any], + **kwargs: Any, + ) -> ChatResponse: + captured_model_calls.append([Message.from_dict(message.to_dict()) for message in 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_call])), + ChatResponse(messages=Message(role="assistant", contents=["ignored stale replay"])), + ChatResponse(messages=Message(role="assistant", contents=["done"])), + ChatResponse(messages=Message(role="assistant", contents=["later"])), + ] + first = await agent.run("run hidden", session=session) + + original_request = first.user_input_requests[0] + security_state = session.state["__agent_framework_fides_security__"] + security_state["variables"][variable_id]["content"] = json.dumps("changed") + + if streaming: + stale_stream = agent.run( + original_request.to_function_approval_response(True), + stream=True, + session=session, + ) + stale_updates = [update async for update in stale_stream] + stale = await stale_stream.get_final_response() + assert [(update.role, [content.type for content in update.contents]) for update in stale_updates] == [ + ("assistant", ["function_approval_request"]), + ] + else: + stale = await agent.run(original_request.to_function_approval_response(True), session=session) + + assert received == [] + assert chat_client_base.call_count == 1 + replacement = stale.user_input_requests[0] + assert replacement.id != original_request.id + assert replacement.function_call is not None + assert original_request.function_call is not None + assert replacement.function_call.id == original_request.function_call.id + assert replacement.function_call.call_id == original_request.function_call.call_id + pending = session.state["tool_approval"]["pending_approval_requests"] + assert [snapshot["id"] for snapshot in pending] == [replacement.id] + + if streaming: + replay_stream = agent.run( + original_request.to_function_approval_response(True), + stream=True, + session=session, + ) + _ = [update async for update in replay_stream] + await replay_stream.get_final_response() + else: + await agent.run(original_request.to_function_approval_response(True), session=session) + + assert received == [] + assert [snapshot["id"] for snapshot in session.state["tool_approval"]["pending_approval_requests"]] == [ + 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 = await approved_stream.get_final_response() + assert [(update.role, [content.type for content in update.contents]) for update in approved_updates] == [ + ("tool", ["function_result"]), + ("assistant", ["text"]), + ] + else: + approved = await agent.run(replacement.to_function_approval_response(True), session=session) + + assert received == ["changed"] + assert chat_client_base.call_count == 3 + assert [(message.role, [content.type for content in message.contents]) for message in approved.messages] == [ + ("tool", ["function_result"]), + ("assistant", ["text"]), + ] + model_contents = [content for message in captured_model_calls[-1] for content in message.contents] + assert [content.type for content in model_contents].count("function_call") == 1 + assert [content.type for content in model_contents].count("function_result") == 1 + assert not any(content.type.startswith("function_approval_") for content in model_contents) + + +async def test_replacement_approval_preserves_unanswered_reused_call_id_sibling( + chat_client_base: MockBaseChatClient, +) -> None: + """Replacing one approval must not discard an unanswered sibling occurrence.""" + calls = 0 + + @tool(name="guarded_sink") + def guarded_sink(value: str) -> str: + nonlocal calls + calls += 1 + return value + + tracker = LabelTrackingFunctionMiddleware() + policy = PolicyEnforcementFunctionMiddleware(approval_on_violation=True) + session = AgentSession(session_id="replacement-sibling") + first_variable = tracker.get_variable_store(session).store( + "first", + ContentLabel(integrity=IntegrityLabel.UNTRUSTED), + ) + second_variable = tracker.get_variable_store(session).store( + "second", + ContentLabel(integrity=IntegrityLabel.UNTRUSTED), + ) + agent = Agent( + client=chat_client_base, + tools=[guarded_sink], + middleware=[tracker, policy], + context_providers=[InMemoryHistoryProvider()], + ) + chat_client_base.run_responses = [ + ChatResponse( + messages=Message( + role="assistant", + contents=[ + Content.from_function_call( + call_id="reused-call", + name="guarded_sink", + arguments={"value": f"[{first_variable}]"}, + id="first-occurrence", + ), + Content.from_function_call( + call_id="reused-call", + name="guarded_sink", + arguments={"value": f"[{second_variable}]"}, + id="second-occurrence", + ), + ], + ) + ) + ] + + first = await agent.run("run both", session=session) + first_request = next(request for request in first.user_input_requests if request.id == "first-occurrence") + second_request = next(request for request in first.user_input_requests if request.id == "second-occurrence") + security_state = session.state["__agent_framework_fides_security__"] + security_state["variables"][first_variable]["content"] = json.dumps("changed") + + resumed = await agent.run(first_request.to_function_approval_response(True), session=session) + + assert calls == 0 + replacement = next( + request + for request in resumed.user_input_requests + if request.function_call is not None and request.function_call.id == "first-occurrence" + ) + assert replacement.id != first_request.id + assert second_request in resumed.user_input_requests + pending_ids = {snapshot["id"] for snapshot in session.state["tool_approval"]["pending_approval_requests"]} + assert pending_ids == {replacement.id, second_request.id} + + +@pytest.mark.parametrize("streaming", [False, True], ids=["non-streaming", "streaming"]) +async def test_dynamic_policy_approval_partitions_safe_sibling_result_roles( + chat_client_base: MockBaseChatClient, + streaming: bool, +) -> None: + """A safe sibling result remains tool-role when dynamic policy asks for approval.""" + safe_calls = 0 + guarded_values: list[str] = [] + + @tool(name="safe_tool", additional_properties={"source_integrity": "trusted"}) + def safe_tool() -> str: + nonlocal safe_calls + safe_calls += 1 + return "safe result" + + @tool(name="guarded_sink", additional_properties={"source_integrity": "trusted"}) + def guarded_sink(value: str) -> str: + guarded_values.append(value) + return "guarded result" + + tracker = LabelTrackingFunctionMiddleware() + policy = PolicyEnforcementFunctionMiddleware(approval_on_violation=True) + session = AgentSession(session_id=f"mixed-policy-{streaming}") + variable_id = tracker.get_variable_store(session).store( + "hidden payload", + ContentLabel(integrity=IntegrityLabel.UNTRUSTED), + ) + agent = Agent( + client=chat_client_base, + tools=[safe_tool, guarded_sink], + middleware=[tracker, policy], + context_providers=[InMemoryHistoryProvider()], + ) + calls = [ + Content.from_function_call(call_id="safe-call", name="safe_tool", arguments="{}", id="safe-occurrence"), + Content.from_function_call( + call_id="guarded-call", + name="guarded_sink", + arguments={"value": f"[{variable_id}]"}, + id="guarded-occurrence", + ), + ] + + if streaming: + chat_client_base.streaming_responses = [ + [ChatResponseUpdate(role="assistant", contents=calls)], + [ChatResponseUpdate(role="assistant", contents=[Content.from_text("done")])], + ] + first_stream = agent.run("run mixed", stream=True, session=session) + first_updates = [update async for update in first_stream] + first = await first_stream.get_final_response() + generated_updates = [ + (update.role, [content.type for content in update.contents]) + for update in first_updates + if any(content.type in {"function_result", "function_approval_request"} for content in update.contents) + ] + assert generated_updates == [ + ("tool", ["function_result"]), + ("assistant", ["function_approval_request"]), + ] + else: + chat_client_base.run_responses = [ + ChatResponse(messages=Message(role="assistant", contents=calls)), + ChatResponse(messages=Message(role="assistant", contents=["done"])), + ] + first = await agent.run("run mixed", session=session) + + assert safe_calls == 1 + assert guarded_values == [] + assert not any( + message.role == "assistant" and any(content.type == "function_result" for content in message.contents) + for message in first.messages + ) + safe_result_message = next( + message + for message in first.messages + if any(content.type == "function_result" and content.call_id == "safe-call" for content in message.contents) + ) + assert safe_result_message.role == "tool" + request = first.user_input_requests[0] + approval_message = next(message for message in first.messages if request in message.contents) + assert approval_message.role == "assistant" + + if streaming: + resumed_stream = agent.run(request.to_function_approval_response(True), stream=True, session=session) + _ = [update async for update in resumed_stream] + await resumed_stream.get_final_response() + else: + await agent.run(request.to_function_approval_response(True), session=session) + + assert safe_calls == 1 + assert guarded_values == ["hidden payload"] + + @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/core/test_sessions.py b/python/packages/core/tests/core/test_sessions.py index 27028dc7315..193859d3aa8 100644 --- a/python/packages/core/tests/core/test_sessions.py +++ b/python/packages/core/tests/core/test_sessions.py @@ -451,6 +451,60 @@ def test_filter_approval_controls_keeps_response_for_pending_placeholder() -> No assert any(placeholder in message.contents for message in filtered) +def _replacement_approval_round( + *, + call_id: str, + occurrence_id: str, + request_id: str, +) -> tuple[Content, Content]: + function_call = Content.from_function_call( + call_id=call_id, + name="guarded", + arguments="{}", + id=occurrence_id, + ) + request = Content.from_function_approval_request( + id=request_id, + function_call=function_call, + additional_properties={"_replacement_approval_request": True}, + ) + return function_call, request + + +def test_filter_approval_controls_correlates_replacement_by_occurrence() -> None: + """Resolving one reused-call-id replacement must leave its sibling pending.""" + first_call, first_request = _replacement_approval_round( + call_id="reused", + occurrence_id="occurrence-1", + request_id="replacement-1", + ) + second_call, second_request = _replacement_approval_round( + call_id="reused", + occurrence_id="occurrence-2", + request_id="replacement-2", + ) + second_response = Content.from_function_approval_response( + approved=True, + id="occurrence-2", + function_call=second_call, + ) + + filtered = _filter_approval_control_messages([ + Message(role="assistant", contents=[first_call, first_request]), + Message(role="assistant", contents=[second_call, second_request]), + Message(role="user", contents=[second_response]), + Message(role="tool", contents=[Content.from_function_result(call_id="reused", result="done")]), + ]) + + controls = [ + content + for message in filtered + for content in message.contents + if content.type in {"function_approval_request", "function_approval_response"} + ] + assert controls == [first_request] + + class TestHistoryProviderBase: def test_default_flags(self) -> None: provider = ConcreteHistoryProvider("mem") diff --git a/python/packages/core/tests/test_security.py b/python/packages/core/tests/test_security.py index 3f1370dc8ec..903d64e0f2c 100644 --- a/python/packages/core/tests/test_security.py +++ b/python/packages/core/tests/test_security.py @@ -4310,3 +4310,514 @@ async def fake_call(**kwargs): wrapped_once = func_tool.func _wrap_mcp_function_for_ifc(func_tool, IntegrityLabel.UNTRUSTED) assert func_tool.func is wrapped_once + + +async def _hide_session_value( + tracker: LabelTrackingFunctionMiddleware, + policy: PolicyEnforcementFunctionMiddleware, + session: AgentSession, + value: str, +) -> str: + """Hide one untrusted result through the public middleware pipeline.""" + + class SourceArgs(BaseModel): + value: str + + async def source(value: str) -> str: + return value + + source_tool = FunctionTool( + func=source, + name="source", + description="Return untrusted data", + input_model=SourceArgs, + additional_properties={"source_integrity": "untrusted"}, + ) + context = FunctionInvocationContext(function=source_tool, arguments={"value": value}, session=session) + + async def execute(_context: FunctionInvocationContext) -> list[Content]: + return [Content.from_text(value)] + + await FunctionMiddlewarePipeline(tracker, policy).execute(context, execute) + result = cast(list[Content], context.result) + assert result[0].text is not None + return cast(str, json.loads(result[0].text)["variable_id"]) + + +class TestManualSecuritySessionSelection: + """Manual middleware wiring must select explicit session state safely.""" + + @staticmethod + def _sink() -> FunctionTool: + class SinkArgs(BaseModel): + value: str + + async def sink(value: str) -> str: + return value + + return FunctionTool( + func=sink, + name="sink", + description="Forward hidden data", + input_model=SinkArgs, + additional_properties={"accepts_untrusted": True, "source_integrity": "trusted"}, + ) + + async def test_shared_middleware_selects_explicit_sessions_a_b_a(self) -> None: + tracker = LabelTrackingFunctionMiddleware() + policy = PolicyEnforcementFunctionMiddleware() + alice = AgentSession(session_id="alice") + bob = AgentSession(session_id="bob") + alice_variable = await _hide_session_value(tracker, policy, alice, "alice secret") + bob_variable = await _hide_session_value(tracker, policy, bob, "bob secret") + sink = self._sink() + received: list[tuple[str, str]] = [] + + async def forward(session: AgentSession, variable_id: str) -> None: + context = FunctionInvocationContext( + function=sink, + arguments={"value": f"[{variable_id}]"}, + session=session, + ) + + async def execute(current: FunctionInvocationContext) -> list[Content]: + value = cast(dict[str, Any], current.arguments)["value"] + received.append((session.session_id, cast(str, value))) + return [Content.from_text("sent")] + + await FunctionMiddlewarePipeline(tracker, policy).execute(context, execute) + + await forward(bob, alice_variable) + await forward(alice, alice_variable) + await forward(alice, bob_variable) + await forward(bob, bob_variable) + await forward(alice, alice_variable) + + assert received == [ + ("bob", f"[{alice_variable}]"), + ("alice", "alice secret"), + ("alice", f"[{bob_variable}]"), + ("bob", "bob secret"), + ("alice", "alice secret"), + ] + + async def test_overlapping_sessions_keep_task_local_variable_scope(self) -> None: + tracker = LabelTrackingFunctionMiddleware() + policy = PolicyEnforcementFunctionMiddleware() + alice = AgentSession(session_id="alice-overlap") + bob = AgentSession(session_id="bob-overlap") + alice_variable = await _hide_session_value(tracker, policy, alice, "alice secret") + bob_variable = await _hide_session_value(tracker, policy, bob, "bob secret") + inspect_tool = next(tool for tool in tracker.get_security_tools() if tool.name == "inspect_variable") + both_started = asyncio.Event() + started = 0 + + async def inspect(session: AgentSession, variable_id: str) -> str: + nonlocal started + context = FunctionInvocationContext( + function=inspect_tool, + arguments={"variable_id": variable_id, "reason": "overlap isolation"}, + session=session, + ) + + async def execute(current: FunctionInvocationContext) -> list[Content]: + nonlocal started + started += 1 + if started == 2: + both_started.set() + await both_started.wait() + await asyncio.sleep(0) + assert get_current_middleware() is tracker + return await inspect_tool.invoke(arguments=current.arguments, context=current) + + await FunctionMiddlewarePipeline(tracker, policy).execute(context, execute) + result = cast(list[Content], context.result) + assert result[0].text is not None + return cast(str, json.loads(result[0].text)["content"]) + + assert await asyncio.gather( + inspect(alice, alice_variable), + inspect(bob, bob_variable), + ) == ["alice secret", "bob secret"] + assert get_current_middleware() is None + + async def test_direct_standalone_invocation_keeps_private_scope(self) -> None: + tracker = LabelTrackingFunctionMiddleware() + policy = PolicyEnforcementFunctionMiddleware() + variable_id = tracker.get_variable_store().store( + "standalone secret", + ContentLabel(integrity=IntegrityLabel.UNTRUSTED), + ) + sink = self._sink() + received: list[str] = [] + context = FunctionInvocationContext(function=sink, arguments={"value": f"[{variable_id}]"}) + + async def execute(current: FunctionInvocationContext) -> list[Content]: + received.append(cast(str, cast(dict[str, Any], current.arguments)["value"])) + return [Content.from_text("sent")] + + await FunctionMiddlewarePipeline(tracker, policy).execute(context, execute) + + assert received == ["standalone secret"] + assert tracker.get_variable_store().retrieve(variable_id)[0] == "standalone secret" + + +class TestVariableArgumentPolicy: + """Expanded hidden values retain their labels for policy enforcement.""" + + @staticmethod + def _sink( + *, + accepts_untrusted: bool, + max_confidentiality: str = "private", + ) -> FunctionTool: + class SinkArgs(BaseModel): + value: Any + + async def sink(value: Any) -> str: + return str(value) + + return FunctionTool( + func=sink, + name="sink", + description="Forward hidden data", + input_model=SinkArgs, + additional_properties={ + "accepts_untrusted": accepts_untrusted, + "max_allowed_confidentiality": max_confidentiality, + "source_integrity": "trusted", + }, + ) + + @pytest.mark.parametrize( + ("make_value", "expected"), + [ + pytest.param(lambda variable_id: f"[{variable_id}]", "payload", id="bracketed"), + pytest.param(lambda variable_id: f" [ {variable_id} ] ", "payload", id="whitespace"), + pytest.param(lambda variable_id: variable_id, "payload", id="bare"), + pytest.param(lambda variable_id: f"prefix [{variable_id}] suffix", "prefix payload suffix", id="embedded"), + pytest.param(lambda variable_id: f"[{variable_id}]/[{variable_id}]", "payload/payload", id="duplicate"), + pytest.param(lambda variable_id: [f"[{variable_id}]", variable_id], ["payload", "payload"], id="list"), + pytest.param(lambda variable_id: {"item": f"[{variable_id}]"}, {"item": "payload"}, id="mapping"), + pytest.param( + lambda variable_id: {"outer": [{"inner": f"value=[{variable_id}]"}]}, + {"outer": [{"inner": "value=payload"}]}, + id="deep", + ), + ], + ) + async def test_all_reference_forms_resolve_and_block_untrusted_arguments( + self, + make_value: Any, + expected: Any, + ) -> None: + tracker = LabelTrackingFunctionMiddleware() + policy = PolicyEnforcementFunctionMiddleware() + variable_id = tracker.get_variable_store().store( + "payload", + ContentLabel(integrity=IntegrityLabel.UNTRUSTED), + ) + value = make_value(variable_id) + context = FunctionInvocationContext( + function=self._sink(accepts_untrusted=False), + arguments={"value": value}, + ) + + with pytest.raises(MiddlewareTermination): + await FunctionMiddlewarePipeline(tracker, policy).execute( + context, + lambda _context: pytest.fail("Blocked tool must not execute"), + ) + + assert cast(dict[str, Any], context.arguments)["value"] == expected + assert cast(dict[str, Any], context.metadata["original_arguments_for_messages"])["value"] == value + assert context.metadata["argument_label"].integrity == IntegrityLabel.UNTRUSTED + assert policy.get_audit_log()[-1]["type"] == "untrusted_arguments" + + @pytest.mark.parametrize("value", ["[var_0123456789abcdef]", "var_0123456789abcdef"]) + async def test_unknown_references_remain_literal(self, value: str) -> None: + tracker = LabelTrackingFunctionMiddleware() + policy = PolicyEnforcementFunctionMiddleware() + received: list[str] = [] + context = FunctionInvocationContext( + function=self._sink(accepts_untrusted=False), + arguments={"value": value}, + ) + + async def execute(current: FunctionInvocationContext) -> list[Content]: + received.append(cast(str, cast(dict[str, Any], current.arguments)["value"])) + return [Content.from_text("sent")] + + await FunctionMiddlewarePipeline(tracker, policy).execute(context, execute) + + assert received == [value] + argument_label = cast(ContentLabel, context.metadata["argument_label"]) + assert argument_label.integrity == IntegrityLabel.TRUSTED + assert argument_label.confidentiality == ConfidentialityLabel.PUBLIC + + async def test_accepts_untrusted_blind_forwards_without_model_context_taint(self) -> None: + tracker = LabelTrackingFunctionMiddleware() + policy = PolicyEnforcementFunctionMiddleware() + variable_id = tracker.get_variable_store().store( + "payload", + ContentLabel(integrity=IntegrityLabel.UNTRUSTED), + ) + received: list[str] = [] + context = FunctionInvocationContext( + function=self._sink(accepts_untrusted=True), + arguments={"value": f"[{variable_id}]"}, + ) + + async def execute(current: FunctionInvocationContext) -> list[Content]: + received.append(cast(str, cast(dict[str, Any], current.arguments)["value"])) + return [Content.from_text("sent")] + + await FunctionMiddlewarePipeline(tracker, policy).execute(context, execute) + + assert received == ["payload"] + assert context.metadata["argument_label"].integrity == IntegrityLabel.UNTRUSTED + assert tracker.get_context_label().integrity == IntegrityLabel.TRUSTED + + async def test_private_hidden_argument_is_blocked_from_public_sink(self) -> None: + tracker = LabelTrackingFunctionMiddleware() + policy = PolicyEnforcementFunctionMiddleware() + variable_id = tracker.get_variable_store().store( + "private payload", + ContentLabel( + integrity=IntegrityLabel.UNTRUSTED, + confidentiality=ConfidentialityLabel.PRIVATE, + ), + ) + context = FunctionInvocationContext( + function=self._sink(accepts_untrusted=True, max_confidentiality="public"), + arguments={"value": f"[{variable_id}]"}, + ) + + with pytest.raises(MiddlewareTermination): + await FunctionMiddlewarePipeline(tracker, policy).execute( + context, + lambda _context: pytest.fail("Confidential value must not reach a public sink"), + ) + + assert context.metadata["effective_invocation_label"].confidentiality == ConfidentialityLabel.PRIVATE + assert policy.get_audit_log()[-1]["subtype"] == "max_allowed_confidentiality" + + async def test_argument_labels_do_not_rewrite_result_labels(self) -> None: + tracker = LabelTrackingFunctionMiddleware() + policy = PolicyEnforcementFunctionMiddleware() + variable_id = tracker.get_variable_store().store( + "private payload", + ContentLabel( + integrity=IntegrityLabel.UNTRUSTED, + confidentiality=ConfidentialityLabel.PRIVATE, + ), + ) + context = FunctionInvocationContext( + function=self._sink(accepts_untrusted=True), + arguments={"value": f"[{variable_id}]"}, + ) + + async def execute(_context: FunctionInvocationContext) -> list[Content]: + return [Content.from_text("public result")] + + await FunctionMiddlewarePipeline(tracker, policy).execute(context, execute) + + assert context.metadata["result_label"].integrity == IntegrityLabel.TRUSTED + assert context.metadata["result_label"].confidentiality == ConfidentialityLabel.PUBLIC + assert tracker.get_context_label().confidentiality == ConfidentialityLabel.PUBLIC + + async def test_policy_approval_allows_exact_resolved_invocation(self) -> None: + tracker = LabelTrackingFunctionMiddleware() + policy = PolicyEnforcementFunctionMiddleware(approval_on_violation=True) + variable_id = tracker.get_variable_store().store( + "payload", + ContentLabel(integrity=IntegrityLabel.UNTRUSTED), + ) + sink = self._sink(accepts_untrusted=False) + request = FunctionInvocationContext(function=sink, arguments={"value": f"[{variable_id}]"}) + request.metadata.update({"call_id": "call-approved", "function_call_occurrence_id": "occurrence-approved"}) + + with pytest.raises(MiddlewareTermination): + await FunctionMiddlewarePipeline(tracker, policy).execute( + request, + lambda _context: pytest.fail("Tool must wait for approval"), + ) + + approval_request = request.result + assert isinstance(approval_request, Content) + assert approval_request.id == "occurrence-approved" + assert approval_request.function_call is not None + assert approval_request.function_call.call_id == "call-approved" + assert approval_request.function_call.parse_arguments() == {"value": f"[{variable_id}]"} + + replay = FunctionInvocationContext(function=sink, arguments={"value": f"[{variable_id}]"}) + replay.metadata.update({ + "call_id": "call-approved", + "function_call_occurrence_id": "occurrence-approved", + "approval_response": approval_request.to_function_approval_response(True), + }) + received: list[str] = [] + + async def execute(current: FunctionInvocationContext) -> list[Content]: + received.append(cast(str, cast(dict[str, Any], current.arguments)["value"])) + return [Content.from_text("sent")] + + await FunctionMiddlewarePipeline(tracker, policy).execute(replay, execute) + + assert received == ["payload"] + assert replay.metadata["user_approved_violation"] is True + + async def test_changed_resolved_arguments_do_not_match_approval(self) -> None: + tracker = LabelTrackingFunctionMiddleware() + policy = PolicyEnforcementFunctionMiddleware(approval_on_violation=True) + session = AgentSession(session_id="resolved-mismatch") + variable_id = tracker.get_variable_store(session).store( + "original", + ContentLabel(integrity=IntegrityLabel.UNTRUSTED), + ) + sink = self._sink(accepts_untrusted=False) + request = FunctionInvocationContext( + function=sink, + arguments={"value": f"[{variable_id}]"}, + session=session, + ) + request.metadata.update({"call_id": "call-resolved", "function_call_occurrence_id": "occurrence-resolved"}) + + with pytest.raises(MiddlewareTermination): + await FunctionMiddlewarePipeline(tracker, policy).execute( + request, + lambda _context: pytest.fail("Tool must wait for approval"), + ) + approval_request = cast(Content, request.result) + security_state = cast(dict[str, Any], session.state["__agent_framework_fides_security__"]) + variable_state = cast(dict[str, Any], security_state["variables"]) + cast(dict[str, Any], variable_state[variable_id])["content"] = json.dumps("changed") + + replay = FunctionInvocationContext( + function=sink, + arguments={"value": f"[{variable_id}]"}, + session=session, + ) + replay.metadata.update({ + "call_id": "call-resolved", + "function_call_occurrence_id": "occurrence-resolved", + "approval_response": approval_request.to_function_approval_response(True), + }) + + with pytest.raises(MiddlewareTermination): + await FunctionMiddlewarePipeline(tracker, policy).execute( + replay, + lambda _context: pytest.fail("Changed resolved arguments need fresh approval"), + ) + assert isinstance(replay.result, Content) + assert replay.result.type == "function_approval_request" + + async def test_changed_runtime_kwargs_do_not_match_approval(self) -> None: + tracker = LabelTrackingFunctionMiddleware() + policy = PolicyEnforcementFunctionMiddleware(approval_on_violation=True) + variable_id = tracker.get_variable_store().store( + "payload", + ContentLabel(integrity=IntegrityLabel.UNTRUSTED), + ) + sink = self._sink(accepts_untrusted=False) + request = FunctionInvocationContext( + function=sink, + arguments={"value": f"[{variable_id}]"}, + kwargs={"tenant": "one"}, + ) + request.metadata["call_id"] = "call-runtime" + + with pytest.raises(MiddlewareTermination): + await FunctionMiddlewarePipeline(tracker, policy).execute( + request, + lambda _context: pytest.fail("Tool must wait for approval"), + ) + approval_request = cast(Content, request.result) + replay = FunctionInvocationContext( + function=sink, + arguments={"value": f"[{variable_id}]"}, + kwargs={"tenant": "two"}, + ) + replay.metadata.update({ + "call_id": "call-runtime", + "approval_response": approval_request.to_function_approval_response(True), + }) + + with pytest.raises(MiddlewareTermination): + await FunctionMiddlewarePipeline(tracker, policy).execute( + replay, + lambda _context: pytest.fail("Changed runtime kwargs need fresh approval"), + ) + assert isinstance(replay.result, Content) + assert replay.result.type == "function_approval_request" + + async def test_opaque_runtime_kwarg_fails_closed(self) -> None: + tracker = LabelTrackingFunctionMiddleware() + policy = PolicyEnforcementFunctionMiddleware(approval_on_violation=True) + variable_id = tracker.get_variable_store().store( + "payload", + ContentLabel(integrity=IntegrityLabel.UNTRUSTED), + ) + context = FunctionInvocationContext( + function=self._sink(accepts_untrusted=False), + arguments={"value": f"[{variable_id}]"}, + kwargs={"opaque": object()}, + ) + context.metadata["call_id"] = "call-opaque" + + with pytest.raises(MiddlewareTermination): + await FunctionMiddlewarePipeline(tracker, policy).execute( + context, + lambda _context: pytest.fail("Opaque invocation must not execute"), + ) + + assert isinstance(context.result, dict) + assert context.result["violation_type"] == "unsafe_approval_binding" + + async def test_pending_policy_approval_survives_session_restore(self) -> None: + config = SecureAgentConfig(approval_on_violation=True) + session = AgentSession(session_id="restored-approval") + tracker, policy = await _get_session_security_middleware(config, session) + variable_id = config.get_variable_store(session).store( + ["durable", 1], + ContentLabel(integrity=IntegrityLabel.UNTRUSTED), + ) + sink = self._sink(accepts_untrusted=False) + request = FunctionInvocationContext( + function=sink, + arguments={"value": f"[{variable_id}]"}, + session=session, + kwargs={"session": session, "tenant": "one"}, + ) + request.metadata.update({"call_id": "call-restored", "function_call_occurrence_id": "occurrence-restored"}) + + with pytest.raises(MiddlewareTermination): + await FunctionMiddlewarePipeline(tracker, policy).execute( + request, + lambda _context: pytest.fail("Tool must wait for approval"), + ) + approval_request = cast(Content, request.result) + restored = AgentSession.from_dict(json.loads(json.dumps(session.to_dict()))) + restored_tracker, restored_policy = await _get_session_security_middleware(config, restored) + replay = FunctionInvocationContext( + function=sink, + arguments={"value": f"[{variable_id}]"}, + session=restored, + kwargs={"session": restored, "tenant": "one"}, + ) + replay.metadata.update({ + "call_id": "call-restored", + "function_call_occurrence_id": "occurrence-restored", + "approval_response": approval_request.to_function_approval_response(True), + }) + executed = False + + async def execute(_context: FunctionInvocationContext) -> list[Content]: + nonlocal executed + executed = True + return [Content.from_text("sent")] + + await FunctionMiddlewarePipeline(restored_tracker, restored_policy).execute(replay, execute) + + assert executed is True + assert replay.metadata["user_approved_violation"] is True diff --git a/python/packages/ollama/tests/test_ollama_chat_client.py b/python/packages/ollama/tests/test_ollama_chat_client.py index 70f578cc37a..a8245bf3ef0 100644 --- a/python/packages/ollama/tests/test_ollama_chat_client.py +++ b/python/packages/ollama/tests/test_ollama_chat_client.py @@ -810,3 +810,21 @@ def test_format_tool_message_strips_unique_suffix(self) -> None: assert formatted[0].tool_name == "search:advanced", ( f"Expected bare name 'search:advanced', got '{formatted[0].tool_name}'" ) + + def test_mixed_policy_approval_roles_preserve_tool_result(self) -> None: + """Role-separated policy output keeps the safe sibling visible to Ollama.""" + client = OllamaChatClient(host="http://localhost:12345", model="test-model") + function_call = Content.from_function_call(call_id="guarded", name="sink", arguments="{}") + approval_request = Content.from_function_approval_request(id="guarded", function_call=function_call) + messages = [ + Message( + role="tool", + contents=[Content.from_function_result(call_id="safe", result="safe result")], + ), + Message(role="assistant", contents=[approval_request]), + ] + + prepared = client._prepare_messages_for_ollama(messages) + + assert [message.role for message in prepared] == ["tool", "assistant"] + assert prepared[0].content == "safe result" diff --git a/python/samples/02-agents/security/FIDES_DEVELOPER_GUIDE.md b/python/samples/02-agents/security/FIDES_DEVELOPER_GUIDE.md index 02a2c5b0970..e28d6a6cfc0 100644 --- a/python/samples/02-agents/security/FIDES_DEVELOPER_GUIDE.md +++ b/python/samples/02-agents/security/FIDES_DEVELOPER_GUIDE.md @@ -635,12 +635,16 @@ agent = Agent( middleware=[label_tracker, policy_enforcer], ) -# Run agent - security is automatic +# Omitting session uses isolated state for this run. response = await agent.run(messages=[ {"role": "user", "content": "Search the web for Python tutorials"} ]) ``` +Reusable manual middleware uses one private security scope for the complete run when `session` is omitted, so tool +chains within that run share hidden variables while separate runs remain isolated. Pass an explicit `AgentSession` +when labels, variables, audit records, or FIDES policy-approval authority must survive across distinct `Agent.run` calls. Ordinary tool approval responses retain their no-session pass-through behavior. + ### Example 3: Agent Processing Hidden Content When an agent encounters hidden content, it uses `quarantined_llm` with variable IDs: @@ -928,6 +932,12 @@ Configure tool security requirements in the `@tool` decorator: ) ``` +Hidden variable references in tool arguments are expanded recursively, but their stored labels remain attached to +the invocation. A tool with `accepts_untrusted=False` is blocked, audited, or sent for policy approval before it can +receive hidden untrusted data. `accepts_untrusted=True` permits blind forwarding without exposing that data to the +model context. It does not bypass `max_allowed_confidentiality`; hidden private data still cannot flow to a public +sink. Argument labels do not rewrite result labels. + **Approval model:** - Use `approval_mode="always_require"` for normal human-in-the-loop approval on a specific tool. - Use `SecureAgentConfig(..., approval_on_violation=True)` to request approval only when a secure-policy check would otherwise block a call. @@ -961,7 +971,7 @@ Configure tool security requirements in the `@tool` decorator: Access the audit log: ```python -audit_log = policy_enforcer.get_audit_log() +audit_log = policy_enforcer.get_audit_log(session) for violation in audit_log: print(f"Type: {violation['type']}") @@ -984,7 +994,7 @@ Access the middleware's variable store to list or inspect stored variables: ```python # Get all stored variables -variables = label_tracker.list_variables() +variables = label_tracker.list_variables(session) print(f"Stored variables: {variables}") # Get variable metadata