diff --git a/docs/specs/004-python-function-calling-loop.md b/docs/specs/004-python-function-calling-loop.md index 85e9e0c657a..ff1c48f8a40 100644 --- a/docs/specs/004-python-function-calling-loop.md +++ b/docs/specs/004-python-function-calling-loop.md @@ -386,6 +386,14 @@ that manually replay messages own the equivalent rule: do not resend an approval - Approval-time `UserInputRequiredException` and `MiddlewareTermination` return immediately without another model call. +### Middleware termination and policy refusals + +- A normal `MiddlewareTermination` stops the function loop after returning its correlated result. +- A correlated `function_result` explicitly marked with `blocked_violation=True` is a policy refusal, not a pause: + it closes the original call and is sent back to the model so it can explain the refusal or choose another action. +- Policy approval requests remain terminal until the user responds. +- Streaming updates, streaming finalization, and non-streaming output follow the same continuation behavior. + ### Approval control content - `function_approval_request` and `function_approval_response` are control-plane contents, not durable model @@ -524,6 +532,7 @@ that manually replay messages own the equivalent rule: do not resend an approval | Consecutive error cap | Error threshold stops repeated failures, submits collected results, and makes only the required final no-tool model call. | `test_function_invocation_config_max_consecutive_errors`, `test_streaming_function_invocation_config_max_consecutive_errors`, `test_approval_resume_error_limit_forces_final_no_tool_response` | | Unknown call handling | Configured false returns an error result; configured true raises. | `test_function_invocation_config_terminate_on_unknown_calls_false`, `test_function_invocation_config_terminate_on_unknown_calls_true`, streaming equivalents | | Middleware termination | Normal non-approval loop stops without a second model call. | `test_terminate_loop_single_function_call`, `test_terminate_loop_multiple_function_calls_one_terminates`, `test_terminate_loop_streaming_single_function_call` | +| Blocked policy result | A correlated result marked `blocked_violation=True` closes the call and continues to a second model turn in streaming and non-streaming modes; the blocked tool does not execute. | `test_blocked_policy_result_continues_function_loop` | | Middleware failure (fatal) | `MiddlewareFailure` from function middleware or a tool body propagates to the caller without becoming a tool result; the tool does not execute (pre-invocation) or its result never feeds another model call (post-invocation); the cause chain is preserved; ordinary exceptions still become tool-error results and the loop continues. | `packages/core/tests/core/test_middleware_with_agent.py::TestMiddlewareFailure::test_failure_before_tool_aborts_run`, `test_failure_after_tool_aborts_run_before_next_model_turn`, `test_failure_cause_chain_reaches_caller`, `test_failure_from_tool_escapes_without_middleware`, `test_failure_streaming_reaches_stream_consumer`, `test_ordinary_exception_still_becomes_tool_error` | | Middleware failure batch cancellation | A fatal signal fails the whole parallel batch: in-flight sibling tool invocations are cancelled and awaited before the failure propagates. Cancellation is cooperative — an async sibling stops at its next suspension point; a synchronous tool body already executing in a worker thread cannot be interrupted and may complete its side effects, but its result is discarded and never reaches the transcript, the model, or history, and failure propagation is not delayed behind it. | `TestMiddlewareFailure::test_failure_cancels_concurrent_sibling_tool`, `test_failure_with_sync_sibling_discards_late_result` | | Middleware failure on a service-managed conversation | The continuation state is already persisted when the batch fails, so before propagating, the loop settles the hosted thread: one error `function_result` per dangling call, sent with `tool_choice="none"` in one extra request; the persisted continuation advances to the settlement response (required for response-ID continuations, a no-op for conversation-object ids) and the settlement response is otherwise discarded; a settlement failure never masks the abort. Without a service-managed conversation no extra request is made. | `TestMiddlewareFailure::test_failure_settles_dangling_calls_on_service_conversation`, `test_failure_settles_service_conversation_streaming`, `test_failure_settlement_advances_response_id_continuation`, `test_failure_without_service_conversation_makes_no_settlement_request` | diff --git a/python/packages/core/agent_framework/_harness/_tool_approval.py b/python/packages/core/agent_framework/_harness/_tool_approval.py index 5c3bf7a2c92..e1c723abeb9 100644 --- a/python/packages/core/agent_framework/_harness/_tool_approval.py +++ b/python/packages/core/agent_framework/_harness/_tool_approval.py @@ -295,9 +295,24 @@ def _function_call_from_request(request: Content) -> Content | None: function_call = request.function_call if function_call is None or function_call.type != "function_call" or function_call.name is None: return None + request_props = request.additional_properties or {} + if request_props: + function_call = copy.copy(function_call) + function_call.additional_properties = { + **(function_call.additional_properties or {}), + **request_props, + } return function_call +def _has_policy_violation(request: Content) -> bool: + """Return whether an approval request represents a FIDES policy violation.""" + properties = request.additional_properties or {} + return bool( + properties.get("policy_violation") or properties.get("blocked_violation") or properties.get("_fides_violations") + ) + + def _arguments_match(rule_arguments: Mapping[str, str], function_call: Content) -> bool: call_arguments = _serialize_arguments(function_call) or {} if len(rule_arguments) != len(call_arguments): @@ -579,7 +594,9 @@ def _inject_collected_responses(self, messages: Sequence[Message], state: ToolAp async def _drain_auto_approvable_queue(self, state: ToolApprovalState) -> None: remaining: list[Content] = [] for request in state.queued_approval_requests: - if _matches_rule(request, state.rules) or await self._matches_auto_rule(request): + if not _has_policy_violation(request) and ( + _matches_rule(request, state.rules) or await self._matches_auto_rule(request) + ): state.collected_approval_responses.append(request.to_function_approval_response(approved=True)) continue remaining.append(request) @@ -603,7 +620,9 @@ async def _process_outbound_messages(self, messages: list[Message], state: ToolA auto_approved: set[int] = set() unresolved: list[Content] = [] for request in approval_requests: - if _matches_rule(request, state.rules) or await self._matches_auto_rule(request): + if not _has_policy_violation(request) and ( + _matches_rule(request, state.rules) or await self._matches_auto_rule(request) + ): state.collected_approval_responses.append(request.to_function_approval_response(approved=True)) auto_approved.add(id(request)) else: diff --git a/python/packages/core/agent_framework/_middleware.py b/python/packages/core/agent_framework/_middleware.py index 45a3dcbf823..a0c8ecb0cec 100644 --- a/python/packages/core/agent_framework/_middleware.py +++ b/python/packages/core/agent_framework/_middleware.py @@ -76,10 +76,18 @@ class MiddlewareTermination(MiddlewareException): """Control-flow exception to terminate middleware execution early.""" result: Any = None # Optional result to return when terminating + blocked_policy: bool = False # Whether this termination represents a FIDES policy block - def __init__(self, message: str = "Middleware terminated execution.", *, result: Any = None) -> None: + def __init__( + self, + message: str = "Middleware terminated execution.", + *, + result: Any = None, + blocked_policy: bool = False, + ) -> None: super().__init__(message, log_level=None) self.result = result + self.blocked_policy = blocked_policy class MiddlewareFailure(MiddlewareException): diff --git a/python/packages/core/agent_framework/_serialization.py b/python/packages/core/agent_framework/_serialization.py index cb7f3a0db89..54df0d25cba 100644 --- a/python/packages/core/agent_framework/_serialization.py +++ b/python/packages/core/agent_framework/_serialization.py @@ -10,6 +10,7 @@ from collections.abc import Mapping, MutableMapping from dataclasses import asdict, is_dataclass from datetime import date, datetime +from enum import Enum from functools import lru_cache from typing import Any, ClassVar, Protocol, TypeGuard, TypeVar, cast, runtime_checkable @@ -650,8 +651,8 @@ def make_json_safe(obj: Any) -> Any: """Recursively convert an object to a JSON-serializable form. Handles dataclasses, Pydantic models, objects with ``to_dict``/``dict``/``__dict__``, - datetimes, bytes (base64), lists, dicts, and primitives. Falls back to ``str()`` for - any remaining non-serializable value so that ``json.dumps`` never raises a + datetimes, bytes/bytearray (base64), enums, sets, lists, dicts, and primitives. Falls back to + ``str()`` for any remaining non-serializable value so that ``json.dumps`` never raises a ``TypeError``. Args: @@ -662,6 +663,8 @@ def make_json_safe(obj: Any) -> Any: """ if isinstance(obj, _JSON_SCALAR_TYPES): return obj + if isinstance(obj, Enum): + return make_json_safe(obj.value) if isinstance(obj, (datetime, date)): return obj.isoformat() if isinstance(obj, (bytes, bytearray)): @@ -691,6 +694,8 @@ def make_json_safe(obj: Any) -> Any: return {str(key): make_json_safe(value) for key, value in obj.items()} # type: ignore[misc] if isinstance(obj, (list, tuple)): return [make_json_safe(item) for item in obj] # type: ignore[misc] + if isinstance(obj, (set, frozenset)): + return [make_json_safe(item) for item in obj] # type: ignore[misc] if hasattr(obj, "__dict__"): return {key: make_json_safe(value) for key, value in vars(obj).items()} return str(obj) diff --git a/python/packages/core/agent_framework/_tools.py b/python/packages/core/agent_framework/_tools.py index 1e4089808ad..f0e057a7c8e 100644 --- a/python/packages/core/agent_framework/_tools.py +++ b/python/packages/core/agent_framework/_tools.py @@ -1616,19 +1616,28 @@ async def final_function_handler(context_obj: Any) -> Any: return Content.from_function_result(call_id=call_id, result=function_result) except MiddlewareTermination as term_exc: # Re-raise to signal loop termination, but first capture any result set by middleware - if middleware_context.result is not None: + middleware_result = middleware_context.result + if middleware_result is not None: + blocked_result = cast(dict[str, Any], middleware_result) if isinstance(middleware_result, dict) else None + if blocked_result is not None and blocked_result.get("blocked_violation") is True: + blocked_properties = dict(function_call_content.additional_properties or {}) + blocked_properties.update({key: value for key, value in blocked_result.items() if key != "error"}) + blocked_error = str(blocked_result.get("error", "Tool blocked by security policy.")) + term_exc.result = Content.from_function_result( + call_id=call_id, + result=blocked_error, + exception=blocked_error, + additional_properties=blocked_properties, + ) # Pass through function_approval_request directly (e.g., from security policy middleware) # so the approval flow in _handle_function_call_results activates correctly. - if ( - isinstance(middleware_context.result, Content) - and middleware_context.result.type == "function_approval_request" - ): - term_exc.result = middleware_context.result + elif isinstance(middleware_result, Content) and middleware_result.type == "function_approval_request": + term_exc.result = middleware_result else: # Store result in exception for caller to extract term_exc.result = Content.from_function_result( call_id=call_id, - result=middleware_context.result, + result=middleware_result, additional_properties=function_call_content.additional_properties, ) raise @@ -1695,7 +1704,11 @@ async def _execute_single_function_call( return [result], False except MiddlewareTermination as exc: if isinstance(exc.result, Content): - return [exc.result], True + # A blocked FIDES call is a normal tool result: the model must receive + # the refusal and get a chance to explain or choose another action. + # Approval requests remain terminal and pause for user input. + is_blocked_policy = exc.blocked_policy + return [exc.result], not is_blocked_policy source_function_call = _underlying_function_call(function_call) return [ Content.from_function_result( diff --git a/python/packages/core/agent_framework/security.py b/python/packages/core/agent_framework/security.py index 397f9af62fc..892becdda33 100644 --- a/python/packages/core/agent_framework/security.py +++ b/python/packages/core/agent_framework/security.py @@ -15,23 +15,26 @@ from __future__ import annotations import asyncio +import base64 +import binascii import contextlib import json import logging import re -import threading import uuid +import weakref from collections.abc import Awaitable, Callable, MutableMapping +from contextvars import ContextVar from copy import 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, cast from pydantic import BaseModel, Field from ._feature_stage import ExperimentalFeature, experimental from ._middleware import FunctionInvocationContext, FunctionMiddleware, MiddlewareTermination -from ._serialization import SerializationMixin +from ._serialization import SerializationMixin, make_json_safe from ._sessions import ContextProvider from ._tools import FunctionTool, tool from ._types import Content, Message @@ -68,6 +71,10 @@ logger = logging.getLogger(__name__) _BRACKETED_VAR_REF_RE = re.compile(r"^\[\s*(var_[0-9a-fA-F]+)\s*\]$") +_FIDES_TURN_NUMBER_KEY = "_fides_turn_number" +_FIDES_CALL_INDEX_KEY = "_fides_call_index" +_FIDES_VARIABLE_VALUE_TYPE_KEY = "__fides_variable_value_type__" +_FIDES_VARIABLE_VALUE_KEY = "value" # Tools that consume variable IDs literally (as opaque references) and therefore # must NOT have ``var_xxx`` arguments expanded to stored content before execution. @@ -83,6 +90,22 @@ def _get_additional_properties(obj: Any) -> dict[str, Any]: return cast(dict[str, Any], props) if isinstance(props, dict) else {} +def _fides_session_state(session: Any) -> dict[str, Any]: + """Return the initialized ``session.state["_fides"]`` durable-state mapping. + + This is the single owner of the FIDES session-state contract. Callers that may not + have a session must guard for ``None`` before calling. + """ + state = cast(dict[str, Any], session.state.setdefault("_fides", {})) + state.setdefault("context_label", None) + state.setdefault("audit_log", []) + state.setdefault("pending_policy_approvals", {}) + state.setdefault("turn_counter", 0) + state.setdefault("variables", {}) + state.setdefault("variable_metadata", {}) + return state + + # ============================================================================= # Core Security Primitives # ============================================================================= @@ -189,7 +212,7 @@ def to_dict(self, *, exclude: set[str] | None = None, exclude_none: bool = True) "confidentiality": str(self.confidentiality), } if self.metadata: - result["metadata"] = self.metadata + result["metadata"] = make_json_safe(self.metadata) return result @classmethod @@ -319,6 +342,52 @@ def send_message(destination: str, message: str, context_label: ContentLabel): return conf_hierarchy[context_label.confidentiality] <= conf_hierarchy[max_allowed] +def _serialize_variable_content(value: Any) -> Any: + """Convert variable content to a durable JSON-compatible representation.""" + if value is None or type(value) in (str, int, float, bool): + return value + if isinstance(value, bytes): + return { + _FIDES_VARIABLE_VALUE_TYPE_KEY: "bytes", + _FIDES_VARIABLE_VALUE_KEY: base64.b64encode(value).decode("ascii"), + } + if isinstance(value, bytearray): + return { + _FIDES_VARIABLE_VALUE_TYPE_KEY: "bytearray", + _FIDES_VARIABLE_VALUE_KEY: base64.b64encode(bytes(value)).decode("ascii"), + } + if isinstance(value, dict): + return {str(key): _serialize_variable_content(item) for key, item in value.items()} + if isinstance(value, (list, tuple, set, frozenset)): + return [_serialize_variable_content(item) for item in value] + + # Preserve the established best-effort behavior for arbitrary values while + # making nested results JSON-compatible for durable session storage. + safe_value = make_json_safe(value) + if safe_value is value: + return safe_value + return _serialize_variable_content(safe_value) + + +def _deserialize_variable_content(value: Any) -> Any: + """Restore reversible binary values from durable variable content.""" + if isinstance(value, dict): + if set(value) == {_FIDES_VARIABLE_VALUE_TYPE_KEY, _FIDES_VARIABLE_VALUE_KEY}: + encoded = value.get(_FIDES_VARIABLE_VALUE_KEY) + value_type = value.get(_FIDES_VARIABLE_VALUE_TYPE_KEY) + if isinstance(encoded, str) and value_type in ("bytes", "bytearray"): + try: + decoded = base64.b64decode(encoded, validate=True) + except (binascii.Error, TypeError, ValueError): + pass + else: + return bytearray(decoded) if value_type == "bytearray" else decoded + return {key: _deserialize_variable_content(item) for key, item in value.items()} + if isinstance(value, list): + return [_deserialize_variable_content(item) for item in value] + return value + + @experimental(feature_id=ExperimentalFeature.FIDES) class ContentVariableStore: """Client-side storage for untrusted content using variable indirection. @@ -342,9 +411,10 @@ class ContentVariableStore: print(content) # "potentially malicious content" """ - def __init__(self) -> None: + def __init__(self, *, _storage: dict[str, Any] | None = None, _durable: bool = False) -> None: """Initialize an empty ContentVariableStore.""" - self._storage: dict[str, tuple[Any, ContentLabel]] = {} + self._storage = _storage if _storage is not None else {} + self._durable = _durable def store(self, content: Any, label: ContentLabel) -> str: """Store content and return a variable ID. @@ -357,7 +427,8 @@ def store(self, content: Any, label: ContentLabel) -> str: A unique variable ID string. """ var_id = f"var_{uuid.uuid4().hex[:16]}" - self._storage[var_id] = (content, label) + stored_content = _serialize_variable_content(content) if self._durable else content + self._storage[var_id] = {"content": stored_content, "label": label.to_dict()} logger.info(f"Stored content in variable {var_id} with label {label}") return var_id @@ -376,7 +447,23 @@ def retrieve(self, var_id: str) -> tuple[Any, ContentLabel]: if var_id not in self._storage: raise KeyError(f"Variable {var_id} not found in store") - content, label = self._storage[var_id] + record = self._storage[var_id] + legacy_record = cast(list[Any] | tuple[Any, ...], record) if isinstance(record, (list, tuple)) else None + if legacy_record is not None and len(legacy_record) == 2: + # Compatibility with stores created by an older in-memory implementation. + content: Any = legacy_record[0] + label_data: Any = legacy_record[1] + elif isinstance(record, dict): + record_map = cast(dict[str, Any], record) + content = _deserialize_variable_content(record_map["content"]) + label_data = record_map["label"] + else: + raise TypeError(f"Variable {var_id} has an invalid stored representation") + label = ( + label_data + if isinstance(label_data, ContentLabel) + else ContentLabel.from_dict(cast(MutableMapping[str, Any], label_data)) + ) logger.info(f"Retrieved content from variable {var_id} with label {label}") return content, label @@ -681,8 +768,40 @@ def from_message(cls, message: dict[str, Any], index: int | None = None) -> Labe # Security Middleware # ============================================================================= -# Thread-local storage for current middleware instance -_current_middleware = threading.local() + +# Async-safe storage for the active FIDES middleware and session. +_current_middleware: ContextVar[Any] = ContextVar("_current_middleware", default=None) +_current_middleware_session: ContextVar[Any] = ContextVar("_current_middleware_session", default=None) +# Remember one strong session per config/middleware owner in each async context. +# Weak owner keys avoid retaining discarded owners, while copy-on-write updates +# prevent child tasks from mutating a mapping inherited from their parent context. +_fides_sessions_by_owner: ContextVar[weakref.WeakKeyDictionary[Any, Any] | None] = ContextVar( + "_fides_sessions_by_owner", default=None +) + + +def _remember_session(owner: Any, session: Any) -> None: + """Remember a session for one config or middleware in this async context.""" + current = _fides_sessions_by_owner.get(None) + sessions = weakref.WeakKeyDictionary[Any, Any]() + if current is not None: + sessions.update(current) + if session is None: + sessions.pop(owner, None) + else: + sessions[owner] = session + _fides_sessions_by_owner.set(sessions) + + +def _remembered_session(owner: Any) -> Any: + """Return the strongly remembered session for one owner in this async context.""" + sessions = _fides_sessions_by_owner.get(None) + return sessions.get(owner) if sessions is not None else None + + +def _resolve_owner_session(owner: Any, session: Any = None) -> Any: + """Prefer an explicit session, otherwise resolve the session remembered for ``owner``.""" + return session if session is not None else _remembered_session(owner) @experimental(feature_id=ExperimentalFeature.FIDES) @@ -751,6 +870,7 @@ def __init__( default_confidentiality: ConfidentialityLabel = ConfidentialityLabel.PUBLIC, auto_hide_untrusted: bool = True, hide_threshold: IntegrityLabel = IntegrityLabel.UNTRUSTED, + tool_labels: dict[str, ContentLabel] | None = None, ) -> None: """Initialize LabelTrackingFunctionMiddleware. @@ -760,27 +880,70 @@ def __init__( 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. + tool_labels: Labels for tools constructed by a harness or MCP client. """ self.default_integrity = default_integrity self.default_confidentiality = default_confidentiality self.auto_hide_untrusted = auto_hide_untrusted self.hide_threshold = hide_threshold - - # Context-level security label that tracks the cumulative security state - # Starts as TRUSTED + PUBLIC and gets updated based on content added to context - self._context_label = ContentLabel( + self._tool_labels = tool_labels if tool_labels is not None else {} + self._fallback_context_label = ContentLabel( integrity=IntegrityLabel.TRUSTED, confidentiality=ConfidentialityLabel.PUBLIC, - metadata={"initialized": True}, + metadata={"fallback": True}, ) - # Stateful variable store for this middleware instance + # Context labels and variable stores are scoped to the active AgentSession. + # The fallback stores are retained only for calls made without a session. self._variable_store = ContentVariableStore() - - # Metadata about stored variables self._variable_metadata: dict[str, dict[str, Any]] = {} - def get_context_label(self) -> ContentLabel: + @property + def _context_label(self) -> ContentLabel: + """Compatibility view for direct, no-session middleware callers.""" + return self._fallback_context_label + + @_context_label.setter + def _context_label(self, value: ContentLabel) -> None: + self._fallback_context_label = value + + def _get_context_label(self, session: Any = None) -> ContentLabel: + """Read the cumulative context label from session state.""" + if session is None: + return self._fallback_context_label + fides_state = _fides_session_state(session) + data = fides_state["context_label"] + if data is None: + label = ContentLabel( + integrity=IntegrityLabel.TRUSTED, + confidentiality=ConfidentialityLabel.PUBLIC, + metadata={"initialized": True}, + ) + fides_state["context_label"] = label.to_dict() + return label + return ContentLabel.from_dict(data) + + def _set_context_label(self, session: Any, label: ContentLabel) -> None: + """Write the cumulative context label to session state.""" + if session is not None: + _fides_session_state(session)["context_label"] = label.to_dict() + else: + self._fallback_context_label = label + + def _resolve_store(self, session: Any = None) -> ContentVariableStore: + """Resolve the variable store for a session, with a no-session fallback.""" + if session is None: + return self._variable_store + storage = cast(dict[str, Any], _fides_session_state(session)["variables"]) + return ContentVariableStore(_storage=storage, _durable=True) + + def _resolve_metadata(self, session: Any = None) -> dict[str, dict[str, Any]]: + """Resolve variable metadata for a session, with a no-session fallback.""" + if session is None: + return self._variable_metadata + return cast(dict[str, dict[str, Any]], _fides_session_state(session)["variable_metadata"]) + + def get_context_label(self, session: Any = None) -> ContentLabel: """Get the current context-level security label. The context label represents the cumulative security state of the conversation. @@ -790,19 +953,23 @@ def get_context_label(self) -> ContentLabel: Returns: The current context security label. """ - return self._context_label + return self._get_context_label(self._resolve_public_session(session)) + + def _resolve_public_session(self, session: Any = None) -> Any: + return _resolve_owner_session(self, session) - def reset_context_label(self) -> None: + def reset_context_label(self, session: Any = None) -> None: """Reset the context label to initial state (TRUSTED + PUBLIC). Call this when starting a new conversation or session. """ - self._context_label = ContentLabel( + label = ContentLabel( integrity=IntegrityLabel.TRUSTED, confidentiality=ConfidentialityLabel.PUBLIC, metadata={"reset": True} ) + self._set_context_label(self._resolve_public_session(session), label) logger.info("Context label reset to TRUSTED + PUBLIC") - def _update_context_label(self, new_content_label: ContentLabel) -> None: + def _update_context_label(self, new_content_label: ContentLabel, session: Any = None) -> None: """Update the context label based on new content added to the context. The context label is updated using the most restrictive policy: @@ -811,21 +978,24 @@ def _update_context_label(self, new_content_label: ContentLabel) -> None: Args: new_content_label: The label of the new content being added to context. + session: Optional session whose context label should be updated. """ - old_label = self._context_label - self._context_label = combine_labels(self._context_label, new_content_label) + session = self._resolve_public_session(session) + old_label = self._get_context_label(session) + merged_label = combine_labels(old_label, new_content_label) + self._set_context_label(session, merged_label) - if old_label != self._context_label: + if old_label != merged_label: logger.info( f">>> CONTEXT TAINT: [{old_label.integrity.value}, {old_label.confidentiality.value}] " - f"-> [{self._context_label.integrity.value}, {self._context_label.confidentiality.value}] " + f"-> [{merged_label.integrity.value}, {merged_label.confidentiality.value}] " f"(new content: [{new_content_label.integrity.value}, {new_content_label.confidentiality.value}])" ) else: logger.debug( "Context label unchanged: [%s, %s]", - self._context_label.integrity.value, - self._context_label.confidentiality.value, + merged_label.integrity.value, + merged_label.confidentiality.value, ) @staticmethod @@ -896,7 +1066,7 @@ def _expand_variable_reference(self, value: Any) -> Any: if whole_bracketed is not None: variable_id = whole_bracketed.group(1) try: - expanded_content, _ = self._variable_store.retrieve(variable_id) + expanded_content, _ = self._resolve_store(get_current_session()).retrieve(variable_id) extracted = self._extract_primary_tool_content(expanded_content) if extracted is not expanded_content: logger.info( @@ -918,7 +1088,7 @@ def _expand_variable_reference(self, value: Any) -> Any: 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) + expanded_content, _ = self._resolve_store(get_current_session()).retrieve(variable_id) extracted = self._extract_primary_tool_content(expanded_content) logger.warning( f"Expanded BARE (non-bracketed) variable reference '{value.strip()}' " @@ -937,7 +1107,7 @@ def _expand_variable_reference(self, value: Any) -> Any: def replace_bracketed(match_obj: Any) -> str: variable_id = match_obj.group(1) try: - expanded_content, _ = self._variable_store.retrieve(variable_id) + expanded_content, _ = self._resolve_store(get_current_session()).retrieve(variable_id) extracted = self._extract_primary_tool_content(expanded_content) if extracted is not expanded_content: logger.info( @@ -956,7 +1126,7 @@ def replace_bracketed(match_obj: Any) -> str: def replace_bare(match_obj: Any) -> str: variable_id = match_obj.group(0) try: - expanded_content, _ = self._variable_store.retrieve(variable_id) + expanded_content, _ = self._resolve_store(get_current_session()).retrieve(variable_id) extracted = self._extract_primary_tool_content(expanded_content) logger.warning( f"Expanded embedded BARE (non-bracketed) variable reference '{variable_id}' " @@ -1144,7 +1314,13 @@ def _ensure_content_list(result: Any) -> list[Content]: text = str(cast(object, result)) return [Content.from_text(text)] - def _should_hide(self, label: ContentLabel, function_name: str | None = None) -> bool: + def _should_hide( + self, + label: ContentLabel, + function_name: str | None = None, + *, + session: Any = None, + ) -> bool: """Decide whether a Content item with *label* should be hidden. An item is hidden when **all four** conditions hold: @@ -1158,7 +1334,7 @@ def _should_hide(self, label: ContentLabel, function_name: str | None = None) -> return ( self.auto_hide_untrusted and label.integrity == self.hide_threshold - and self._context_label.integrity == IntegrityLabel.TRUSTED + and self._get_context_label(self._resolve_public_session(session)).integrity == IntegrityLabel.TRUSTED and function_name != "inspect_variable" ) @@ -1204,10 +1380,13 @@ async def process( call_next: Callback to continue to next middleware or function execution. """ # Set thread-local middleware reference for tools to access - _current_middleware.instance = self + middleware_token = _current_middleware.set(self) + _remember_session(self, context.session) + session_token = _current_middleware_session.set(context.session) try: function_name = context.function.name + configured_label = self._tool_labels.get(function_name) # ========== Tiered Label Propagation ========== # Step 1: Extract labels from input arguments @@ -1215,9 +1394,13 @@ async def process( # Step 2: Get tool's source_integrity declaration (may be None) declared_source_integrity = self._get_source_integrity(context) + if declared_source_integrity is None and configured_label is not None: + declared_source_integrity = configured_label.integrity # Get confidentiality from function additional_properties or use default confidentiality = self._get_function_confidentiality(context) + if configured_label is not None and "confidentiality" not in _get_additional_properties(context.function): + confidentiality = configured_label.confidentiality # Step 3: Build tiered fallback_label # This label is used for result items that have NO embedded labels. @@ -1250,7 +1433,7 @@ async def process( # context_label: cumulative conversation security state (cross-call). # Used by PolicyEnforcementFunctionMiddleware to validate tool calls. - context.metadata["context_label"] = self._context_label + context.metadata["context_label"] = self._get_context_label(context.session) logger.info( f"Tool call '{function_name}' fallback label (tiered): " @@ -1259,8 +1442,8 @@ async def process( 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}" + f"Current context label: {self._get_context_label(context.session).integrity.value}, " + f"{self._get_context_label(context.session).confidentiality.value}" ) # Store original unexpanded arguments for message reconstruction before expanding @@ -1285,7 +1468,8 @@ async def process( self._label_result(context, function_name, fallback_label) finally: # Clear thread-local reference - _current_middleware.instance = None + _current_middleware.reset(middleware_token) + _current_middleware_session.reset(session_token) def _label_result( self, @@ -1320,6 +1504,7 @@ def _label_result( original_items, function_name, fallback_label=fallback_label, + session=context.session, ) context.result = processed @@ -1329,13 +1514,14 @@ def _label_result( # may affect integrity taint. Confidentiality still reflects the most # restrictive label across the entire tool result, including hidden items. if visible_result_label is None: - if result_label.confidentiality != self._context_label.confidentiality: - old_conf = self._context_label.confidentiality + context_label = self._get_context_label(context.session) + if result_label.confidentiality != context_label.confidentiality: + old_conf = context_label.confidentiality hidden_label = ContentLabel( - integrity=self._context_label.integrity, + integrity=context_label.integrity, confidentiality=result_label.confidentiality, ) - self._update_context_label(hidden_label) + self._update_context_label(hidden_label, context.session) logger.info( f"Result from '{function_name}' hidden (integrity clean) but " f"confidentiality updated: {old_conf.value} -> " @@ -1344,8 +1530,8 @@ def _label_result( else: logger.info( f"Result from '{function_name}' fully hidden - context label " - f"unchanged: {self._context_label.integrity.value}, " - f"{self._context_label.confidentiality.value}" + f"unchanged: {context_label.integrity.value}, " + f"{context_label.confidentiality.value}" ) else: # Only visible content can taint integrity; hidden content still @@ -1354,11 +1540,12 @@ def _label_result( integrity=visible_result_label.integrity, confidentiality=result_label.confidentiality, ) - self._update_context_label(exposed_label) + self._update_context_label(exposed_label, context.session) + context_label = self._get_context_label(context.session) logger.info( f"Context label after processing '{function_name}': " - f"{self._context_label.integrity.value}, " - f"{self._context_label.confidentiality.value}" + f"{context_label.integrity.value}, " + f"{context_label.confidentiality.value}" ) def _get_function_confidentiality(self, context: FunctionInvocationContext) -> ConfidentialityLabel: @@ -1390,6 +1577,7 @@ def _process_result_with_embedded_labels( items: list[Content], function_name: str, fallback_label: ContentLabel, + session: Any = None, ) -> tuple[list[Content], ContentLabel, ContentLabel | None]: """Process Content items, respecting per-item embedded labels. @@ -1410,6 +1598,7 @@ def _process_result_with_embedded_labels( ``_ensure_content_list``). function_name: Name of the function that produced the result. fallback_label: Label to use when an item has no embedded label. + session: Optional session owning the invocation and its variable store. Returns: Tuple of (processed_content_list, combined_label, visible_combined_label). @@ -1427,8 +1616,8 @@ def _process_result_with_embedded_labels( item_label = self._extract_content_label(item, fallback_label) item_labels.append(item_label) - if self._should_hide(item_label, function_name): - hidden = self._hide_item(item, item_label, function_name) + if self._should_hide(item_label, function_name, session=session): + hidden = self._hide_item(item, item_label, function_name, session=session) processed.append(hidden) else: # Attach this item's own label (preserves per-item granularity) @@ -1476,6 +1665,8 @@ def _hide_item( item: Content, label: ContentLabel, function_name: str, + *, + session: Any = None, ) -> Content: """Replace an untrusted Content item with a variable-reference placeholder. @@ -1487,6 +1678,7 @@ def _hide_item( item: The original Content item to hide. label: The security label for the item. function_name: Name of the function that produced the item. + session: Optional session owning the invocation and its variable store. Returns: A Content item containing the variable reference. @@ -1496,10 +1688,13 @@ 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) + session = self._resolve_public_session(session) + store = self._resolve_store(session) + metadata = self._resolve_metadata(session) + var_id = store.store(stored_value, label) # Store metadata about this variable - self._variable_metadata[var_id] = { + metadata[var_id] = { "function_name": function_name, "original_type": item.type, "timestamp": datetime.now().isoformat(), @@ -1521,32 +1716,33 @@ def _hide_item( additional_properties={"_variable_reference": True, "security_label": label.to_dict()}, ) - def get_variable_store(self) -> ContentVariableStore: + def get_variable_store(self, session: Any = None) -> ContentVariableStore: """Get the variable store for this middleware instance. Returns: The ContentVariableStore instance. """ - return self._variable_store + return self._resolve_store(self._resolve_public_session(session)) - def get_variable_metadata(self, var_id: str) -> dict[str, Any] | None: + def get_variable_metadata(self, var_id: str, session: Any = None) -> dict[str, Any] | None: """Get metadata for a stored variable. Args: var_id: The variable ID. + session: Optional session containing the variable. Returns: Metadata dictionary or None if not found. """ - return self._variable_metadata.get(var_id) + return self._resolve_metadata(self._resolve_public_session(session)).get(var_id) - def list_variables(self) -> list[str]: + def list_variables(self, session: Any = None) -> list[str]: """Get a list of all stored variable IDs. Returns: List of variable ID strings. """ - return self._variable_store.list_variables() + return self._resolve_store(self._resolve_public_session(session)).list_variables() def get_security_tools(self) -> list[FunctionTool]: """Get the list of security tools for agent integration. @@ -1599,7 +1795,7 @@ def _set_as_current(self) -> None: This is primarily for testing and debugging purposes. In normal operation, the middleware is automatically set during process(). """ - _current_middleware.instance = self + _current_middleware.set(self) def _clear_current(self) -> None: """Clear the current thread-local middleware instance. @@ -1607,36 +1803,53 @@ def _clear_current(self) -> None: This is primarily for testing and debugging purposes. In normal operation, the middleware is automatically cleared after process(). """ - _current_middleware.instance = None + _current_middleware.set(None) def get_current_middleware() -> LabelTrackingFunctionMiddleware | None: - """Get the current middleware instance from thread-local storage. + """Get the current middleware instance from async-safe context storage. This function allows tools to access the middleware's variable store. Returns: The current LabelTrackingFunctionMiddleware instance, or None if not set. """ - return getattr(_current_middleware, "instance", None) + return _current_middleware.get(None) -class _PendingPolicyApproval(NamedTuple): - """Immutable binding record for a pending policy-violation approval. +def get_current_session() -> Any: + """Get the current AgentSession from async-safe context storage.""" + return _current_middleware_session.get(None) + + +class _FidesAuditRunContext: + """Mutable audit counters shared by calls in one async run context.""" + + def __init__(self, scope: object, turn_number: int) -> None: + self.scope = scope + self.turn_number = turn_number + self.call_index = 0 - 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. - """ - body_signature: str - label_key: str - session_key: str - disclosed_violations: tuple[str, ...] +_fides_audit_runs_by_owner: ContextVar[weakref.WeakKeyDictionary[Any, _FidesAuditRunContext] | None] = ContextVar( + "_fides_audit_runs_by_owner", default=None +) + + +def _set_direct_audit_run(owner: Any, run: _FidesAuditRunContext) -> None: + """Bind a direct-call audit run to one middleware in this async context.""" + current = _fides_audit_runs_by_owner.get(None) + runs = weakref.WeakKeyDictionary[Any, _FidesAuditRunContext]() + if current is not None: + runs.update(current) + runs[owner] = run + _fides_audit_runs_by_owner.set(runs) + + +def _get_direct_audit_run(owner: Any) -> _FidesAuditRunContext | None: + """Return the direct-call audit run bound to one middleware.""" + runs = _fides_audit_runs_by_owner.get(None) + return runs.get(owner) if runs is not None else None @experimental(feature_id=ExperimentalFeature.FIDES) @@ -1651,8 +1864,9 @@ class PolicyEnforcementFunctionMiddleware(FunctionMiddleware): Attributes: allow_untrusted_tools: Set of tool names allowed to execute in an untrusted context. + deny_untrusted_tools: Set of tool names denied in an untrusted context. block_on_violation: Whether to block execution on policy violations. - audit_log: List of policy violation events for audit purposes. + enable_audit_log: Whether to record policy violations in the active session. Examples: .. code-block:: python @@ -1674,34 +1888,122 @@ class PolicyEnforcementFunctionMiddleware(FunctionMiddleware): def __init__( self, allow_untrusted_tools: set[str] | None = None, + deny_untrusted_tools: set[str] | None = None, block_on_violation: bool = True, enable_audit_log: bool = True, approval_on_violation: bool = False, + max_audit_log_entries: int | None = 1000, ) -> None: """Initialize PolicyEnforcementFunctionMiddleware. Args: allow_untrusted_tools: Set of tool names allowed to execute in an untrusted context. + deny_untrusted_tools: Set of tool names denied in an untrusted context. Deny takes precedence. block_on_violation: Whether to block execution on policy violations. Ignored if approval_on_violation is True. + A blocked call is returned as a correlated function result and counts toward the + framework's consecutive-error guard. With the default limit, three consecutive + blocked calls disable tools for the remainder of that request to prevent retry loops. enable_audit_log: Whether to maintain an audit log of violations. + max_audit_log_entries: Maximum retained audit records per session. Set to ``None`` + for an unlimited log. 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. """ self.allow_untrusted_tools = allow_untrusted_tools or set() + self.deny_untrusted_tools = deny_untrusted_tools or set() self.approval_on_violation = approval_on_violation # If approval_on_violation is True, we don't block - we request approval instead self.block_on_violation = block_on_violation if not approval_on_violation else False self.enable_audit_log = enable_audit_log - self.audit_log: list[dict[str, Any]] = [] - # Track call_ids awaiting approval, each mapped to a binding record capturing the exact - # invocation the approval was requested for: the function name + arguments, the security - # label (integrity/confidentiality) shown for review, and the session. Combined with the - # call_id key and consume-on-use, an approval cannot re-authorize a repeated call, a - # different function, changed arguments, a different security label, or a different session. - self._pending_policy_approvals: dict[str, _PendingPolicyApproval] = {} + if max_audit_log_entries is not None and max_audit_log_entries < 1: + raise ValueError("max_audit_log_entries must be at least 1 or None") + self.max_audit_log_entries = max_audit_log_entries + # State is stored in AgentSession.state so a shared middleware instance cannot + # cross-contaminate conversations. These fallbacks only serve calls made without a session. + self._fallback_audit_log: list[dict[str, Any]] = [] + self._fallback_pending_policy_approvals: dict[str, dict[str, Any]] = {} + self._fallback_turn_counter = 0 + self._fallback_audit_scope = object() + + @property + def _pending_policy_approvals(self) -> dict[str, dict[str, Any]]: + """Compatibility view for direct, no-session middleware callers.""" + return self._fallback_pending_policy_approvals + + @_pending_policy_approvals.setter + def _pending_policy_approvals(self, value: dict[str, dict[str, Any]]) -> None: + self._fallback_pending_policy_approvals = value + + @staticmethod + def _fides_state(session: Any) -> dict[str, Any] | None: + if session is None: + return None + return _fides_session_state(session) + + def _resolve_audit_log(self, session: Any = None) -> list[dict[str, Any]]: + state = self._fides_state(session) + if state is None: + return self._fallback_audit_log + return cast(list[dict[str, Any]], state["audit_log"]) + + def _resolve_pending_approvals(self, session: Any = None) -> dict[str, dict[str, Any]]: + state = self._fides_state(session) + if state is None: + return self._fallback_pending_policy_approvals + approvals = cast(dict[str, Any], state["pending_policy_approvals"]) + # Older in-memory versions used a NamedTuple, which durable session state serializes + # as a positional list. Normalize that shape when a session is resumed. + for call_id, record in list(approvals.items()): + legacy_record = cast(list[Any] | tuple[Any, ...], record) if isinstance(record, (list, tuple)) else None + if legacy_record is not None and len(legacy_record) == 4: + approvals[call_id] = { + "body_signature": legacy_record[0], + "label_key": legacy_record[1], + "session_key": legacy_record[2], + "disclosed_violations": list(legacy_record[3]), + } + return cast(dict[str, dict[str, Any]], approvals) + + def begin_turn(self, session: Any) -> int: + """Advance the turn counter and bind direct middleware calls to the new turn.""" + turn = self._advance_turn(session) + _set_direct_audit_run(self, _FidesAuditRunContext(self._audit_scope(session), turn)) + return turn + + def _audit_scope(self, session: Any) -> object: + state = self._fides_state(session) + return state if state is not None else self._fallback_audit_scope + + def _advance_turn(self, session: Any) -> int: + """Return the next turn without binding it to the current async context.""" + state = self._fides_state(session) + if state is None: + self._fallback_turn_counter += 1 + return self._fallback_turn_counter + turn = int(state["turn_counter"]) + 1 + state["turn_counter"] = turn + return turn + + def _resolve_direct_audit_run(self, session: Any = None) -> _FidesAuditRunContext: + scope = self._audit_scope(session) + run = _get_direct_audit_run(self) + if run is None or run.scope is not scope: + self.begin_turn(session) + run = _get_direct_audit_run(self) + if run is None: # pragma: no cover - begin_turn always binds a context + raise RuntimeError("FIDES audit run context was not initialized") + return run + + def _resolve_turn_counter(self, session: Any = None) -> int: + return self._resolve_direct_audit_run(session).turn_number + + def _resolve_call_counter(self, session: Any = None) -> int: + run = self._resolve_direct_audit_run(session) + run.call_index += 1 + return run.call_index def _get_call_id(self, context: FunctionInvocationContext) -> str: """Get the tool call id for this invocation context.""" @@ -1722,14 +2024,23 @@ def _current_arguments(self, context: FunctionInvocationContext) -> dict[str, An return context.arguments.model_dump() return dict(context.arguments) - def _build_function_call_content(self, context: FunctionInvocationContext) -> Content: - """Reconstruct the current function call as Content for approval requests.""" + def build_function_call_content(self, context: FunctionInvocationContext) -> Content: + """Build the function call shown in an approval request. + + This is a public extension point for hosts that need to add provider- or + application-specific metadata to approval requests. Subclasses may override + it without depending on a private method. + """ return Content.from_function_call( call_id=self._get_call_id(context), name=context.function.name, arguments=self._current_arguments(context), ) + def _build_function_call_content(self, context: FunctionInvocationContext) -> Content: + """Compatibility alias for the former private extension point.""" + return self.build_function_call_content(context) + 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: @@ -1783,14 +2094,14 @@ def _pending_record( self, context: FunctionInvocationContext, violations: list[dict[str, Any]], - ) -> _PendingPolicyApproval: + ) -> dict[str, Any]: """Build the binding record for the current invocation and disclosed violation set.""" - return _PendingPolicyApproval( - body_signature=self._call_body_signature(context), - label_key=self._context_label_key(context), - session_key=self._session_key(context), - disclosed_violations=self._violation_set_key(violations), - ) + return { + "body_signature": self._call_body_signature(context), + "label_key": self._context_label_key(context), + "session_key": self._session_key(context), + "disclosed_violations": list(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.""" @@ -1840,7 +2151,7 @@ def _matches_pending_approval( call_id = self._get_call_id(context) if not call_id: return False - pending = self._pending_policy_approvals.get(call_id) + pending = self._resolve_pending_approvals(context.session).get(call_id) if pending is None: return False approval_response = context.metadata.get("approval_response") @@ -1854,11 +2165,11 @@ def _matches_pending_approval( # 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, 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 + self._response_matches_pending(approval_response, call_id, pending.get("body_signature", "")) + and self._call_body_signature(context) == pending.get("body_signature") + and self._context_label_key(context) == pending.get("label_key") + and self._session_key(context) == pending.get("session_key") + and self._violation_set_key(current_violations) == tuple(pending.get("disclosed_violations", ())) ) def _consume_pending_approval(self, context: FunctionInvocationContext) -> None: @@ -1867,7 +2178,7 @@ def _consume_pending_approval(self, context: FunctionInvocationContext) -> None: Idempotent: safe to call for both the integrity and confidentiality checks of a single invocation. """ - self._pending_policy_approvals.pop(self._get_call_id(context), None) + self._resolve_pending_approvals(context.session).pop(self._get_call_id(context), None) def _mark_policy_violation_approved( self, @@ -1900,7 +2211,7 @@ def _request_policy_violation_approval( ) call_id = self._get_call_id(context) if call_id: - self._pending_policy_approvals[call_id] = self._pending_record(context, violations) + self._resolve_pending_approvals(context.session)[call_id] = self._pending_record(context, violations) additional_properties: dict[str, Any] = { "policy_violation": True, "violation_type": primary["violation_type"], @@ -1916,6 +2227,7 @@ def _request_policy_violation_approval( additional_properties["violations"] = [ {"violation_type": v["violation_type"], "reason": v["approval_reason"]} for v in violations ] + additional_properties["_fides_violations"] = [v["violation_type"] for v in violations] context.result = Content.from_function_approval_request( id=call_id, function_call=self._build_function_call_content(context), @@ -1932,17 +2244,21 @@ def _block_policy_violation( ) -> None: """Block the tool call and surface the detected policy violation(s).""" primary = violations[0] + violation_types = [v["violation_type"] for v in violations] result: dict[str, Any] = { "error": primary["block_error"], "function": context.function.name, "context_label": context_label.to_dict(), + "blocked_violation": True, + "policy_violation": True, + "_fides_violations": violation_types, } if primary["block_violation_type"] is not None: result["violation_type"] = primary["block_violation_type"] if len(violations) > 1: - result["violations"] = [v["violation_type"] for v in violations] + result["violations"] = violation_types context.result = result - raise MiddlewareTermination("Policy violation blocked tool execution") + raise MiddlewareTermination("Policy violation blocked tool execution", blocked_policy=True) async def process( self, @@ -1960,6 +2276,18 @@ async def process( call_next: Callback to continue to next middleware or function execution. """ function_name = context.function.name + _remember_session(self, context.session) + turn_number: int | None = None + call_index: int | None = None + if self.enable_audit_log: + run_turn = context.metadata.get(_FIDES_TURN_NUMBER_KEY) + run_call_index = context.metadata.get(_FIDES_CALL_INDEX_KEY) + if type(run_turn) is int and type(run_call_index) is int: + turn_number = run_turn + call_index = run_call_index + else: + turn_number = self._resolve_turn_counter(context.session) + call_index = self._resolve_call_counter(context.session) # Get the context label (cumulative security state of the conversation) # This is set by LabelTrackingFunctionMiddleware and represents the @@ -1999,10 +2327,9 @@ async def process( # 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 ( + function_name in self.deny_untrusted_tools + or (function_name not in self.allow_untrusted_tools and not function_props.get("accepts_untrusted", False)) ): violations.append({ "violation_type": "untrusted_context", @@ -2018,7 +2345,8 @@ async def process( "type": "untrusted_context", "function": function_name, "context_label": context_label.to_dict(), - "turn": context.metadata.get("turn_number", -1), + "turn": turn_number, + "call_index": call_index, "reason": "Context is UNTRUSTED and tool is not allowed to execute in an untrusted context", }, }) @@ -2041,7 +2369,8 @@ async def process( "function": function_name, "context_label": context_label.to_dict(), "reason": conf_result["reason"], - "turn": context.metadata.get("turn_number", -1), + "turn": turn_number, + "call_index": call_index, }, }) @@ -2052,7 +2381,7 @@ async def process( return for violation in violations: - self._log_violation(violation["audit"]) + self._log_violation(violation["audit"], context.session) # 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 @@ -2160,28 +2489,69 @@ def _check_confidentiality_policy_detailed( return {"passed": True, "failure_type": None, "reason": None} - def _log_violation(self, violation: dict[str, Any]) -> None: + def _log_violation(self, violation: dict[str, Any], session: Any = None) -> None: """Log a policy violation. Args: violation: Dictionary containing violation details. + session: Optional session whose audit log should receive the violation. """ if self.enable_audit_log: - self.audit_log.append(violation) + audit_log = self._resolve_audit_log(session) + audit_log.append(violation) + if self.max_audit_log_entries is not None and len(audit_log) > self.max_audit_log_entries: + del audit_log[: -self.max_audit_log_entries] logger.warning(f"Policy violation detected: {violation}") - def get_audit_log(self) -> list[dict[str, Any]]: + def get_audit_log(self, session: Any = None) -> list[dict[str, Any]]: """Get the audit log of policy violations. Returns: List of violation records. """ - return self.audit_log.copy() + active_session = _resolve_owner_session(self, session) + return self._resolve_audit_log(active_session).copy() - def clear_audit_log(self) -> None: + def clear_audit_log(self, session: Any = None) -> None: """Clear the audit log.""" - self.audit_log.clear() + active_session = _resolve_owner_session(self, session) + self._resolve_audit_log(active_session).clear() + + +class _FidesRunContextFunctionMiddleware(FunctionMiddleware): + """Bind per-run audit metadata and the configured quarantine client.""" + + def __init__( + self, + *, + turn_number: int | None, + quarantine_client: SupportsChatGetResponse | None, + bind_quarantine_client: bool, + ) -> None: + self._turn_number = turn_number + self._call_index = 0 + self._quarantine_client = quarantine_client + self._bind_quarantine_client = bind_quarantine_client + + async def process( + self, + context: FunctionInvocationContext, + call_next: Callable[[], Awaitable[None]], + ) -> None: + if self._turn_number is not None: + self._call_index += 1 + context.metadata[_FIDES_TURN_NUMBER_KEY] = self._turn_number + context.metadata[_FIDES_CALL_INDEX_KEY] = self._call_index + + quarantine_token = ( + _quarantine_chat_client.set(self._quarantine_client) if self._bind_quarantine_client else None + ) + try: + await call_next() + finally: + if quarantine_token is not None: + _quarantine_chat_client.reset(quarantine_token) @experimental(feature_id=ExperimentalFeature.FIDES) @@ -2197,15 +2567,9 @@ class SecureAgentConfig(ContextProvider): auto_hide_untrusted: Whether to automatically hide untrusted content. Note: - The quarantine chat client is stored per-instance (see ``get_quarantine_client``) - but is *also* registered in a single process-global slot via - ``set_quarantine_client``. The ``quarantined_llm`` tool always reads that global - slot, so the behavior is **last-writer-wins**: when multiple ``SecureAgentConfig`` - instances are constructed in the same process with different ``quarantine_chat_client`` - values, the most recently constructed instance's client is the one every agent's - ``quarantined_llm`` tool will use. Running multiple instances is supported, but they - share this one global quarantine client rather than each using their own. If you need - distinct quarantine clients per agent, run them in separate processes. + Each quarantine client is bound and restored around tool execution. Concurrent + agents therefore resolve their own client without a process-global + last-writer-wins race or leaking it into a later run. Examples: .. code-block:: python @@ -2234,25 +2598,38 @@ class SecureAgentConfig(ContextProvider): def __init__( self, auto_hide_untrusted: bool = True, + enable_quarantine: bool = True, default_integrity: IntegrityLabel = IntegrityLabel.UNTRUSTED, default_confidentiality: ConfidentialityLabel = ConfidentialityLabel.PUBLIC, allow_untrusted_tools: set[str] | None = None, + deny_untrusted_tools: set[str] | None = None, + tool_labels: dict[str, ContentLabel] | None = None, block_on_violation: bool = True, approval_on_violation: bool = False, enable_audit_log: bool = True, enable_policy_enforcement: bool = True, quarantine_chat_client: SupportsChatGetResponse | None = None, source_id: str | None = None, + max_audit_log_entries: int | None = 1000, ) -> None: """Initialize secure agent configuration. Args: auto_hide_untrusted: Whether to automatically hide UNTRUSTED content. + enable_quarantine: Whether to inject the quarantine tools and instructions. + Set this to False when labels and policy enforcement are wanted without + variable hiding or quarantined LLM calls. In that mode, + ``auto_hide_untrusted`` must also be False. default_integrity: Default integrity label for tool calls. default_confidentiality: Default confidentiality label for tool calls. allow_untrusted_tools: Set of tool names allowed to execute in an untrusted context. + deny_untrusted_tools: Set of tool names denied in an untrusted context. Deny takes precedence. + tool_labels: Labels for tools constructed by a harness or MCP client. block_on_violation: Whether to block execution on policy violations. - Ignored if approval_on_violation is True. + Ignored if approval_on_violation is True. A blocked call is returned as a + correlated function result and counts toward the framework's consecutive-error + guard; with the default limit, three consecutive blocked calls disable tools for + the remainder of that request to prevent retry loops. 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 @@ -2263,41 +2640,49 @@ def __init__( If provided, the quarantined_llm tool will make actual isolated LLM calls instead of returning placeholder responses. This client should ideally be a separate instance using a cheaper model (e.g., gpt-4o-mini) since it - processes untrusted content. Note: this client is registered in a - process-global slot shared by all instances (last-writer-wins); see the - class docstring for details on running multiple instances. + processes untrusted content. It is bound to the current async context + while tools execute. source_id: Optional source identifier for context provider attribution. Defaults to "secure_agent". + max_audit_log_entries: Maximum retained audit records per session. Set to ``None`` + for an unlimited log. """ super().__init__(source_id or self.DEFAULT_SOURCE_ID) + if not enable_quarantine and auto_hide_untrusted: + raise ValueError("auto_hide_untrusted must be False when enable_quarantine is False") + self._tool_labels = dict(tool_labels or {}) + self.enable_quarantine = enable_quarantine self.label_tracker = LabelTrackingFunctionMiddleware( - auto_hide_untrusted=auto_hide_untrusted, + auto_hide_untrusted=auto_hide_untrusted and enable_quarantine, default_integrity=default_integrity, default_confidentiality=default_confidentiality, + tool_labels=self._tool_labels, ) self.enable_policy_enforcement = enable_policy_enforcement if enable_policy_enforcement: # Always allow security tools to execute in an untrusted context - tools_allowing_untrusted = {"quarantined_llm", "inspect_variable"} + tools_allowing_untrusted: set[str] = set() + if enable_quarantine: + tools_allowing_untrusted.update({"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, + deny_untrusted_tools=deny_untrusted_tools, block_on_violation=block_on_violation, approval_on_violation=approval_on_violation, enable_audit_log=enable_audit_log, + max_audit_log_entries=max_audit_log_entries, ) else: self.policy_enforcer = None - # Store and configure quarantine client for real LLM calls + # Store the client on this configuration. It is bound and restored by the + # per-run function middleware, not at construction time. self._quarantine_chat_client = quarantine_chat_client - if quarantine_chat_client is not None: - set_quarantine_client(quarantine_chat_client) - logger.info("Quarantine chat client configured for real LLM calls") async def before_run( self, @@ -2319,9 +2704,30 @@ async def before_run( context: The invocation context - tools, instructions, and middleware are added here. state: The provider-scoped mutable state dict. """ + fides_state = _fides_session_state(session) + if fides_state["context_label"] is None: + self.label_tracker.reset_context_label(session) + _remember_session(self, session) + _remember_session(self.label_tracker, session) + if self.policy_enforcer is not None: + _remember_session(self.policy_enforcer, session) + turn_number = None + if self.policy_enforcer is not None and self.policy_enforcer.enable_audit_log: + turn_number = self.policy_enforcer.begin_turn(session) + context.extend_tools(self.source_id, self.get_tools()) context.extend_instructions(self.source_id, self.get_instructions()) - context.extend_middleware(self.source_id, self.get_middleware()) + middleware = self.get_middleware() + if turn_number is not None or self.enable_quarantine: + middleware = [ + _FidesRunContextFunctionMiddleware( + turn_number=turn_number, + quarantine_client=self._quarantine_chat_client, + bind_quarantine_client=self.enable_quarantine, + ), + *middleware, + ] + context.extend_middleware(self.source_id, middleware) def get_tools(self) -> list[FunctionTool]: """Get the security tools for agent integration. @@ -2329,7 +2735,24 @@ def get_tools(self) -> list[FunctionTool]: Returns: List containing quarantined_llm and inspect_variable tools. """ - return self.label_tracker.get_security_tools() + return self.label_tracker.get_security_tools() if self.enable_quarantine else [] + + @property + def tool_labels(self) -> dict[str, ContentLabel]: + """Return labels configured for harness and MCP tools.""" + return self._tool_labels + + def set_tool_label( + self, + tool_name: str, + integrity: IntegrityLabel = IntegrityLabel.UNTRUSTED, + confidentiality: ConfidentialityLabel = ConfidentialityLabel.PUBLIC, + ) -> None: + """Configure the provenance label for a tool by name.""" + self._tool_labels[tool_name] = ContentLabel( + integrity=integrity, + confidentiality=confidentiality, + ) def get_instructions(self) -> str: """Get the security instructions for agent integration. @@ -2337,7 +2760,7 @@ def get_instructions(self) -> str: Returns: String containing security tool usage instructions. """ - return self.label_tracker.get_security_instructions() + return self.label_tracker.get_security_instructions() if self.enable_quarantine else "" def get_middleware(self) -> list[FunctionMiddleware]: """Get the middleware stack for agent integration. @@ -2350,31 +2773,34 @@ def get_middleware(self) -> list[FunctionMiddleware]: middleware.append(self.policy_enforcer) return middleware - def get_audit_log(self) -> list[dict[str, Any]]: + def get_audit_log(self, session: Any = None) -> list[dict[str, Any]]: """Get the audit log from policy enforcement. Returns: List of violation records, or empty list if policy enforcement disabled. """ if self.policy_enforcer: - return self.policy_enforcer.get_audit_log() + active_session = _resolve_owner_session(self, session) + return self.policy_enforcer.get_audit_log(active_session) return [] - def get_variable_store(self) -> ContentVariableStore: + def get_variable_store(self, session: Any = None) -> ContentVariableStore: """Get the variable store for this configuration. Returns: The ContentVariableStore instance. """ - return self.label_tracker.get_variable_store() + active_session = _resolve_owner_session(self, session) + return self.label_tracker.get_variable_store(active_session) - def list_variables(self) -> list[str]: + def list_variables(self, session: Any = None) -> list[str]: """Get a list of all stored variable IDs. Returns: List of variable ID strings. """ - return self.label_tracker.list_variables() + active_session = _resolve_owner_session(self, session) + return self.label_tracker.list_variables(active_session) def get_quarantine_client(self) -> SupportsChatGetResponse | None: """Get the quarantine chat client. @@ -2392,12 +2818,13 @@ def get_quarantine_client(self) -> SupportsChatGetResponse | None: # Global variable store instance (can be made per-session or injected) _global_variable_store = ContentVariableStore() -# Global quarantine chat client (set via set_quarantine_client or SecureAgentConfig) -_quarantine_chat_client: SupportsChatGetResponse | None = None +# Async-context-local quarantine chat client. Each SecureAgentConfig binds and +# restores its own client around tool execution. +_quarantine_chat_client: ContextVar[Any] = ContextVar("_quarantine_chat_client", default=None) def set_quarantine_client(client: SupportsChatGetResponse | None) -> None: - """Set the global quarantine chat client. + """Set the quarantine chat client for the current async context. This client will be used by quarantined_llm to make actual LLM calls in an isolated context. The client should ideally be a separate instance @@ -2421,8 +2848,7 @@ def set_quarantine_client(client: SupportsChatGetResponse | None) -> None: ) set_quarantine_client(quarantine_client) """ - global _quarantine_chat_client - _quarantine_chat_client = client + _quarantine_chat_client.set(client) if client: logger.info("Quarantine chat client set") else: @@ -2435,7 +2861,7 @@ def get_quarantine_client() -> SupportsChatGetResponse | None: Returns: The quarantine chat client, or None if not set. """ - return _quarantine_chat_client + return _quarantine_chat_client.get(None) # Security instructions that teach the agent how to handle variable references @@ -3029,9 +3455,10 @@ def _map_mcp_annotations_to_labels( Mapping rules (conservative - when in doubt, default to UNTRUSTED *source* and PUBLIC-only *sink*): - * ``readOnlyHint=True`` -> ``accepts_untrusted=True`` (pure data source, - safe to call even when the context is tainted - it cannot exfiltrate) - and **no** ``max_allowed_confidentiality`` cap. + * ``readOnlyHint=True`` -> ``accepts_untrusted=True`` and + ``max_allowed_confidentiality=PUBLIC``. A read-only tool may be safe to call + in a tainted context, but its query arguments still leave the process and can + carry private data to the remote server. * ``readOnlyHint`` is anything other than ``True`` (``False`` *or* missing) -> treated as a potential write / exfiltration sink: ``max_allowed_confidentiality = PUBLIC`` and ``accepts_untrusted = False``. @@ -3050,9 +3477,9 @@ def _map_mcp_annotations_to_labels( Returns: A ``(integrity, max_confidentiality, accepts_untrusted)`` tuple. - ``max_confidentiality`` is ``None`` for read-only / source tools and - ``PUBLIC`` for sinks. ``accepts_untrusted`` is ``True`` for read-only - tools that are safe to invoke in a tainted context. + ``max_confidentiality`` is ``PUBLIC`` by default because both read and write + tool arguments can be an exfiltration channel. ``accepts_untrusted`` is + ``True`` for read-only tools that are safe to invoke in a tainted context. """ if annotations is None: # No annotations at all - treat as both UNTRUSTED-by-default and a @@ -3074,19 +3501,12 @@ def _map_mcp_annotations_to_labels( # Closed-world tool (e.g., local memory) -> data is trusted integrity = IntegrityLabel.TRUSTED - # --- Determine max_allowed_confidentiality (sink detection) --- - # Conservative rule: only tools that *explicitly* declare ``readOnlyHint=True`` - # are treated as pure data sources. Everything else - including tools whose - # server omits the hint entirely - is treated as a potential write / sink - # and capped at PUBLIC confidentiality. This matters because many real - # servers (notably GitHub's MCP) declare ``readOnlyHint=True`` on read - # tools but leave *all* hints as ``None`` on their write tools - # (``push_files``, ``create_or_update_file``, ``create_pull_request``, - # ``create_repository``, ``merge_pull_request``, ...). Without this default, - # those write tools would bypass the exfiltration gate entirely. - max_confidentiality: ConfidentialityLabel | None = None - if read_only is not True: - max_confidentiality = ConfidentialityLabel.PUBLIC + # --- Determine max_allowed_confidentiality (argument egress) --- + # Read-only describes the result side of a tool, not its argument side. A + # search query or document identifier can still carry private information to + # the remote MCP server, so all annotation-derived tools receive a PUBLIC cap + # unless the caller supplies an explicit override. + max_confidentiality: ConfidentialityLabel | None = ConfidentialityLabel.PUBLIC # --- Determine accepts_untrusted --- # Read-only tools are pure data sources; they cannot exfiltrate data, @@ -3103,6 +3523,7 @@ async def apply_mcp_security_labels( default_integrity: IntegrityLabel = IntegrityLabel.UNTRUSTED, annotation_overrides: dict[str, tuple[IntegrityLabel, ConfidentialityLabel | None]] | None = None, mark_write_tools_as_sinks: bool = True, + mark_read_tools_as_sinks: bool = True, ) -> None: """Auto-assign FIDES security labels to every tool loaded from an MCP server. @@ -3128,6 +3549,9 @@ async def apply_mcp_security_labels( mark_write_tools_as_sinks: When ``True`` (default), non-read-only tools get ``max_allowed_confidentiality=PUBLIC`` to prevent data exfiltration via tool arguments. + mark_read_tools_as_sinks: When ``True`` (default), read-only tools also + get a PUBLIC argument cap. Set this to ``False`` only when the host + has independently established that read-only query arguments are safe. Raises: RuntimeError: If the ``MCPTool`` is not connected. @@ -3184,6 +3608,7 @@ async def apply_mcp_security_labels( continue # Check for explicit per-tool override first + has_explicit_override = remote_name in overrides if remote_name in overrides: integrity, max_conf = overrides[remote_name] accepts_untrusted = False # overrides must opt-in explicitly @@ -3196,9 +3621,27 @@ async def apply_mcp_security_labels( # Patch source_integrity (Tier 2 - read by LabelTrackingFunctionMiddleware) props["source_integrity"] = integrity.value - # Patch sink constraint - if mark_write_tools_as_sinks and max_conf is not None: + # Patch sink constraint. Read-only describes the result, not the arguments, + # so the secure default caps both classes while allowing a granular opt-out. + # An explicit confidentiality cap is a host decision and must not be + # discarded by either sink opt-out flag. The flags only control the + # annotation-derived defaults. + if has_explicit_override: + should_mark_sink = max_conf is not None + else: + is_read_only = getattr(annotation_map.get(remote_name), "readOnlyHint", None) is True + should_mark_sink = mark_read_tools_as_sinks if is_read_only else mark_write_tools_as_sinks + if should_mark_sink and max_conf is not None: props["max_allowed_confidentiality"] = max_conf.value + if has_explicit_override: + # This cap belongs to the host. Do not let a stale marker from + # an earlier annotation-derived pass remove it later. + props.pop("_fides_mcp_auto_max_confidentiality", None) + else: + props["_fides_mcp_auto_max_confidentiality"] = True + elif props.get("_fides_mcp_auto_max_confidentiality"): + props.pop("max_allowed_confidentiality", None) + props.pop("_fides_mcp_auto_max_confidentiality", None) # Allow read-only tools to execute even when context is tainted; # explicitly block write tools in untrusted contexts. @@ -3388,6 +3831,7 @@ def __init__( default_integrity: IntegrityLabel = IntegrityLabel.UNTRUSTED, annotation_overrides: dict[str, tuple[IntegrityLabel, ConfidentialityLabel | None]] | None = None, mark_write_tools_as_sinks: bool = True, + mark_read_tools_as_sinks: bool = True, ) -> None: """Initialize a secure proxy for an MCP tool or MCP URL endpoint. @@ -3409,6 +3853,8 @@ def __init__( annotation_overrides: Per-tool-name label overrides keyed by remote MCP tool name. mark_write_tools_as_sinks: Whether to restrict write tools to PUBLIC confidentiality. Defaults to ``True``. + mark_read_tools_as_sinks: Whether to restrict read-only tool arguments + to PUBLIC confidentiality. Defaults to ``True``. Raises: ValueError: If both ``mcp_tool`` and ``url`` are provided, or if neither is provided. @@ -3451,6 +3897,7 @@ def __init__( self._default_integrity = default_integrity self._annotation_overrides = annotation_overrides self._mark_write_tools_as_sinks = mark_write_tools_as_sinks + self._mark_read_tools_as_sinks = mark_read_tools_as_sinks # -- Async context manager -- @@ -3521,6 +3968,7 @@ async def _apply_labels(self) -> None: default_integrity=self._default_integrity, annotation_overrides=self._annotation_overrides, mark_write_tools_as_sinks=self._mark_write_tools_as_sinks, + mark_read_tools_as_sinks=self._mark_read_tools_as_sinks, ) # After static labels are stamped on each FunctionTool, install a # per-tool wrapper that consumes any server-provided ``_meta.ifc`` 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 1d6c70fb395..99de2bd45b4 100644 --- a/python/packages/core/tests/core/test_function_invocation_logic.py +++ b/python/packages/core/tests/core/test_function_invocation_logic.py @@ -4965,6 +4965,109 @@ async def process(self, context: FunctionInvocationContext, next_handler: Callab raise MiddlewareTermination +class BlockedPolicyMiddleware(FunctionMiddleware): + """Return a correlated policy refusal without executing the tool.""" + + async def process(self, context: FunctionInvocationContext, call_next: Callable[[], Awaitable[None]]) -> None: + context.result = { + "error": "Policy violation: tool execution blocked", + "blocked_violation": True, + "policy_violation": True, + } + raise MiddlewareTermination("Policy violation blocked tool execution", blocked_policy=True) + + +@pytest.mark.parametrize("streaming", [False, True], ids=["non-streaming", "streaming"]) +async def test_blocked_policy_result_continues_function_loop( + chat_client_base: SupportsChatGetResponse, + streaming: bool, +) -> None: + """A blocked policy result closes its call and gives the model another turn.""" + executions = 0 + + @tool(name="blocked_tool", approval_mode="never_require") + def blocked_tool() -> str: + nonlocal executions + executions += 1 + return "should not execute" + + function_call = Content.from_function_call(call_id="blocked-call", name="blocked_tool", arguments="{}") + final_text = "I could not run that tool because policy blocked it." + messages = [Message(role="user", contents=["run the blocked tool"])] + options: ChatOptions = {"tools": [blocked_tool]} + client_kwargs: dict[str, Any] = {"middleware": [BlockedPolicyMiddleware()]} + if streaming: + chat_client_base.streaming_responses = [ # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + [ChatResponseUpdate(role="assistant", contents=[function_call])], + [ChatResponseUpdate(role="assistant", contents=[Content.from_text(final_text)])], + ] + stream = chat_client_base.get_response( + messages, + stream=True, + options=options, + client_kwargs=client_kwargs, + ) + async for _ in stream: + pass + response = await stream.get_final_response() + else: + chat_client_base.run_responses = [ # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + ChatResponse(messages=Message(role="assistant", contents=[function_call])), + ChatResponse(messages=Message(role="assistant", contents=[final_text])), + ] + response = await chat_client_base.get_response( + messages, + options=options, + client_kwargs=client_kwargs, + ) + + function_results = [ + content for message in response.messages for content in message.contents if content.type == "function_result" + ] + assert executions == 0 + assert chat_client_base.call_count == 2 # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + assert len(function_results) == 1 + assert function_results[0].call_id == "blocked-call" + assert function_results[0].exception == "Policy violation: tool execution blocked" + assert function_results[0].additional_properties["blocked_violation"] is True + assert response.text == final_text + + +async def test_middleware_termination_does_not_trust_blocked_metadata( + chat_client_base: SupportsChatGetResponse, +) -> None: + """Generic termination must remain terminal even if call metadata says blocked.""" + executions = 0 + + @tool(name="terminated_tool", approval_mode="never_require") + def terminated_tool() -> str: + nonlocal executions + executions += 1 + return "should not execute" + + function_call = Content.from_function_call( + call_id="terminated-call", + name="terminated_tool", + arguments="{}", + additional_properties={"blocked_violation": True}, + ) + chat_client_base.run_responses = [ # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + ChatResponse(messages=Message(role="assistant", contents=[function_call])), + ChatResponse(messages=Message(role="assistant", contents=["must remain queued"])), + ] + + response = await chat_client_base.get_response( # type: ignore[call-overload, var-annotated] # pyrefly: ignore[no-matching-overload] # ty: ignore[no-matching-overload] + "run the tool", + options={"tool_choice": "auto", "tools": [terminated_tool]}, + client_kwargs={"middleware": [TerminateLoopMiddleware()]}, + ) + + assert executions == 0 + assert chat_client_base.call_count == 1 # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + assert len(chat_client_base.run_responses) == 1 # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + assert response.messages[-1].contents[0].result == "terminated by middleware" + + @pytest.mark.parametrize("approved", [True, False], ids=["approved", "rejected"]) async def test_streaming_approval_resume_yields_terminal_result_before_model_text( chat_client_base: SupportsChatGetResponse, diff --git a/python/packages/core/tests/core/test_harness_tool_approval.py b/python/packages/core/tests/core/test_harness_tool_approval.py index b637bc751e5..6475a7624b8 100644 --- a/python/packages/core/tests/core/test_harness_tool_approval.py +++ b/python/packages/core/tests/core/test_harness_tool_approval.py @@ -25,6 +25,11 @@ tool, ) from agent_framework._feature_stage import ExperimentalWarning +from agent_framework._harness._tool_approval import ( + ToolApprovalRule, + _function_call_from_request, + _has_policy_violation, +) from .conftest import MockBaseChatClient @@ -1426,3 +1431,24 @@ def optional_args_tool(value: str = "default") -> str: requests = _approval_requests(second_response.messages) assert [_function_call(request).arguments for request in requests] == ['{"value": "custom"}'] assert calls == 1 + + +async def test_policy_violations_are_not_auto_approved_by_standing_rules() -> None: + request = Content.from_function_approval_request( + id="call-1", + function_call=Content.from_function_call(call_id="call-1", name="write_file", arguments={}), + additional_properties={"policy_violation": True}, + ) + function_call = _function_call_from_request(request) + assert function_call is not None + assert function_call.additional_properties["policy_violation"] is True + assert _has_policy_violation(request) is True + + middleware = ToolApprovalMiddleware() + state = ToolApprovalState(rules=[ToolApprovalRule(tool_name="write_file")]) + messages = [Message(role="assistant", contents=[request])] + + all_auto_approved = await middleware._process_outbound_messages(messages, state) + + assert all_auto_approved is False + assert messages[0].contents == [request] diff --git a/python/packages/core/tests/core/test_observability.py b/python/packages/core/tests/core/test_observability.py index b2cd1603ad6..e3eab102197 100644 --- a/python/packages/core/tests/core/test_observability.py +++ b/python/packages/core/tests/core/test_observability.py @@ -2695,6 +2695,21 @@ def test_make_json_safe_dict_with_non_string_keys(): assert parsed["str_key"] == "normal" +def test_make_json_safe_bytes_are_base64_encoded(): + """Test make_json_safe converts bytes and bytearray values to base64 strings.""" + result = make_json_safe({ + "bytes": b"binary data", + "bytearray": bytearray(b"more binary data"), + "nested": [b"nested bytes"], + }) + + assert result == { + "bytes": "YmluYXJ5IGRhdGE=", + "bytearray": "bW9yZSBiaW5hcnkgZGF0YQ==", + "nested": ["bmVzdGVkIGJ5dGVz"], + } + + def test_to_otel_part_function_result(): """Test _to_otel_part with function_result content.""" from agent_framework import Content diff --git a/python/packages/core/tests/test_security.py b/python/packages/core/tests/test_security.py index 3b9932e9100..6080bb2e999 100644 --- a/python/packages/core/tests/test_security.py +++ b/python/packages/core/tests/test_security.py @@ -2,16 +2,37 @@ """Unit tests for prompt injection defense system.""" +import asyncio +import contextlib +import gc import json +import weakref from types import SimpleNamespace +from typing import Any, cast import pytest from pydantic import BaseModel -from agent_framework import AgentSession, ExperimentalFeature, FunctionInvocationContext, FunctionMiddleware +from agent_framework import ( + Agent, + AgentSession, + ExperimentalFeature, + FunctionInvocationContext, + FunctionMiddleware, + SessionContext, +) 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._sessions import FileSessionStore +from agent_framework._tools import ( + FunctionInvocationConfiguration, + FunctionTool, + _auto_invoke_function, + _execute_single_function_call, + _extract_function_calls, + _handle_function_call_results, + normalize_function_invocation_configuration, +) +from agent_framework._types import ChatResponse, Content, Message from agent_framework.security import ( ConfidentialityLabel, ContentLabel, @@ -23,10 +44,41 @@ PolicyEnforcementFunctionMiddleware, SecureAgentConfig, VariableReferenceContent, + _fides_session_state, combine_labels, + get_quarantine_client, + set_quarantine_client, store_untrusted_content, ) +TRUSTED = ContentLabel( + integrity=IntegrityLabel.TRUSTED, + confidentiality=ConfidentialityLabel.PUBLIC, +) +UNTRUSTED = ContentLabel( + integrity=IntegrityLabel.UNTRUSTED, + confidentiality=ConfidentialityLabel.PUBLIC, +) + + +class _Session: + def __init__(self, session_id: str) -> None: + self.session_id = session_id + self.state: dict[str, object] = {} + + +async def _noop() -> None: + pass + + +def _context(session: AgentSession | _Session, tool: FunctionTool, label: ContentLabel) -> FunctionInvocationContext: + return FunctionInvocationContext( + function=tool, + arguments={}, + session=cast(AgentSession, session), + metadata={"context_label": label}, + ) + class TestContentLabel: """Tests for ContentLabel class.""" @@ -74,6 +126,16 @@ def test_label_deserialization(self): assert label.confidentiality == ConfidentialityLabel.PRIVATE assert label.metadata["key"] == "value" + def test_label_metadata_is_durable_state_safe(self) -> None: + session = AgentSession(session_id="safe-label") + session.state["_fides"] = { + "context_label": ContentLabel(metadata={"opaque": object()}).to_dict(), + } + + restored = AgentSession.from_dict(session.to_dict()) + + assert restored.state["_fides"]["context_label"]["metadata"]["opaque"] + class TestSecurityFeatureStage: """Tests for security feature-stage annotations.""" @@ -615,6 +677,38 @@ def test_audit_log_recording(self, middleware, mock_function): initial_count = len(middleware.get_audit_log()) assert initial_count == 0 + async def test_audit_accessors_use_remembered_session(self) -> None: + tool = FunctionTool(name="get_config", description="test", fn=lambda: "secret") + middleware = PolicyEnforcementFunctionMiddleware(block_on_violation=True) + session = _Session("implicit-audit") + context = _context(session, tool, UNTRUSTED) + + with contextlib.suppress(MiddlewareTermination): + await middleware.process(context, _noop) + + assert middleware.get_audit_log() + middleware.clear_audit_log() + assert middleware.get_audit_log(session) == [] + + async def test_audit_accessors_remember_their_own_middleware_session(self) -> None: + middleware_a = PolicyEnforcementFunctionMiddleware(block_on_violation=True) + middleware_b = PolicyEnforcementFunctionMiddleware(block_on_violation=True) + session_a = _Session("policy-a") + session_b = _Session("policy-b") + tool_a = FunctionTool(name="tool_a", description="test", fn=lambda: "a") + tool_b = FunctionTool(name="tool_b", description="test", fn=lambda: "b") + + with contextlib.suppress(MiddlewareTermination): + await middleware_a.process(_context(session_a, tool_a, UNTRUSTED), _noop) + with contextlib.suppress(MiddlewareTermination): + await middleware_b.process(_context(session_b, tool_b, UNTRUSTED), _noop) + + assert [entry["function"] for entry in middleware_a.get_audit_log()] == ["tool_a"] + assert [entry["function"] for entry in middleware_b.get_audit_log()] == ["tool_b"] + middleware_a.clear_audit_log() + assert middleware_a.get_audit_log() == [] + assert [entry["function"] for entry in middleware_b.get_audit_log()] == ["tool_b"] + async def test_untrusted_call_requests_policy_approval(self, mock_function): """Test that policy violations can become approval requests.""" middleware = PolicyEnforcementFunctionMiddleware(approval_on_violation=True) @@ -1435,6 +1529,182 @@ async def execute() -> None: assert isinstance(replay_context.result, Content) assert replay_context.result.type == "function_approval_request" + async def test_policy_block_returns_content_and_audit_is_per_session(self) -> None: + tool = FunctionTool(name="get_config", description="test", fn=lambda: "secret") + middleware = PolicyEnforcementFunctionMiddleware(block_on_violation=True) + session_a = _Session("a") + session_b = _Session("b") + + async def must_not_execute() -> None: + raise AssertionError("blocked tool was executed") + + for _ in range(2): + context = _context(session_a, tool, UNTRUSTED) + with contextlib.suppress(MiddlewareTermination): + await middleware.process(context, must_not_execute) + assert isinstance(context.result, dict) + assert context.result["blocked_violation"] is True + + trusted_context = _context(session_b, tool, TRUSTED) + executed = False + + async def execute_trusted() -> None: + nonlocal executed + executed = True + + await middleware.process(trusted_context, execute_trusted) + + assert executed is True + audit = middleware.get_audit_log(session_a) + assert [entry["turn"] for entry in audit] == [1, 1] + assert [entry["call_index"] for entry in audit] == [1, 2] + assert middleware.get_audit_log(session_b) == [] + + async def test_blocked_result_is_visible_through_function_invoker(self) -> None: + async def get_config() -> str: + raise AssertionError("blocked tool was executed") + + tool = FunctionTool(name="get_config", description="test", fn=get_config) + session = AgentSession(session_id="a2") + tracker = LabelTrackingFunctionMiddleware(auto_hide_untrusted=False) + tracker.reset_context_label(session) + tracker._set_context_label(session, UNTRUSTED) + pipeline = FunctionMiddlewarePipeline( + tracker, + PolicyEnforcementFunctionMiddleware(block_on_violation=True), + ) + + result_groups, terminated = await _execute_single_function_call( + Content.from_function_call(call_id="call-a2", name="get_config", arguments={}), + custom_args={}, + config=FunctionInvocationConfiguration(), + tool_map={"get_config": tool}, + invocation_session=session, + middleware_pipeline=pipeline, + live_tools=None, + ) + + result = result_groups[0] + assert terminated is False + assert result.type == "function_result" + assert result.call_id == "call-a2" + assert result.exception and "Policy violation" in result.exception + assert result.result and "Policy violation" in result.result + assert result.additional_properties["blocked_violation"] is True + + # Exercise the same response path used by the agent loop: the blocked + # result must close the original call_id so the next model request is valid. + response = ChatResponse( + messages=[ + Message( + role="assistant", + contents=[Content.from_function_call(call_id="call-a2", name="get_config", arguments={})], + ) + ] + ) + processing = _handle_function_call_results( + response=response, + execution_results=[result], + function_call_count=1, + function_call_messages=None, + errors_in_a_row=0, + had_errors=True, + max_errors=3, + ) + assert processing.action == "continue" + assert response.messages[-1].contents[0].call_id == "call-a2" + assert _extract_function_calls(response) == [] + + async def test_public_approval_builder_is_used_for_policy_requests(self) -> None: + class CustomPolicy(PolicyEnforcementFunctionMiddleware): + def build_function_call_content(self, context): # type: ignore[no-untyped-def] + function_call = super().build_function_call_content(context) + function_call.additional_properties["host_marker"] = "custom" + return function_call + + tool = FunctionTool(name="write_file", description="test", fn=lambda: "written") + middleware = CustomPolicy(approval_on_violation=True) + context = _context(_Session("b4"), tool, UNTRUSTED) + context.metadata["call_id"] = "call-b4" + + with contextlib.suppress(MiddlewareTermination): + await middleware.process(context, _noop) + + assert context.result.type == "function_approval_request" + assert context.result.function_call.additional_properties["host_marker"] == "custom" + + async def test_private_approval_builder_override_remains_compatible(self) -> None: + class LegacyPolicy(PolicyEnforcementFunctionMiddleware): + def _build_function_call_content(self, context): # type: ignore[no-untyped-def] + function_call = super()._build_function_call_content(context) + function_call.additional_properties["legacy_host_marker"] = "custom" + return function_call + + tool = FunctionTool(name="write_file", description="test", fn=lambda: "written") + middleware = LegacyPolicy(approval_on_violation=True) + context = _context(_Session("b4-legacy"), tool, UNTRUSTED) + context.metadata["call_id"] = "call-b4-legacy" + + with contextlib.suppress(MiddlewareTermination): + await middleware.process(context, _noop) + + assert context.result.type == "function_approval_request" + assert context.result.function_call.additional_properties["legacy_host_marker"] == "custom" + + def test_no_session_fallback_counters_are_independent(self) -> None: + middleware = PolicyEnforcementFunctionMiddleware() + + assert middleware._resolve_turn_counter() == 1 + assert middleware._resolve_call_counter() == 1 + assert middleware._resolve_turn_counter() == 1 + assert middleware._resolve_call_counter() == 2 + assert middleware.begin_turn(None) == 2 + assert middleware._resolve_turn_counter() == 2 + assert middleware._resolve_call_counter() == 1 + + async def test_pending_approvals_are_json_persistible(self) -> None: + tool = FunctionTool(name="write_file", description="test", fn=lambda: "written") + middleware = PolicyEnforcementFunctionMiddleware(approval_on_violation=True) + session = AgentSession(session_id="approval") + context = _context(session, tool, UNTRUSTED) + context.metadata["call_id"] = "call-approval" + + with contextlib.suppress(MiddlewareTermination): + await middleware.process(context, _noop) + + pending = session.state["_fides"]["pending_policy_approvals"]["call-approval"] + assert isinstance(pending, dict) + restored = AgentSession.from_dict(session.to_dict()) + assert isinstance( + restored.state["_fides"]["pending_policy_approvals"]["call-approval"], + dict, + ) + + async def test_denylist_takes_precedence_over_allowlist(self) -> None: + tool = FunctionTool(name="sensitive_tool", description="test", fn=lambda: "secret") + middleware = PolicyEnforcementFunctionMiddleware( + allow_untrusted_tools={"sensitive_tool"}, + deny_untrusted_tools={"sensitive_tool"}, + ) + context = _context(_Session("deny"), tool, UNTRUSTED) + + async def must_not_execute() -> None: + raise AssertionError("denylisted tool was executed") + + with contextlib.suppress(MiddlewareTermination): + await middleware.process(context, must_not_execute) + assert isinstance(context.result, dict) + assert context.result["blocked_violation"] is True + + def test_audit_log_has_a_bounded_default(self) -> None: + middleware = PolicyEnforcementFunctionMiddleware(max_audit_log_entries=2) + session = _Session("audit-cap") + + for index in range(3): + middleware._log_violation({"index": index}, session) + + assert middleware.get_audit_log(session) == [{"index": 1}, {"index": 2}] + class TestAutomaticHiding: """Tests for automatic variable hiding functionality.""" @@ -1802,9 +2072,97 @@ async def next_fn(current_context=context, data=f"data_{i}"): assert len(variables) == 5 +@pytest.mark.asyncio +async def test_nested_hidden_results_use_their_invocation_session() -> None: + """Nested calls must not store an outer result in the inner session.""" + tracker = LabelTrackingFunctionMiddleware() + tool = FunctionTool(name="nested_tool", description="test", fn=lambda: "unused") + outer_session = AgentSession(session_id="nested-outer") + inner_session = AgentSession(session_id="nested-inner") + outer_context = _context(outer_session, tool, TRUSTED) + inner_context = _context(inner_session, tool, TRUSTED) + + async def inner_next() -> None: + inner_context.result = [Content.from_text("inner secret")] + + async def outer_next() -> None: + outer_context.result = [Content.from_text("outer secret")] + await tracker.process(inner_context, inner_next) + + await tracker.process(outer_context, outer_next) + + outer_variables = tracker.list_variables(outer_session) + inner_variables = tracker.list_variables(inner_session) + assert len(outer_variables) == 1 + assert len(inner_variables) == 1 + assert tracker.get_variable_store(outer_session).retrieve(outer_variables[0])[0] == "outer secret" + assert tracker.get_variable_store(inner_session).retrieve(inner_variables[0])[0] == "inner secret" + + class TestSecureAgentConfig: """Tests for SecureAgentConfig helper class.""" + def test_fides_session_state_initializes_the_complete_schema(self) -> None: + session = AgentSession(session_id="fides-state-schema") + + state = _fides_session_state(session) + + assert state == { + "context_label": None, + "audit_log": [], + "pending_policy_approvals": {}, + "turn_counter": 0, + "variables": {}, + "variable_metadata": {}, + } + + assert _fides_session_state(session) is state + + def test_variable_store_survives_session_round_trip(self) -> None: + tracker = LabelTrackingFunctionMiddleware() + session = AgentSession(session_id="durable-variables") + label = ContentLabel( + integrity=IntegrityLabel.UNTRUSTED, + confidentiality=ConfidentialityLabel.PRIVATE, + ) + variable_id = tracker.get_variable_store(session).store({"secret": "value"}, label) + tracker._resolve_metadata(session)[variable_id] = {"function_name": "load_secret"} + + restored = AgentSession.from_dict(session.to_dict()) + content, restored_label = tracker.get_variable_store(restored).retrieve(variable_id) + + assert content == {"secret": "value"} + assert restored_label.integrity == IntegrityLabel.UNTRUSTED + assert restored_label.confidentiality == ConfidentialityLabel.PRIVATE + assert tracker.get_variable_metadata(variable_id, restored) == {"function_name": "load_secret"} + + @pytest.mark.asyncio + async def test_variable_store_round_trips_binary_content_through_file_store(self, tmp_path) -> None: + tracker = LabelTrackingFunctionMiddleware() + session = AgentSession(session_id="durable-binary-variables") + payload = b"private-bytes" + variable_id = tracker.get_variable_store(session).store(payload, UNTRUSTED) + file_store = FileSessionStore(tmp_path) + + await file_store.set(session.session_id, session) + restored = await file_store.get(session.session_id) + + assert restored is not None + content, label = tracker.get_variable_store(restored).retrieve(variable_id) + assert content == payload + assert label.to_dict() == UNTRUSTED.to_dict() + + def test_variable_store_is_shared_with_a_propagated_child_session(self) -> None: + tracker = LabelTrackingFunctionMiddleware() + parent = AgentSession(session_id="shared-variables") + variable_id = tracker.get_variable_store(parent).store("hidden", UNTRUSTED) + child = AgentSession(session_id=parent.session_id) + child.state = parent.state + + content, label = tracker.get_variable_store(child).retrieve(variable_id) + assert content == "hidden" + assert label.to_dict() == UNTRUSTED.to_dict() + def test_create_config_defaults(self): """Test creating config with default values.""" from agent_framework.security import SecureAgentConfig @@ -1870,6 +2228,288 @@ def test_inspect_variable_uses_generic_approval_mode(self): assert inspect_variable.approval_mode == "never_require" assert "requires_approval" not in inspect_variable.additional_properties # type: ignore[operator] # pyrefly: ignore[not-iterable] # ty: ignore[unsupported-operator] + def test_secure_config_can_disable_quarantine_without_disabling_policy(self) -> None: + config = SecureAgentConfig(enable_quarantine=False, auto_hide_untrusted=False) + + assert config.get_tools() == [] + assert config.get_instructions() == "" + assert len(config.get_middleware()) == 2 + + def test_secure_config_rejects_hidden_content_without_quarantine_tooling(self) -> None: + try: + SecureAgentConfig(enable_quarantine=False) + except ValueError as exc: + assert "auto_hide_untrusted" in str(exc) + else: + raise AssertionError("Disabling quarantine must not leave hidden content without a handler") + + async def test_tool_labels_apply_to_harness_tools(self) -> None: + tool = FunctionTool(name="external_tool", description="test", fn=lambda: "external") + config = SecureAgentConfig(tool_labels={"external_tool": UNTRUSTED}) + labeler = config.get_middleware()[0] + context = _context(_Session("labels"), tool, TRUSTED) + + async def execute() -> None: + context.result = Content.from_text("external") + + await labeler.process(context, execute) + + assert tool.additional_properties is None + result = context.result[0] if isinstance(context.result, list) else context.result + assert result.additional_properties["security_label"]["integrity"] == "untrusted" + + async def test_before_run_does_not_clear_an_external_quarantine_client(self) -> None: + sentinel = object() + set_quarantine_client(cast(Any, sentinel)) + config = SecureAgentConfig(enable_quarantine=False, auto_hide_untrusted=False) + session = AgentSession(session_id="quarantine-preserve") + + class _Context: + def extend_tools(self, *args: object) -> None: + pass + + def extend_instructions(self, *args: object) -> None: + pass + + def extend_middleware(self, *args: object) -> None: + pass + + await config.before_run(agent=None, session=session, context=_Context(), state={}) + + assert get_quarantine_client() is sentinel + + async def test_public_accessors_use_the_last_session_in_the_async_context(self) -> None: + config = SecureAgentConfig(enable_quarantine=False, auto_hide_untrusted=False) + session = AgentSession(session_id="public-accessors") + + class _Context: + def extend_tools(self, *args: object) -> None: + pass + + def extend_instructions(self, *args: object) -> None: + pass + + def extend_middleware(self, *args: object) -> None: + pass + + await config.before_run(agent=None, session=session, context=_Context(), state={}) + policy_enforcer = config.policy_enforcer + assert policy_enforcer is not None + policy_enforcer._log_violation({"type": "test"}, session) + + assert config.get_audit_log() == [{"type": "test"}] + + async def test_public_accessors_remember_each_configs_own_session(self) -> None: + config_a = SecureAgentConfig(enable_quarantine=False, auto_hide_untrusted=False) + config_b = SecureAgentConfig(enable_quarantine=False, auto_hide_untrusted=False) + session_a = AgentSession(session_id="public-config-a") + session_b = AgentSession(session_id="public-config-b") + run_a = SessionContext(session_id=session_a.session_id, input_messages=[]) + run_b = SessionContext(session_id=session_b.session_id, input_messages=[]) + + await config_a.before_run(agent=None, session=session_a, context=run_a, state={}) + await config_b.before_run(agent=None, session=session_b, context=run_b, state={}) + policy_a = config_a.policy_enforcer + policy_b = config_b.policy_enforcer + assert policy_a is not None + assert policy_b is not None + policy_a._log_violation({"owner": "a"}, session_a) + policy_b._log_violation({"owner": "b"}, session_b) + variable_a = config_a.get_variable_store(session_a).store("a", UNTRUSTED) + variable_b = config_b.get_variable_store(session_b).store("b", TRUSTED) + config_a.label_tracker._set_context_label(session_a, UNTRUSTED) + config_b.label_tracker._set_context_label(session_b, TRUSTED) + + assert config_a.get_audit_log() == [{"owner": "a"}] + assert config_b.get_audit_log() == [{"owner": "b"}] + assert policy_a.get_audit_log() == [{"owner": "a"}] + assert policy_b.get_audit_log() == [{"owner": "b"}] + assert config_a.list_variables() == [variable_a] + assert config_b.list_variables() == [variable_b] + assert config_a.label_tracker.list_variables() == [variable_a] + assert config_b.label_tracker.list_variables() == [variable_b] + assert config_a.label_tracker.get_context_label().to_dict() == UNTRUSTED.to_dict() + assert config_b.label_tracker.get_context_label().to_dict() == TRUSTED.to_dict() + + async def test_public_accessors_keep_an_implicit_session_alive(self) -> None: + class RecordingConfig(SecureAgentConfig): + session_reference: weakref.ReferenceType[Any] | None = None + + async def before_run( + self, + *, + agent: Any, + session: Any, + context: Any, + state: dict[str, Any], + ) -> None: + await super().before_run(agent=agent, session=session, context=context, state=state) + self.session_reference = weakref.ref(session) + policy = self.policy_enforcer + assert policy is not None + policy._log_violation({"type": "retained"}, session) + + class ChatClient: + async def get_response(self, messages: Any, **kwargs: Any) -> ChatResponse: + return ChatResponse(messages=Message(role="assistant", contents=["done"])) + + config = RecordingConfig(enable_quarantine=False, auto_hide_untrusted=False) + agent = Agent(client=cast(Any, ChatClient()), context_providers=[config]) + await agent.run("hello") + gc.collect() + + assert config.session_reference is not None + assert config.session_reference() is not None + assert config.get_audit_log() == [{"type": "retained"}] + + async def test_public_accessors_remain_task_local_for_a_shared_config(self) -> None: + config = SecureAgentConfig(enable_quarantine=False, auto_hide_untrusted=False) + policy = config.policy_enforcer + assert policy is not None + + async def run(marker: str) -> list[dict[str, Any]]: + session = AgentSession(session_id=f"shared-config-{marker}") + run_context = SessionContext(session_id=session.session_id, input_messages=[]) + await config.before_run(agent=None, session=session, context=run_context, state={}) + policy._log_violation({"owner": marker}, session) + await asyncio.sleep(0) + return config.get_audit_log() + + audit_a, audit_b = await asyncio.gather(run("a"), run("b")) + + assert audit_a == [{"owner": "a"}] + assert audit_b == [{"owner": "b"}] + + async def test_tool_labels_preserve_server_supplied_result_labels(self) -> None: + server_label = ContentLabel( + integrity=IntegrityLabel.TRUSTED, + confidentiality=ConfidentialityLabel.PRIVATE, + ) + configured_label = ContentLabel( + integrity=IntegrityLabel.UNTRUSTED, + confidentiality=ConfidentialityLabel.PUBLIC, + ) + tool = FunctionTool(name="external_tool", description="test", fn=lambda: "external") + config = SecureAgentConfig(tool_labels={"external_tool": configured_label}) + context = _context(_Session("server-label"), tool, TRUSTED) + + async def execute() -> None: + context.result = Content.from_text( + "external", additional_properties={"security_label": server_label.to_dict()} + ) + + await config.get_middleware()[0].process(context, execute) + + assert tool.additional_properties is None + result = context.result[0] + assert result.additional_properties["security_label"] == server_label.to_dict() + + async def test_tool_labels_do_not_leak_between_configs(self) -> None: + tool = FunctionTool(name="external_tool", description="test", fn=lambda: "external") + config_a = SecureAgentConfig(tool_labels={"external_tool": UNTRUSTED}) + config_b = SecureAgentConfig(tool_labels={"external_tool": TRUSTED}) + + async def run(config: SecureAgentConfig, session_id: str) -> Content: + context = _context(_Session(session_id), tool, TRUSTED) + + async def execute() -> None: + context.result = Content.from_text("external") + + await config.get_middleware()[0].process(context, execute) + return context.result[0] + + result_a = await run(config_a, "labels-a") + result_b = await run(config_b, "labels-b") + + assert result_a.additional_properties["security_label"]["integrity"] == "untrusted" + assert result_b.additional_properties["security_label"]["integrity"] == "trusted" + assert tool.additional_properties is None + + async def test_propagated_child_run_does_not_change_parent_audit_coordinates(self) -> None: + config = SecureAgentConfig(enable_quarantine=False, auto_hide_untrusted=False) + policy = config.policy_enforcer + assert policy is not None + tool = FunctionTool(name="write_file", description="test", fn=lambda: "written") + parent = AgentSession(session_id="shared-audit") + parent_run = SessionContext(session_id=parent.session_id, input_messages=[]) + await config.before_run(agent=None, session=parent, context=parent_run, state={}) + parent_binder = cast(FunctionMiddleware, parent_run.get_middleware()[0]) + + async def invoke(binder: FunctionMiddleware, session: AgentSession) -> None: + context = _context(session, tool, UNTRUSTED) + + async def enforce() -> None: + await policy.process(context, _noop) + + with contextlib.suppress(MiddlewareTermination): + await binder.process(context, enforce) + + await invoke(parent_binder, parent) + + child = AgentSession(session_id=parent.session_id) + child.state = parent.state + child_run = SessionContext(session_id=child.session_id, input_messages=[]) + await config.before_run(agent=None, session=child, context=child_run, state={}) + child_binder = cast(FunctionMiddleware, child_run.get_middleware()[0]) + await invoke(child_binder, child) + await invoke(parent_binder, parent) + + assert [(entry["turn"], entry["call_index"]) for entry in policy.get_audit_log(parent)] == [ + (1, 1), + (2, 1), + (1, 2), + ] + + async def test_run_quarantine_binding_is_nested_and_restored(self) -> None: + external_client = object() + configured_client = object() + set_quarantine_client(cast(Any, external_client)) + try: + configured = SecureAgentConfig(quarantine_chat_client=cast(Any, configured_client)) + configured_session = AgentSession(session_id="configured-quarantine") + configured_run = SessionContext(session_id=configured_session.session_id, input_messages=[]) + await configured.before_run( + agent=None, + session=configured_session, + context=configured_run, + state={}, + ) + configured_binder = cast(FunctionMiddleware, configured_run.get_middleware()[0]) + + no_client = SecureAgentConfig() + no_client_session = AgentSession(session_id="empty-quarantine") + no_client_run = SessionContext(session_id=no_client_session.session_id, input_messages=[]) + await no_client.before_run( + agent=None, + session=no_client_session, + context=no_client_run, + state={}, + ) + no_client_binder = cast(FunctionMiddleware, no_client_run.get_middleware()[0]) + tool = FunctionTool(name="observe_client", description="test", fn=lambda: None) + observed: list[object | None] = [] + + async def observe_empty_client() -> None: + observed.append(get_quarantine_client()) + + async def observe_configured_client() -> None: + observed.append(get_quarantine_client()) + await no_client_binder.process( + _context(no_client_session, tool, TRUSTED), + observe_empty_client, + ) + observed.append(get_quarantine_client()) + + await configured_binder.process( + _context(configured_session, tool, TRUSTED), + observe_configured_client, + ) + + assert observed == [configured_client, None, configured_client] + assert get_quarantine_client() is external_client + finally: + set_quarantine_client(None) + class TestGetSecurityTools: """Tests for get_security_tools function.""" @@ -2198,6 +2838,16 @@ async def next_fn(): # Context should STILL be UNTRUSTED (once tainted, stays tainted) assert middleware.get_context_label().integrity == IntegrityLabel.UNTRUSTED + def test_context_label_isolated_by_session(self) -> None: + tracker = LabelTrackingFunctionMiddleware() + session_a = _Session("a") + session_b = _Session("b") + + tracker._set_context_label(session_a, UNTRUSTED) + + assert tracker._get_context_label(session_a).integrity == IntegrityLabel.UNTRUSTED + assert tracker._get_context_label(session_b).integrity == IntegrityLabel.TRUSTED + class TestPolicyEnforcementWithContextLabel: """Tests for policy enforcement using context labels.""" @@ -2403,7 +3053,7 @@ async def test_quarantined_llm_returns_response(self): ) # Set middleware context - _current_middleware.instance = middleware + _current_middleware.set(middleware) try: result = await quarantined_llm(prompt="Summarize this data", variable_ids=[var_id]) @@ -2413,7 +3063,7 @@ async def test_quarantined_llm_returns_response(self): assert result["quarantined"] is True assert "auto_hidden" not in result finally: - _current_middleware.instance = None + _current_middleware.set(None) @pytest.mark.asyncio async def test_quarantined_llm_trusted_input(self): @@ -2427,7 +3077,7 @@ async def test_quarantined_llm_trusted_input(self): "trusted system data", ContentLabel(integrity=IntegrityLabel.TRUSTED) ) - _current_middleware.instance = middleware + _current_middleware.set(middleware) try: result = await quarantined_llm( @@ -2439,7 +3089,7 @@ async def test_quarantined_llm_trusted_input(self): assert "response" in result assert result["quarantined"] is True finally: - _current_middleware.instance = None + _current_middleware.set(None) @pytest.mark.asyncio async def test_quarantined_llm_multiple_variables(self): @@ -2451,7 +3101,7 @@ async def test_quarantined_llm_multiple_variables(self): var1 = middleware.get_variable_store().store("data1", ContentLabel(integrity=IntegrityLabel.UNTRUSTED)) var2 = middleware.get_variable_store().store("data2", ContentLabel(integrity=IntegrityLabel.UNTRUSTED)) - _current_middleware.instance = middleware + _current_middleware.set(middleware) try: result = await quarantined_llm(prompt="Compare these", variable_ids=[var1, var2]) @@ -2460,7 +3110,7 @@ async def test_quarantined_llm_multiple_variables(self): assert result["quarantined"] is True assert result["variables_processed"] == [var1, var2] finally: - _current_middleware.instance = None + _current_middleware.set(None) def test_quarantined_llm_declares_source_integrity(self): """Test that quarantined_llm declares source_integrity='untrusted'.""" @@ -2514,8 +3164,8 @@ async def get_response(self, messages, **kwargs): # Create config with quarantine client config = SecureAgentConfig(quarantine_chat_client=mock_client) # type: ignore[arg-type] # pyrefly: ignore[bad-argument-type] # ty: ignore[invalid-argument-type] - # Should have set the global client - assert get_quarantine_client() is mock_client + # Construction must not mutate process/task-global state. The client is bound by before_run. + assert get_quarantine_client() is None # Config should also return the client assert config.get_quarantine_client() is mock_client @@ -2571,7 +3221,7 @@ async def test_quarantined_llm_uses_real_client_when_set(self): "Some email content with [INJECTION ATTEMPT]", ContentLabel(integrity=IntegrityLabel.UNTRUSTED) ) - _current_middleware.instance = middleware + _current_middleware.set(middleware) try: result = await quarantined_llm(prompt="Summarize this email", variable_ids=[var_id]) @@ -2597,7 +3247,7 @@ async def test_quarantined_llm_uses_real_client_when_set(self): assert result["response"] == "This is a safe summary of the content." finally: - _current_middleware.instance = None + _current_middleware.set(None) set_quarantine_client(None) @pytest.mark.asyncio @@ -2621,7 +3271,7 @@ async def test_quarantined_llm_fallback_without_client(self): ContentLabel(integrity=IntegrityLabel.TRUSTED), # Use trusted to see response directly ) - _current_middleware.instance = middleware + _current_middleware.set(middleware) try: result = await quarantined_llm( @@ -2634,7 +3284,7 @@ async def test_quarantined_llm_fallback_without_client(self): assert "[Quarantined LLM Response] Processed:" in result["response"] finally: - _current_middleware.instance = None + _current_middleware.set(None) @pytest.mark.asyncio async def test_quarantined_llm_handles_client_error(self): @@ -2659,7 +3309,7 @@ async def test_quarantined_llm_handles_client_error(self): middleware = LabelTrackingFunctionMiddleware() var_id = middleware.get_variable_store().store("Some content", ContentLabel(integrity=IntegrityLabel.TRUSTED)) - _current_middleware.instance = middleware + _current_middleware.set(middleware) try: result = await quarantined_llm(prompt="Process this", variable_ids=[var_id]) @@ -2670,7 +3320,7 @@ async def test_quarantined_llm_handles_client_error(self): assert "API Error" in result["response"] finally: - _current_middleware.instance = None + _current_middleware.set(None) set_quarantine_client(None) @pytest.mark.asyncio @@ -2706,7 +3356,7 @@ async def test_quarantined_llm_builds_correct_messages(self): ContentLabel(integrity=IntegrityLabel.UNTRUSTED), ) - _current_middleware.instance = middleware + _current_middleware.set(middleware) try: await quarantined_llm(prompt="Summarize both emails", variable_ids=[var1, var2]) @@ -2722,9 +3372,17 @@ async def test_quarantined_llm_builds_correct_messages(self): assert '"subject": "Test"' in user_message # Dict should be JSON serialized finally: - _current_middleware.instance = None + _current_middleware.set(None) set_quarantine_client(None) + async def test_quarantine_client_isolated_between_async_tasks(self) -> None: + async def run(client: object) -> None: + set_quarantine_client(cast(Any, client)) + await asyncio.sleep(0) + assert get_quarantine_client() is client + + await asyncio.gather(run(object()), run(object())) + # ========== Per-Item Embedded Label Tests ========== @@ -3478,9 +4136,9 @@ class TestMCPAnnotationMapping: @pytest.mark.parametrize( ("read_only", "open_world", "default_integrity", "expected_integrity", "expected_max_conf", "expected_accepts"), [ - (True, None, IntegrityLabel.UNTRUSTED, IntegrityLabel.UNTRUSTED, None, True), - (True, True, IntegrityLabel.TRUSTED, IntegrityLabel.UNTRUSTED, None, True), - (True, False, IntegrityLabel.UNTRUSTED, IntegrityLabel.TRUSTED, None, True), + (True, None, IntegrityLabel.UNTRUSTED, IntegrityLabel.UNTRUSTED, ConfidentialityLabel.PUBLIC, True), + (True, True, IntegrityLabel.TRUSTED, IntegrityLabel.UNTRUSTED, ConfidentialityLabel.PUBLIC, True), + (True, False, IntegrityLabel.UNTRUSTED, IntegrityLabel.TRUSTED, ConfidentialityLabel.PUBLIC, True), (False, None, IntegrityLabel.UNTRUSTED, IntegrityLabel.UNTRUSTED, ConfidentialityLabel.PUBLIC, False), (False, True, IntegrityLabel.TRUSTED, IntegrityLabel.UNTRUSTED, ConfidentialityLabel.PUBLIC, False), (False, False, IntegrityLabel.UNTRUSTED, IntegrityLabel.TRUSTED, ConfidentialityLabel.PUBLIC, False), @@ -3520,6 +4178,16 @@ def test_map_missing_annotations_defaults_to_sink(self): assert max_conf == ConfidentialityLabel.PUBLIC assert accepts_untrusted is False + def test_read_only_mcp_tools_cap_argument_confidentiality(self) -> None: + from agent_framework.security import _map_mcp_annotations_to_labels + + _, max_confidentiality, accepts_untrusted = _map_mcp_annotations_to_labels( + SimpleNamespace(readOnlyHint=True, openWorldHint=True) + ) + + assert max_confidentiality == ConfidentialityLabel.PUBLIC + assert accepts_untrusted is True + # --------------------------------------------------------------------------- # IFC labels from MCP _meta payload @@ -3835,3 +4503,47 @@ 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 test_mcp_read_only_sink_cap_has_a_granular_opt_out(self) -> None: + from agent_framework.security import apply_mcp_security_labels + + read_only = SimpleNamespace(readOnlyHint=True, openWorldHint=True) + tool = FunctionTool(name="search", description="test", fn=lambda: "result") + tool.additional_properties = {"_mcp_remote_name": "search"} + + class _McpSession: + async def list_tools(self, params: object = None) -> object: + return SimpleNamespace(tools=[SimpleNamespace(name="search", annotations=read_only)], nextCursor=None) + + mcp = SimpleNamespace(is_connected=True, session=_McpSession(), functions=[tool]) + await apply_mcp_security_labels(mcp, mark_read_tools_as_sinks=False) + + assert "max_allowed_confidentiality" not in tool.additional_properties + + async def test_mcp_explicit_override_cap_survives_read_only_sink_opt_out(self) -> None: + from agent_framework.security import apply_mcp_security_labels + + read_only = SimpleNamespace(readOnlyHint=True, openWorldHint=True) + tool = FunctionTool(name="search", description="test", fn=lambda: "result") + tool.additional_properties = {"_mcp_remote_name": "search"} + + class _McpSession: + async def list_tools(self, params: object = None) -> object: + return SimpleNamespace(tools=[SimpleNamespace(name="search", annotations=read_only)], nextCursor=None) + + mcp = SimpleNamespace(is_connected=True, session=_McpSession(), functions=[tool]) + await apply_mcp_security_labels( + mcp, + annotation_overrides={"search": (IntegrityLabel.TRUSTED, ConfidentialityLabel.PRIVATE)}, + mark_read_tools_as_sinks=False, + ) + + assert tool.additional_properties["max_allowed_confidentiality"] == "private" + assert "_fides_mcp_auto_max_confidentiality" not in tool.additional_properties + + # A later annotation-only pass must not mistake the host-owned cap for an + # auto-generated one and remove it when read-only sinks are opted out. + await apply_mcp_security_labels(mcp, mark_read_tools_as_sinks=False) + + assert tool.additional_properties["max_allowed_confidentiality"] == "private" + assert "_fides_mcp_auto_max_confidentiality" not in tool.additional_properties