From 8c58607fea25f8bc66fd6ffedcbec486d732b470 Mon Sep 17 00:00:00 2001 From: pratikwayase Date: Fri, 18 Sep 2026 15:10:02 +0530 Subject: [PATCH] feat(security): publish rewritten argument positions from variable expansion --- .../packages/core/agent_framework/security.py | 71 +++++- python/packages/core/tests/test_security.py | 214 ++++++++++++++++++ 2 files changed, 283 insertions(+), 2 deletions(-) diff --git a/python/packages/core/agent_framework/security.py b/python/packages/core/agent_framework/security.py index b7d87a5f43c..496bc6fac8d 100644 --- a/python/packages/core/agent_framework/security.py +++ b/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,6 +1601,7 @@ 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, @@ -1585,6 +1609,7 @@ def _expand_variable_references_in_context(self, context: FunctionInvocationCont 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. diff --git a/python/packages/core/tests/test_security.py b/python/packages/core/tests/test_security.py index 7f371b61a1d..780ef6659b1 100644 --- a/python/packages/core/tests/test_security.py +++ b/python/packages/core/tests/test_security.py @@ -46,6 +46,7 @@ VariableReferenceContent, combine_labels, get_current_middleware, + rewritten_arguments, store_untrusted_content, ) @@ -7569,3 +7570,216 @@ async def execute(_context: FunctionInvocationContext) -> list[Content]: assert executed is True assert replay.metadata["user_approved_violation"] is True + + +@pytest.mark.asyncio +async def test_rewritten_arguments_no_rewrites(): + """Verify normal/non-expanded path returns empty dict.""" + tracker = LabelTrackingFunctionMiddleware() + captured = {} + + async def my_tool(files: list[str]): + captured["rewritten"] = rewritten_arguments() + return "ok" + + tool = FunctionTool( + name="my_tool", + func=my_tool, + additional_properties={"accepts_untrusted": True}, + ) + + context = FunctionInvocationContext( + function=tool, + arguments={"files": ["one.txt", "two.txt"]}, + ) + + async def call_next(): + await tool.invoke(arguments=context.arguments) + + await tracker.process(context, call_next) + + assert captured["rewritten"] == {} + + +@pytest.mark.asyncio +async def test_rewritten_arguments_explicit_context(): + """Verify the explicit context API works inside a tool.""" + tracker = LabelTrackingFunctionMiddleware() + store = tracker.get_variable_store() + var_id = store.store("payload", ContentLabel(integrity=IntegrityLabel.UNTRUSTED)) + + captured = {} + + async def my_tool(files: list[str], context: FunctionInvocationContext): + captured["explicit"] = rewritten_arguments(context) + captured["implicit"] = rewritten_arguments() + return "ok" + + tool = FunctionTool(name="my_tool", func=my_tool, additional_properties={"accepts_untrusted": True}) + context = FunctionInvocationContext(function=tool, arguments={"files": [f"[{var_id}]", "safe.txt"]}) + + async def call_next(): + await tool.func(files=context.arguments["files"], context=context) + + await tracker.process(context, call_next) + + assert captured["explicit"] == {"files": {0}} + assert captured["implicit"] == {"files": {0}} + + +@pytest.mark.asyncio +async def test_rewritten_arguments_multiple_args(): + """Verify tracking works across multiple top-level arguments.""" + tracker = LabelTrackingFunctionMiddleware() + store = tracker.get_variable_store() + var_id1 = store.store("file_content", ContentLabel(integrity=IntegrityLabel.UNTRUSTED)) + var_id2 = store.store("msg_content", ContentLabel(integrity=IntegrityLabel.UNTRUSTED)) + + captured = {} + + async def my_tool(files: list[str], message: str): + captured["rewritten"] = rewritten_arguments() + return "ok" + + tool = FunctionTool(name="my_tool", func=my_tool, additional_properties={"accepts_untrusted": True}) + context = FunctionInvocationContext( + function=tool, arguments={"files": [f"[{var_id1}]", "safe.txt"], "message": f"[{var_id2}]"} + ) + + async def call_next(): + await tool.invoke(arguments=context.arguments) + + await tracker.process(context, call_next) + + assert captured["rewritten"] == {"files": {0}, "message": {-1}} + + +@pytest.mark.asyncio +async def test_rewritten_arguments_duplicate_equal_values(): + """Test that duplicate/equal final values are tracked per position.""" + tracker = LabelTrackingFunctionMiddleware() + store = tracker.get_variable_store() + + var_id1 = store.store("same_string", ContentLabel(integrity=IntegrityLabel.UNTRUSTED)) + var_id2 = store.store("same_string", ContentLabel(integrity=IntegrityLabel.UNTRUSTED)) + + captured = {} + + async def my_tool(files: list[str]): + captured["rewritten"] = rewritten_arguments() + captured["received"] = files + return "ok" + + tool = FunctionTool(name="my_tool", func=my_tool, additional_properties={"accepts_untrusted": True}) + context = FunctionInvocationContext( + function=tool, arguments={"files": [f"[{var_id1}]", f"[{var_id2}]", "normal.txt"]} + ) + + async def call_next(): + await tool.invoke(arguments=context.arguments) + + await tracker.process(context, call_next) + + assert captured["received"] == ["same_string", "same_string", "normal.txt"] + assert captured["rewritten"] == {"files": {0, 1}} + + +@pytest.mark.asyncio +async def test_rewritten_arguments_multiple_list_positions(): + """Test multiple list positions alongside untouched positions.""" + tracker = LabelTrackingFunctionMiddleware() + store = tracker.get_variable_store() + var_id = store.store("payload", ContentLabel(integrity=IntegrityLabel.UNTRUSTED)) + + captured = {} + + async def my_tool(files: list[str]): + captured["rewritten"] = rewritten_arguments() + return "ok" + + tool = FunctionTool(name="my_tool", func=my_tool, additional_properties={"accepts_untrusted": True}) + context = FunctionInvocationContext( + function=tool, arguments={"files": [f"[{var_id}]", "untouched.txt", f"[{var_id}]"]} + ) + + async def call_next(): + await tool.invoke(arguments=context.arguments) + + await tracker.process(context, call_next) + + assert captured["rewritten"] == {"files": {0, 2}} + + +@pytest.mark.asyncio +async def test_rewritten_arguments_scalar(): + """Test scalar (non-list) arguments.""" + tracker = LabelTrackingFunctionMiddleware() + store = tracker.get_variable_store() + var_id = store.store("payload", ContentLabel(integrity=IntegrityLabel.UNTRUSTED)) + + captured = {} + + async def my_tool(text: str): + captured["rewritten"] = rewritten_arguments() + return "ok" + + tool = FunctionTool(name="my_tool", func=my_tool, additional_properties={"accepts_untrusted": True}) + context = FunctionInvocationContext(function=tool, arguments={"text": f"[{var_id}]"}) + + async def call_next(): + await tool.invoke(arguments=context.arguments) + + await tracker.process(context, call_next) + + assert captured["rewritten"] == {"text": {-1}} + + +@pytest.mark.asyncio +async def test_rewritten_arguments_nested_dict_semantics(): + """Test that nested rewrites are reported against the top-level argument.""" + tracker = LabelTrackingFunctionMiddleware() + store = tracker.get_variable_store() + var_id = store.store("payload", ContentLabel(integrity=IntegrityLabel.UNTRUSTED)) + + captured = {} + + async def my_tool(config: dict): + captured["rewritten"] = rewritten_arguments() + return "ok" + + tool = FunctionTool(name="my_tool", func=my_tool, additional_properties={"accepts_untrusted": True}) + context = FunctionInvocationContext(function=tool, arguments={"config": {"path": f"[{var_id}]", "safe": "txt"}}) + + async def call_next(): + await tool.invoke(arguments=context.arguments) + + await tracker.process(context, call_next) + + assert captured["rewritten"] == {"config": {-1}} + + +@pytest.mark.asyncio +async def test_rewritten_arguments_asyncio_to_thread(): + """Verify async/thread/context behavior with asyncio.to_thread.""" + tracker = LabelTrackingFunctionMiddleware() + store = tracker.get_variable_store() + var_id = store.store("thread_content", ContentLabel(integrity=IntegrityLabel.UNTRUSTED)) + + captured = {} + + async def threaded_tool(files: list[str]): + def worker(): + return rewritten_arguments() + + captured["rewritten"] = await asyncio.to_thread(worker) + return "ok" + + tool = FunctionTool(name="threaded_tool", func=threaded_tool, additional_properties={"accepts_untrusted": True}) + context = FunctionInvocationContext(function=tool, arguments={"files": [f"[{var_id}]", "safe.txt"]}) + + async def call_next(): + await tool.invoke(arguments=context.arguments) + + await tracker.process(context, call_next) + + assert captured["rewritten"] == {"files": {0}}