Python: [Feature]: Publish which arguments an invocation's variable expansion rewrote

python/packages/core/agent_framework/security.py

@@ -73,6 +73,7 @@
     "get_security_tools",
     "inspect_variable",
     "quarantined_llm",
+    "rewritten_arguments",
     "set_quarantine_client",
     "store_untrusted_content",
 ]
@@ -101,6 +102,7 @@
 # ``variable_ids`` list internally. Expanding their arguments would replace the ID
 # with the content and break the lookup.
 _VARIABLE_ID_CONSUMERS = frozenset({"inspect_variable", "quarantined_llm"})
+_REWRITTEN_ARGUMENT_INDICES_KEY = "_rewritten_argument_indices"
 
 
 def _get_additional_properties(obj: Any) -> dict[str, Any]:
@@ -109,6 +111,35 @@ def _get_additional_properties(obj: Any) -> dict[str, Any]:
     return cast(dict[str, Any], props) if isinstance(props, dict) else {}
 
 
+def _top_level_argument_value(context: FunctionInvocationContext, arg_name: str) -> tuple[Any, str | None]:
+    """Locate a top-level argument value in either context.arguments or context.kwargs.
+
+    Returns the value and a string indicating its source ('arguments' or 'kwargs'),
+    or (None, None) if not found.
+    """
+    args = cast(Any, context.arguments)
+    if isinstance(args, Mapping) and arg_name in args:
+        return cast(Any, args[arg_name]), "arguments"
+
+    kwargs = cast(Any, context.kwargs)
+    if isinstance(kwargs, Mapping) and arg_name in kwargs:
+        return cast(Any, kwargs[arg_name]), "kwargs"
+
+    return None, None
+
+
+def _top_level_argument_keys(context: FunctionInvocationContext) -> set[str]:
+    """Union of top level keys in context.arguments and context.kwargs."""
+    keys: set[str] = set()
+    args = cast(Any, context.arguments)
+    if isinstance(args, Mapping):
+        keys.update(cast(Mapping[str, Any], args).keys())
+    kwargs = cast(Any, context.kwargs)
+    if isinstance(kwargs, Mapping):
+        keys.update(cast(Mapping[str, Any], kwargs).keys())
+    return keys
+
+
 @dataclass(frozen=True, order=True, slots=True)
 class _Principal:
     """Canonical tenant/user identity used internally for comparisons."""
@@ -1252,6 +1283,11 @@ def list_variables(self) -> list[str]:
     default=None,
 )
 
+_current_context: ContextVar[FunctionInvocationContext | None] = ContextVar(
+    "agent_framework_current_security_context",
+    default=None,
+)
+
 
 @experimental(feature_id=ExperimentalFeature.FIDES)
 class LabelTrackingFunctionMiddleware(FunctionMiddleware, _SecurityScopeBinding):
@@ -1391,6 +1427,38 @@ def _update_context_label(self, new_content_label: ContentLabel) -> None:
                 self._context_label.confidentiality.value,
             )
 
+    @staticmethod
+    def _degrade_rewritten_arguments(context: FunctionInvocationContext) -> None:
+        """Degrade rewritten argument indices when validation mutates arguments.
+
+        For list arguments, marks every final list position as rewritten
+        (fail-closed) rather than using -1, which callers may misinterpret as
+        "non-list argument." A validator that moves hidden content from index 0
+        to index 1 must not let a tool treat index 1 as safe.
+
+        For non-list arguments, uses -1 as before.
+        If validation changes top-level keys (e.g., a Pydantic field alias
+        ``fileNames`` normalizes to the callable parameter ``files``), entries
+        keyed by the pre-validation name cannot be located under the final
+        callable argument name. In that case, fail closed under the final key
+        set: mark every final top-level argument as fully rewritten so a tool
+        looking up the actual parameter does not miss hidden content.
+        """
+        rewritten = context.metadata.get(_REWRITTEN_ARGUMENT_INDICES_KEY)
+        if not rewritten:
+            return
+        final_keys = _top_level_argument_keys(context)
+        pre_validation_keys = set(cast(dict[str, set[int]], rewritten).keys())
+        iterable_keys = final_keys if pre_validation_keys != final_keys else pre_validation_keys
+        degraded: dict[str, set[int]] = {}
+        for arg_name in iterable_keys:
+            value, _ = _top_level_argument_value(context, arg_name)
+            if isinstance(value, (list, tuple)):
+                degraded[arg_name] = set(range(len(cast(Sequence[Any], value))))
+            else:
+                degraded[arg_name] = {-1}
+        context.metadata[_REWRITTEN_ARGUMENT_INDICES_KEY] = degraded
+
     @staticmethod
     def _extract_primary_tool_content(expanded_content: Any, *, from_quarantined_llm: bool) -> Any:
         """Return the primary response from a proven quarantined LLM result.
@@ -1478,6 +1546,8 @@ def _resolve_string(
         depth: int,
         active_variables: set[str],
         reference_count: list[int],
+        rewritten_paths: set[tuple[str | int, ...]] | None = None,
+        current_path: tuple[str | int, ...] = (),
     ) -> Any:
         if not _EMBEDDED_VAR_REF_RE.search(value):
             return value
@@ -1496,6 +1566,8 @@ def _resolve_string(
                 return value
             if whole.group("bare"):
                 logger.warning(_BARE_REFERENCE_WARNING)
+            if rewritten_paths is not None:
+                rewritten_paths.add(current_path)
             return resolved
 
         def replace(match: re.Match[str]) -> str:
@@ -1511,6 +1583,8 @@ def replace(match: re.Match[str]) -> str:
                 return match.group(0)
             if match.group("bare"):
                 logger.warning(_BARE_REFERENCE_WARNING)
+            if rewritten_paths is not None:
+                rewritten_paths.add(current_path)
             return str(resolved)
 
         return _EMBEDDED_VAR_REF_RE.sub(replace, value)
@@ -1523,6 +1597,8 @@ def _resolve_value(
         depth: int,
         active_variables: set[str],
         reference_count: list[int],
+        rewritten_paths: set[tuple[str | int, ...]] | None = None,
+        current_path: tuple[str | int, ...] = (),
     ) -> Any:
         if isinstance(value, str):
             return self._resolve_string(
@@ -1531,6 +1607,8 @@ def _resolve_value(
                 depth=depth,
                 active_variables=active_variables,
                 reference_count=reference_count,
+                rewritten_paths=rewritten_paths,
+                current_path=current_path,
             )
         if isinstance(value, BaseModel):
             value = value.model_dump()
@@ -1543,6 +1621,8 @@ def _resolve_value(
                     depth=depth,
                     active_variables=active_variables,
                     reference_count=reference_count,
+                    rewritten_paths=rewritten_paths,
+                    current_path=(*current_path, key),
                 )
                 for key, item in value_dict.items()
             }
@@ -1554,8 +1634,10 @@ def _resolve_value(
                     depth=depth,
                     active_variables=active_variables,
                     reference_count=reference_count,
+                    rewritten_paths=rewritten_paths,
+                    current_path=(*current_path, index),
                 )
-                for item in cast(list[Any], value)
+                for index, item in enumerate(cast(list[Any], value))
             ]
         if isinstance(value, tuple):
             return tuple(
@@ -1565,8 +1647,10 @@ def _resolve_value(
                     depth=depth,
                     active_variables=active_variables,
                     reference_count=reference_count,
+                    rewritten_paths=rewritten_paths,
+                    current_path=(*current_path, index),
                 )
-                for item in cast(tuple[Any, ...], value)
+                for index, item in enumerate(cast(tuple[Any, ...], value))
             )
         return value
 
@@ -1578,13 +1662,16 @@ def _expand_variable_references_in_context(self, context: FunctionInvocationCont
         labels: list[ContentLabel] = []
         active_variables: set[str] = set()
         reference_count = [0]
+        rewritten_paths: set[tuple[str | int, ...]] = set()
+        kwargs_rewritten_paths: set[tuple[str | int, ...]] = set()
         if context.arguments:
             context.arguments = self._resolve_value(
                 context.arguments,
                 labels,
                 depth=0,
                 active_variables=active_variables,
                 reference_count=reference_count,
+                rewritten_paths=rewritten_paths,
             )
         if context.kwargs:
             context.kwargs = cast(
@@ -1595,8 +1682,48 @@ def _expand_variable_references_in_context(self, context: FunctionInvocationCont
                     depth=0,
                     active_variables=active_variables,
                     reference_count=reference_count,
+                    rewritten_paths=kwargs_rewritten_paths,
                 ),
             )
+
+        args_mapping = cast(Any, context.arguments)
+        arg_names_in_arguments: set[str] = (
+            set(cast(dict[str, Any], args_mapping).keys()) if isinstance(args_mapping, Mapping) else set()
+        )
+
+        rewritten_args: dict[str, set[int]] = {}
+        for path in rewritten_paths | kwargs_rewritten_paths:
+            if not path or not isinstance(path[0], str):
+                continue
+
+            arg_name = path[0]
+            if path not in rewritten_paths and arg_name in arg_names_in_arguments:
+                continue
+
+            if arg_name not in rewritten_args:
+                rewritten_args[arg_name] = set()
+
+            arg_value, arg_source = _top_level_argument_value(context, arg_name)
+
+            if arg_source is None:
+                rewritten_args[arg_name].add(-1)
+            elif len(path) == 1:
+                if isinstance(arg_value, (list, tuple)):
+                    resolved_list = cast(Sequence[Any], arg_value)
+                    rewritten_args[arg_name].update(range(len(resolved_list)))
+                else:
+                    rewritten_args[arg_name].add(-1)
+            elif (
+                len(path) > 1
+                and isinstance(path[1], int)
+                and not isinstance(path[1], bool)
+                and isinstance(arg_value, (list, tuple))
+            ):
+                rewritten_args[arg_name].add(path[1])
+            else:
+                rewritten_args[arg_name].add(-1)
+
+        context.metadata[_REWRITTEN_ARGUMENT_INDICES_KEY] = rewritten_args
         return labels
 
     def _get_input_labels(self, context: FunctionInvocationContext) -> list[ContentLabel]:
@@ -1756,6 +1883,7 @@ async def process(
         """Resolve hidden arguments, publish their labels, and label the result."""
         scope_token = self._activate_security_scope(context)
         middleware_token = _current_middleware.set(self)
+        context_token = _current_context.set(context)
         try:
             function_name = context.function.name
             if "original_arguments_for_messages" not in context.metadata:
@@ -1778,14 +1906,16 @@ async def process(
                 boundary="security policy",
             )
             if context.metadata.get(_AUTO_ARGUMENT_PREPARATION_CONTEXT_KEY) is True:
+                pre_validation_token = _argument_authority_token(context.arguments, boundary="security policy")
                 context.function._prepare_context_arguments(  # pyright: ignore[reportPrivateUsage]
                     context,
                     context.arguments,
                 )
-                context.metadata[_SECURITY_ARGUMENTS_SNAPSHOT_CONTEXT_KEY] = _argument_authority_token(
-                    context.arguments,
-                    boundary="security policy",
-                )
+                post_validation_token = _argument_authority_token(context.arguments, boundary="security policy")
+                context.metadata[_SECURITY_ARGUMENTS_SNAPSHOT_CONTEXT_KEY] = post_validation_token
+                if pre_validation_token != post_validation_token:
+                    self._degrade_rewritten_arguments(context)
+
             argument_labels = [*input_labels, *resolved_labels]
             argument_label = combine_labels(*argument_labels) if argument_labels else ContentLabel()
 
@@ -1839,6 +1969,7 @@ async def process(
                 return
             self._label_result(context, function_name, fallback_label)
         finally:
+            _current_context.reset(context_token)
             _current_middleware.reset(middleware_token)
             self._active_security_scope.reset(scope_token)
 
@@ -2245,6 +2376,54 @@ def get_current_middleware() -> LabelTrackingFunctionMiddleware | None:
     return _current_middleware.get()
 
 
+def rewritten_arguments(context: FunctionInvocationContext | None = None) -> dict[str, set[int]]:
+    """Get a mapping of argument names to the set of rewritten positions.
+
+    Returns a dictionary where keys are argument names and values are sets of
+    indices. For list arguments, the set contains the indices of the items
+    that were rewritten by variable expansion. For non-list arguments, the set
+    contains -1.
+
+    If the arguments were mutated after the indices were published (e.g., by a
+    Pydantic validator that reorders or filters a list), all final list
+    positions are marked as rewritten (fail-closed) because the original
+    positions are stale and a validator may have moved hidden content to any
+    position. For non-list arguments, -1 is used. This covers both the
+    auto-preparation path and the public direct-middleware path.
+
+    Args:
+        context: The function invocation context. If None, the context from
+            the current execution flow is used.
+
+    Returns:
+        A dictionary mapping argument names to sets of rewritten indices.
+    """
+    if context is None:
+        context = _current_context.get()
+    if context is None:
+        return {}
+    rewritten = context.metadata.get(_REWRITTEN_ARGUMENT_INDICES_KEY)
+    if rewritten is None:
+        return {}
+
+    indices_snapshot = context.metadata.get(_SECURITY_ARGUMENTS_SNAPSHOT_CONTEXT_KEY)
+    if indices_snapshot is not None:
+        current_snapshot = _argument_authority_token(context.arguments, boundary="security policy")
+        if current_snapshot != indices_snapshot:
+            rewritten_keys = set(cast(dict[str, set[int]], rewritten).keys())
+            final_keys = _top_level_argument_keys(context)
+            iterable_keys = final_keys if rewritten_keys != final_keys else rewritten_keys
+            result: dict[str, set[int]] = {}
+            for arg_name in iterable_keys:
+                value, _ = _top_level_argument_value(context, arg_name)
+                if isinstance(value, (list, tuple)):
+                    result[arg_name] = set(range(len(cast(Sequence[Any], value))))
+                else:
+                    result[arg_name] = {-1}
+            return result
+    return {k: set(v) for k, v in cast(dict[str, set[int]], rewritten).items()}
+
+
 @dataclass(frozen=True, slots=True)
 class _PendingPolicyApproval:
     """Immutable binding record for a pending policy-violation approval.

python/packages/core/tests/core/test_harness_tool_approval.py

@@ -4,6 +4,7 @@
 
 import asyncio
 import json
+import logging
 import warnings
 from collections.abc import Awaitable, Callable, MutableSequence
 from contextlib import nullcontext
@@ -1133,6 +1134,7 @@ async def test_approval_resume_returns_result_without_mutating_inputs(
     conflicting_duplicate: bool,
 ) -> None:
     """Approval resume should return its terminal result without changing caller-owned messages."""
+    caplog.set_level(logging.WARNING)
     calls = 0
 
     @tool(name="guarded_tool", approval_mode="always_require")