Python: Harden FIDES core policy foundation

python/packages/core/agent_framework/_middleware.py

@@ -2,6 +2,8 @@
 
 from __future__ import annotations
 
+import asyncio
+import contextvars
 import inspect
 import logging
 import sys
@@ -38,6 +40,10 @@
 
 logger = logging.getLogger(__name__)
 
+_ACTIVE_FUNCTION_MIDDLEWARE_CONTEXT: contextvars.ContextVar[tuple[Any, asyncio.Task[Any] | None] | None] = (
+    contextvars.ContextVar("agent_framework_core_active_function_middleware_context", default=None)
+)
+
 if sys.version_info >= (3, 13):
     from typing import TypeVar  # pragma: no cover
 else:
@@ -1570,9 +1576,20 @@ def create_next_handler(index: int) -> Callable[[], Awaitable[None]]:
             if index >= len(self._middleware):
 
                 async def final_wrapper() -> None:
-                    context.result = final_handler(context)
-                    if inspect.isawaitable(context.result):
-                        context.result = await context.result
+                    from ._tools import (
+                        _SECURITY_ARGUMENT_MATERIALIZATION_CONTEXT,  # pyright: ignore[reportPrivateUsage]
+                        _suspend_security_argument_materialization,  # pyright: ignore[reportPrivateUsage]
+                    )
+
+                    security_materialization = _SECURITY_ARGUMENT_MATERIALIZATION_CONTEXT.get()
+                    if security_materialization is not None:
+                        active_context, materialize, origin_task = security_materialization
+                        if context is active_context and asyncio.current_task() is origin_task:
+                            materialize(context)
+                    with _suspend_security_argument_materialization():
+                        context.result = final_handler(context)
+                        if inspect.isawaitable(context.result):
+                            context.result = await context.result
 
                 return final_wrapper
 
@@ -1584,7 +1601,11 @@ async def current_handler() -> None:
 
         first_handler = create_next_handler(0)
         # Don't suppress MiddlewareTermination - let it propagate to signal loop termination
-        await first_handler()
+        pipeline_token = _ACTIVE_FUNCTION_MIDDLEWARE_CONTEXT.set((context, asyncio.current_task()))
+        try:
+            await first_handler()
+        finally:
+            _ACTIVE_FUNCTION_MIDDLEWARE_CONTEXT.reset(pipeline_token)
 
         return context.result
 

python/packages/core/agent_framework/_tools.py

@@ -17,11 +17,12 @@
     AsyncIterable,
     Awaitable,
     Callable,
+    Generator,
     Iterable,
     Mapping,
     Sequence,
 )
-from contextlib import suppress
+from contextlib import contextmanager, suppress
 from dataclasses import dataclass
 from datetime import date, datetime, timedelta, timezone, tzinfo
 from datetime import time as datetime_time
@@ -138,12 +139,25 @@ def _has_authoritative_approval_session(invocation_session: AgentSession | None)
 _SECURITY_ARGUMENTS_SNAPSHOT_CONTEXT_KEY: Final[str] = "_security_function_arguments"
 _PREPARED_ARGUMENTS_CONTEXT_KEY: Final[str] = "_prepared_function_arguments"
 _AUTO_ARGUMENT_PREPARATION_CONTEXT_KEY: Final[str] = "_auto_prepare_function_arguments"
+_SECURITY_ARGUMENT_MATERIALIZATION_CONTEXT: Final[
+    contextvars.ContextVar[tuple[Any, Callable[[Any], None], asyncio.Task[Any] | None] | None]
+] = contextvars.ContextVar("agent_framework_core_security_argument_materialization_context", default=None)
 _FUNCTION_INVOCATION_LIMIT_FALLBACK_TEXT: Final[str] = (
     "Function invocation limit reached before a final answer could be produced."
 )
 _USER_VISIBLE_CONTENT_TYPES: Final[set[str]] = {"data", "uri", "error", "hosted_file", "hosted_vector_store"}
 
 
+@contextmanager
+def _suspend_security_argument_materialization() -> Generator[None]:
+    """Prevent nested tool calls from inheriting one invocation's materialization capability."""
+    token = _SECURITY_ARGUMENT_MATERIALIZATION_CONTEXT.set(None)
+    try:
+        yield
+    finally:
+        _SECURITY_ARGUMENT_MATERIALIZATION_CONTEXT.reset(token)
+
+
 class _FunctionArgumentValidationError(TypeError):
     """An argument-validation failure raised before the function body starts."""
 
@@ -1167,7 +1181,6 @@ async def invoke(
         direct_argument_kwargs = (
             {key: value for key, value in kwargs.items() if key in parameter_names} if arguments is None else {}
         )
-        runtime_kwargs = dict(context.kwargs) if context is not None else {}
         unexpected_kwargs = {key: value for key, value in kwargs.items() if key not in direct_argument_kwargs}
         if unexpected_kwargs:
             unexpected_names = ", ".join(sorted(unexpected_kwargs))
@@ -1180,6 +1193,23 @@ async def invoke(
         if arguments is None and context is not None:
             arguments = context.arguments
 
+        security_materialization = _SECURITY_ARGUMENT_MATERIALIZATION_CONTEXT.get()
+        if security_materialization is not None:
+            active_context, materialize, origin_task = security_materialization
+            if (
+                asyncio.current_task() is origin_task
+                and active_context.function is self
+                and (context is None or context is active_context)
+            ):
+                direct_arguments = self._arguments_as_mapping(arguments)
+                if direct_arguments is not None:
+                    active_context.arguments = direct_arguments
+                active_context.metadata[_AUTO_ARGUMENT_PREPARATION_CONTEXT_KEY] = True
+                materialize(active_context)
+                arguments = active_context.arguments
+                context = active_context
+
+        runtime_kwargs = dict(context.kwargs) if context is not None else {}
         current_arguments = self._arguments_as_mapping(arguments)
         approval_visible_arguments = self._approval_visible_arguments(arguments, context)
         self._ensure_security_arguments_unchanged(context, current_arguments)
@@ -1211,7 +1241,8 @@ async def invoke(
         if not OBSERVABILITY_SETTINGS.ENABLED:
             logger.info(f"Function name: {self.name}")
             logger.debug(f"Function arguments: {observable_kwargs}")
-            result = await self._invoke_function(call_kwargs)
+            with _suspend_security_argument_materialization():
+                result = await self._invoke_function(call_kwargs)
             if skip_parsing:
                 logger.info(f"Function {self.name} succeeded.")
                 logger.debug(f"Function result: {type(result).__name__}")
@@ -1290,7 +1321,8 @@ async def invoke(
             start_time_stamp = perf_counter()
             end_time_stamp: float | None = None
             try:
-                result = await self._invoke_function(call_kwargs)
+                with _suspend_security_argument_materialization():
+                    result = await self._invoke_function(call_kwargs)
                 end_time_stamp = perf_counter()
             except Exception as exception:
                 end_time_stamp = perf_counter()