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]:
@@ -1252,6 +1254,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):
@@ -1478,6 +1485,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 +1505,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 +1522,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 +1536,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 +1546,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 +1560,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 +1573,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 +1586,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 +1601,15 @@ 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()
         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 +1620,23 @@ def _expand_variable_references_in_context(self, context: FunctionInvocationCont
                     depth=0,
                     active_variables=active_variables,
                     reference_count=reference_count,
+                    rewritten_paths=rewritten_paths,
                 ),
             )
+
+        rewritten_args: dict[str, set[int]] = {}
+        for path in rewritten_paths:
+            if not path or not isinstance(path[0], str):
+                continue
+            arg_name = path[0]
+            if arg_name not in rewritten_args:
+                rewritten_args[arg_name] = set()
+            if len(path) > 1 and isinstance(path[1], int):
+                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 +1796,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:
@@ -1839,6 +1880,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 +2287,31 @@ 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.
+
+    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 {}
+    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.