From 862887e395417c145145aa8a070aad75ebb789dc Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Mon, 7 Sep 2026 20:52:43 +0200 Subject: [PATCH 1/3] Python: enforce FIDES labels on expanded variables Bind expanded hidden values to their stored integrity and confidentiality labels, enforce sink policy without tainting model context, and persist exact approval digests across session restore. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../packages/core/agent_framework/security.py | 949 ++++++++---------- python/packages/core/tests/test_security.py | 525 +++++++++- .../security/FIDES_DEVELOPER_GUIDE.md | 20 +- 3 files changed, 965 insertions(+), 529 deletions(-) diff --git a/python/packages/core/agent_framework/security.py b/python/packages/core/agent_framework/security.py index f2ccf1f673..19ea856e5e 100644 --- a/python/packages/core/agent_framework/security.py +++ b/python/packages/core/agent_framework/security.py @@ -16,22 +16,23 @@ 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 from ._feature_stage import ExperimentalFeature, experimental -from ._middleware import FunctionInvocationContext, FunctionMiddleware, MiddlewareTermination +from ._middleware import FunctionInvocationContext, FunctionMiddleware, MiddlewareFailure, MiddlewareTermination from ._serialization import SerializationMixin from ._sessions import AgentSession, ContextProvider from ._tools import FunctionTool, tool @@ -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,95 @@ 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 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 +954,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 +995,45 @@ 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: + if context.session is not None: + scope = self._scope_for_session(context.session) + elif context.tools is not None: + raise MiddlewareFailure( + "Reusable FIDES middleware requires an AgentSession. Pass session=... to Agent.run(), " + "or configure SecureAgentConfig as a context provider." + ) + return self._active_security_scope.set(scope) + + class _ScopedVariableStore(ContentVariableStore): """Variable store backed by one security scope's serializable state.""" @@ -983,7 +1121,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: @@ -1048,64 +1186,45 @@ def __init__( 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 +1281,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 +1494,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 +1526,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 +1756,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 +1782,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 +1874,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: @@ -2022,49 +1950,37 @@ def __init__( 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 +2018,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 +2062,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 +2093,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 +2123,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 +2146,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,7 +2162,7 @@ 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) additional_properties: dict[str, Any] = { "policy_violation": True, "violation_type": primary["violation_type"], @@ -2322,73 +2206,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 +2274,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", }, }) - # Confidentiality policy: block writing higher-confidentiality data to a lower - # confidentiality destination (data exfiltration). - conf_result = self._check_confidentiality_policy_detailed(context, context_label) + 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", + }, + }) + + 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 +2313,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 +2440,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 +2545,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 +2834,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/test_security.py b/python/packages/core/tests/test_security.py index 3f1370dc8e..da01de12a2 100644 --- a/python/packages/core/tests/test_security.py +++ b/python/packages/core/tests/test_security.py @@ -21,7 +21,7 @@ Message, SessionContext, ) -from agent_framework._middleware import FunctionMiddlewarePipeline, MiddlewareTermination +from agent_framework._middleware import FunctionMiddlewarePipeline, MiddlewareFailure, MiddlewareTermination from agent_framework._tools import FunctionTool, _auto_invoke_function, normalize_function_invocation_configuration from agent_framework._types import Content from agent_framework.security import ( @@ -4310,3 +4310,526 @@ 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_manual_agent_loop_without_session_fails_closed(self) -> None: + tracker = LabelTrackingFunctionMiddleware() + sink = self._sink() + context = FunctionInvocationContext( + function=sink, + arguments={"value": "test"}, + tools=[sink], + ) + + with pytest.raises(MiddlewareFailure, match="requires an AgentSession"): + await tracker.process(context, lambda: pytest.fail("Tool must not execute without a session")) + + 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/samples/02-agents/security/FIDES_DEVELOPER_GUIDE.md b/python/samples/02-agents/security/FIDES_DEVELOPER_GUIDE.md index 02a2c5b097..5bdd2b20d6 100644 --- a/python/samples/02-agents/security/FIDES_DEVELOPER_GUIDE.md +++ b/python/samples/02-agents/security/FIDES_DEVELOPER_GUIDE.md @@ -612,6 +612,7 @@ response = await agent.run(messages=[ ### Example 2: Manual Setup (More Control) ```python +from agent_framework import AgentSession from agent_framework.security import ( LabelTrackingFunctionMiddleware, PolicyEnforcementFunctionMiddleware, @@ -635,12 +636,17 @@ agent = Agent( middleware=[label_tracker, policy_enforcer], ) -# Run agent - security is automatic -response = await agent.run(messages=[ +# Manual agent-loop wiring requires an explicit session so security state cannot leak between runs. +session = AgentSession() +response = await agent.run(session=session, messages=[ {"role": "user", "content": "Search the web for Python tutorials"} ]) ``` +Reusable manual middleware selects labels, variables, audit records, and approvals from the explicit +`AgentSession`. Concurrent runs remain task-local. Calling an agent loop without a session fails closed; only +direct standalone `FunctionTool` invocation (where no agent tool list is present) may use the middleware private scope. + ### Example 3: Agent Processing Hidden Content When an agent encounters hidden content, it uses `quarantined_llm` with variable IDs: @@ -928,6 +934,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 +973,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 +996,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 From 504c4c3423bc84f89419578b0a06cb3c0965da4d Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Wed, 9 Sep 2026 09:20:03 +0200 Subject: [PATCH 2/3] Python: address FIDES argument policy review Preserve no-session middleware compatibility with run-local state, support enum approval arguments, recover superseded policy approvals safely, and keep mixed function results in valid tool-role messages. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../specs/004-python-function-calling-loop.md | 6 +- .../core/agent_framework/_sessions.py | 43 +- .../packages/core/agent_framework/_tools.py | 196 ++++-- .../packages/core/agent_framework/_types.py | 1 + .../packages/core/agent_framework/security.py | 34 +- .../core/test_function_invocation_logic.py | 42 ++ .../tests/core/test_harness_tool_approval.py | 620 +++++++++++++++++- .../packages/core/tests/core/test_sessions.py | 54 ++ python/packages/core/tests/test_security.py | 14 +- .../ollama/tests/test_ollama_chat_client.py | 18 + .../security/FIDES_DEVELOPER_GUIDE.md | 12 +- 11 files changed, 961 insertions(+), 79 deletions(-) diff --git a/docs/specs/004-python-function-calling-loop.md b/docs/specs/004-python-function-calling-loop.md index d7d4605818..e96dd9d8ee 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 c11a8ae9c2..da8cfbcb31 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 e76c37d665..6e8756022b 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) @@ -2253,16 +2264,17 @@ def _extract_tools( def _get_tool_approval_state(invocation_session: AgentSession | None) -> 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( @@ -2270,7 +2282,7 @@ def _get_tool_approval_state(invocation_session: AgentSession | None) -> dict[st f"got {type(raw_state).__name__}." ) 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 @@ -2367,7 +2379,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 +2393,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 +2409,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 +2563,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 +2590,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 +2606,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 +2617,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 +2640,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 +2786,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 +2902,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 +3185,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 +3204,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 +3249,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 +3303,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 +4087,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 +4122,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 93a5679c88..7b0efa634e 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 19ea856e5e..1a46bd3687 100644 --- a/python/packages/core/agent_framework/security.py +++ b/python/packages/core/agent_framework/security.py @@ -32,7 +32,7 @@ from pydantic import BaseModel, Field from ._feature_stage import ExperimentalFeature, experimental -from ._middleware import FunctionInvocationContext, FunctionMiddleware, MiddlewareFailure, MiddlewareTermination +from ._middleware import FunctionInvocationContext, FunctionMiddleware, MiddlewareTermination from ._serialization import SerializationMixin from ._sessions import AgentSession, ContextProvider from ._tools import FunctionTool, tool @@ -708,6 +708,14 @@ def _strict_json_value( """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: @@ -1023,14 +1031,8 @@ def _scope_for_session(self, session: AgentSession | None) -> _SecurityScope: def _activate_security_scope(self, context: FunctionInvocationContext) -> Token[_SecurityScope | None]: scope = self._default_security_scope - if not self._security_scope_is_fixed: - if context.session is not None: - scope = self._scope_for_session(context.session) - elif context.tools is not None: - raise MiddlewareFailure( - "Reusable FIDES middleware requires an AgentSession. Pass session=... to Agent.run(), " - "or configure SecureAgentConfig as a context provider." - ) + 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) @@ -1180,6 +1182,8 @@ 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, @@ -1944,6 +1948,8 @@ class PolicyEnforcementFunctionMiddleware(FunctionMiddleware, _SecurityScopeBind ) """ + _requires_session_state = True + def __init__( self, allow_untrusted_tools: set[str] | None = None, @@ -2163,7 +2169,15 @@ def _request_policy_violation_approval( approval_id = self._get_approval_id(context) if approval_id: 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": ( @@ -2179,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, ) 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 18615b7c60..e4d36d8e07 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 ( 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 b637bc751e..032309accd 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 27028dc731..193859d3aa 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 da01de12a2..903d64e0f2 100644 --- a/python/packages/core/tests/test_security.py +++ b/python/packages/core/tests/test_security.py @@ -21,7 +21,7 @@ Message, SessionContext, ) -from agent_framework._middleware import FunctionMiddlewarePipeline, MiddlewareFailure, MiddlewareTermination +from agent_framework._middleware import FunctionMiddlewarePipeline, MiddlewareTermination from agent_framework._tools import FunctionTool, _auto_invoke_function, normalize_function_invocation_configuration from agent_framework._types import Content from agent_framework.security import ( @@ -4441,18 +4441,6 @@ async def execute(current: FunctionInvocationContext) -> list[Content]: ) == ["alice secret", "bob secret"] assert get_current_middleware() is None - async def test_manual_agent_loop_without_session_fails_closed(self) -> None: - tracker = LabelTrackingFunctionMiddleware() - sink = self._sink() - context = FunctionInvocationContext( - function=sink, - arguments={"value": "test"}, - tools=[sink], - ) - - with pytest.raises(MiddlewareFailure, match="requires an AgentSession"): - await tracker.process(context, lambda: pytest.fail("Tool must not execute without a session")) - async def test_direct_standalone_invocation_keeps_private_scope(self) -> None: tracker = LabelTrackingFunctionMiddleware() policy = PolicyEnforcementFunctionMiddleware() diff --git a/python/packages/ollama/tests/test_ollama_chat_client.py b/python/packages/ollama/tests/test_ollama_chat_client.py index 70f578cc37..a8245bf3ef 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 5bdd2b20d6..e28d6a6cfc 100644 --- a/python/samples/02-agents/security/FIDES_DEVELOPER_GUIDE.md +++ b/python/samples/02-agents/security/FIDES_DEVELOPER_GUIDE.md @@ -612,7 +612,6 @@ response = await agent.run(messages=[ ### Example 2: Manual Setup (More Control) ```python -from agent_framework import AgentSession from agent_framework.security import ( LabelTrackingFunctionMiddleware, PolicyEnforcementFunctionMiddleware, @@ -636,16 +635,15 @@ agent = Agent( middleware=[label_tracker, policy_enforcer], ) -# Manual agent-loop wiring requires an explicit session so security state cannot leak between runs. -session = AgentSession() -response = await agent.run(session=session, messages=[ +# 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 selects labels, variables, audit records, and approvals from the explicit -`AgentSession`. Concurrent runs remain task-local. Calling an agent loop without a session fails closed; only -direct standalone `FunctionTool` invocation (where no agent tool list is present) may use the middleware private scope. +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 From 052698ce523342bcdff5def5921e1dfabece09d8 Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Wed, 9 Sep 2026 10:37:00 +0200 Subject: [PATCH 3/3] Python: avoid creating approval state on reads Keep pending-approval lookup side-effect free so runs without approvals do not persist a server-owned approval bucket.\n\nCo-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- python/packages/core/agent_framework/_tools.py | 6 ++++-- .../core/tests/core/test_function_invocation_logic.py | 9 +++++++++ 2 files changed, 13 insertions(+), 2 deletions(-) diff --git a/python/packages/core/agent_framework/_tools.py b/python/packages/core/agent_framework/_tools.py index 6e8756022b..5fa270ab80 100644 --- a/python/packages/core/agent_framework/_tools.py +++ b/python/packages/core/agent_framework/_tools.py @@ -2262,7 +2262,7 @@ 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 not _has_authoritative_approval_session(invocation_session): return None @@ -2281,6 +2281,8 @@ def _get_tool_approval_state(invocation_session: AgentSession | None) -> dict[st 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] = {} authoritative_session.state[_TOOL_APPROVAL_STATE_KEY] = new_state return new_state @@ -2299,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, []) 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 e4d36d8e07..02a0f05b8a 100644 --- a/python/packages/core/tests/core/test_function_invocation_logic.py +++ b/python/packages/core/tests/core/test_function_invocation_logic.py @@ -581,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,