From 7dbd951b3efce2745e1ca14e02b7ebdd51f7f544 Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Mon, 7 Sep 2026 20:41:29 +0200 Subject: [PATCH 1/5] Python: preserve confidentiality through FIDES security tools Keep resolved-input confidentiality on transformed and embedded-label results, publish authoritative inspect/quarantine labels, and avoid unwrapping ordinary response-shaped JSON. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../packages/core/agent_framework/security.py | 104 +++++++-- python/packages/core/tests/test_security.py | 203 +++++++++++++++++- 2 files changed, 289 insertions(+), 18 deletions(-) diff --git a/python/packages/core/agent_framework/security.py b/python/packages/core/agent_framework/security.py index 1a46bd3687..55ffa5feac 100644 --- a/python/packages/core/agent_framework/security.py +++ b/python/packages/core/agent_framework/security.py @@ -1262,14 +1262,25 @@ def _update_context_label(self, new_content_label: ContentLabel) -> None: def _extract_primary_tool_content(expanded_content: Any) -> Any: """Return the tool-visible content for an expanded variable payload. - Some hidden results are stored as rich payloads containing fields such as - ``response``, ``security_label``, and ``metadata``. Tool arguments should - only receive the primary content they would normally have seen without - variable indirection. + Hidden ``quarantined_llm`` results are stored as rich payloads containing + the response plus security metadata. Tool arguments should receive only the + primary response they would have seen without variable indirection. Ordinary + mappings or JSON text that happen to contain a ``response`` key remain intact. """ + + def is_quarantine_payload(payload: dict[str, Any]) -> bool: + return ( + payload.get("quarantined") is True + and "response" in payload + and isinstance(payload.get("security_label"), MutableMapping) + and isinstance(payload.get("metadata"), MutableMapping) + and isinstance(payload.get("variables_processed"), list) + and isinstance(payload.get("content_summary"), list) + ) + if isinstance(expanded_content, dict): content_map = cast(dict[str, Any], expanded_content) - if "response" in content_map: + if is_quarantine_payload(content_map): return content_map["response"] return content_map @@ -1280,7 +1291,7 @@ def _extract_primary_tool_content(expanded_content: Any) -> Any: parsed = json.loads(stripped) if isinstance(parsed, dict): parsed_map = cast(dict[str, Any], parsed) - if "response" in parsed_map: + if is_quarantine_payload(parsed_map): return parsed_map["response"] return expanded_content @@ -1516,23 +1527,37 @@ async def process( declared_source_integrity = self._get_source_integrity(context) confidentiality = self._get_function_confidentiality(context) + # Expand hidden references before execution and retain their stored labels. + resolved_labels = self._expand_variable_references_in_context(context) + argument_labels = [*input_labels, *resolved_labels] + argument_label = combine_labels(*argument_labels) if argument_labels else ContentLabel() + + # Integrity may be declared by the source, but a transformer cannot + # implicitly declassify data derived from its inputs. + result_confidentiality = combine_labels( + ContentLabel(confidentiality=confidentiality), argument_label + ).confidentiality + + # Step 3: Build tiered fallback_label + # This label is used for result items that have NO embedded labels. + # Priority: source_integrity declaration (tier 2) > input labels join (tier 3) if declared_source_integrity is not None: fallback_label = ContentLabel( integrity=declared_source_integrity, - confidentiality=confidentiality, + confidentiality=result_confidentiality, metadata={"source": "source_integrity", "function_name": function_name}, ) elif input_labels: combined = combine_labels(*input_labels) fallback_label = ContentLabel( integrity=combined.integrity, - confidentiality=confidentiality, + confidentiality=result_confidentiality, metadata={"source": "input_labels_join", "function_name": function_name}, ) else: fallback_label = ContentLabel( integrity=self.default_integrity, - confidentiality=confidentiality, + confidentiality=result_confidentiality, metadata={"source": "default", "function_name": function_name}, ) @@ -1724,11 +1749,18 @@ def _extract_content_label( """ additional_props = _get_additional_properties(item) - # Check for standard security_label + # Embedded labels remain authoritative for integrity, but cannot + # declassify data inherited from the tool's inputs. label_data = additional_props.get("security_label") if label_data and isinstance(label_data, dict): try: - return ContentLabel.from_dict(cast(dict[str, Any], label_data)) + embedded_label = ContentLabel.from_dict(cast(dict[str, Any], label_data)) + combined_label = combine_labels(fallback_label, embedded_label) + return ContentLabel( + integrity=embedded_label.integrity, + confidentiality=combined_label.confidentiality, + metadata=combined_label.metadata, + ) except Exception as e: logger.warning(f"Failed to parse security_label from Content: {e}") @@ -2874,6 +2906,34 @@ def get_quarantine_client() -> SupportsChatGetResponse | None: """ +def _quarantined_llm_result_parser(result: Any) -> list[Content]: + """Publish quarantine output integrity and combined input confidentiality.""" + contents = FunctionTool.parse_result(result) + if not contents or not isinstance(result, dict): + return contents + + label_data = cast(dict[str, Any], result).get("security_label") + if not isinstance(label_data, MutableMapping): + return contents + + try: + parsed_label = ContentLabel.from_dict(cast(MutableMapping[str, Any], label_data)) + except (TypeError, ValueError) as exc: + logger.warning("Failed to parse quarantined_llm result label: %s", exc) + return contents + + quarantine_label = ContentLabel( + integrity=IntegrityLabel.UNTRUSTED, + confidentiality=parsed_label.confidentiality, + metadata=parsed_label.metadata, + ) + first = contents[0] + props = first.additional_properties or {} + props["security_label"] = quarantine_label.to_dict() + first.additional_properties = props + return contents + + @tool( description=( "Make an isolated LLM call with labeled data in a quarantined context. " @@ -2883,6 +2943,7 @@ def get_quarantine_client() -> SupportsChatGetResponse | None: "You can pass variable_ids directly to reference hidden content from VariableReferenceContent objects. " "UNTRUSTED results are automatically hidden by the middleware." ), + result_parser=_quarantined_llm_result_parser, additional_properties={ "confidentiality": "private", "accepts_untrusted": True, @@ -3122,6 +3183,9 @@ class InspectVariableInput(BaseModel): reason: str | None = Field(default=None, description="Reason for inspecting this variable (for audit purposes)") +_INSPECT_VARIABLE_CONFIDENTIALITY = ConfidentialityLabel.PRIVATE + + def _inspect_variable_result_parser(result: Any) -> list[Content]: """Parse ``inspect_variable``'s dict result while preserving its security label. @@ -3134,13 +3198,21 @@ def _inspect_variable_result_parser(result: Any) -> list[Content]: downgrading the real label. This parser stamps the inspected label back onto the produced Content so - ``LabelTrackingFunctionMiddleware`` propagates it faithfully. The error path - (``security_label`` is ``None``) is left unstamped so it safely falls back to - the tool's default label. + ``LabelTrackingFunctionMiddleware`` propagates it faithfully. Missing-variable + errors are labeled as trusted tool output so probing an absent or foreign id + does not falsely taint cumulative integrity. """ contents = FunctionTool.parse_result(result) + if not contents: + return contents + label = cast(dict[str, Any], result).get("security_label") if isinstance(result, dict) else None - if label and contents: + if label is None and isinstance(result, dict) and "error" in cast(dict[str, Any], result): + label = ContentLabel( + integrity=IntegrityLabel.TRUSTED, + confidentiality=_INSPECT_VARIABLE_CONFIDENTIALITY, + ).to_dict() + if label: first = contents[0] props = first.additional_properties or {} props["security_label"] = label @@ -3158,7 +3230,7 @@ def _inspect_variable_result_parser(result: Any) -> list[Content]: approval_mode="never_require", result_parser=_inspect_variable_result_parser, additional_properties={ - "confidentiality": "private", + "confidentiality": _INSPECT_VARIABLE_CONFIDENTIALITY.value, # No source_integrity declared: output inherits the label of the # inspected content via Tier 3. The variable store is just a # container — the data inside it is untrusted external content. diff --git a/python/packages/core/tests/test_security.py b/python/packages/core/tests/test_security.py index 903d64e0f2..fb64e96d98 100644 --- a/python/packages/core/tests/test_security.py +++ b/python/packages/core/tests/test_security.py @@ -536,6 +536,39 @@ async def next_fn() -> None: await middleware.process(context, next_fn) + async def test_json_response_object_is_not_mistaken_for_quarantine_payload(self, middleware) -> None: + """Ordinary JSON objects containing ``response`` survive expansion intact.""" + + class MessageArgs(BaseModel): + summary: str + + async def send_message(summary: str) -> str: + return summary + + message_tool = FunctionTool( + fn=send_message, + name="SendMessagetoSelf", + description="Send message", + args_schema=MessageArgs, + ) + stored_payload = json.dumps({"response": "keep", "other": "also keep"}) + variable_id = middleware.get_variable_store().store( + stored_payload, + ContentLabel(integrity=IntegrityLabel.UNTRUSTED), + ) + context = FunctionInvocationContext( + function=message_tool, + arguments=MessageArgs(summary=f"[{variable_id}]"), + ) + + async def next_fn() -> None: + current_args = context.arguments + assert isinstance(current_args, dict) + assert current_args["summary"] == stored_payload + context.result = [Content.from_text("sent")] + + await middleware.process(context, next_fn) + class TestPolicyEnforcementMiddleware: """Tests for PolicyEnforcementFunctionMiddleware.""" @@ -1750,7 +1783,9 @@ async def next_fn(): await middleware_no_auto_hide.process(context, next_fn) result_label = context.metadata["result_label"] + assert result_label.integrity == IntegrityLabel.UNTRUSTED assert result_label.confidentiality == ConfidentialityLabel.USER_IDENTITY + assert middleware_no_auto_hide.get_context_label().integrity == IntegrityLabel.UNTRUSTED assert middleware_no_auto_hide.get_context_label().confidentiality == ConfidentialityLabel.USER_IDENTITY @pytest.mark.asyncio @@ -1780,7 +1815,7 @@ async def next_fn(): @pytest.mark.asyncio async def test_inspect_variable_missing_var_does_not_crash(self, middleware_no_auto_hide): - """A missing variable id returns an error result and falls back safely.""" + """A missing variable id returns a trusted tool-generated error.""" from agent_framework.security import get_security_tools inspect_tool = next(tool for tool in get_security_tools() if tool.name == "inspect_variable") @@ -1798,8 +1833,37 @@ async def next_fn(): payload = json.loads(context.result[0].text) assert payload["security_label"] is None assert "error" in payload - # No embedded label -> falls back to the tool's default confidentiality. + assert context.metadata["result_label"].integrity == IntegrityLabel.TRUSTED + assert context.metadata["result_label"].confidentiality == ConfidentialityLabel.PRIVATE + assert middleware_no_auto_hide.get_context_label().integrity == IntegrityLabel.TRUSTED + + async def test_inspect_variable_foreign_var_does_not_taint_integrity(self) -> None: + """An id owned by another store returns a trusted tool-generated error.""" + from agent_framework.security import get_security_tools + + owner = LabelTrackingFunctionMiddleware() + foreign_id = owner.get_variable_store().store( + "foreign secret", + ContentLabel(integrity=IntegrityLabel.UNTRUSTED, confidentiality=ConfidentialityLabel.USER_IDENTITY), + ) + middleware = LabelTrackingFunctionMiddleware(auto_hide_untrusted=False) + inspect_tool = next(tool for tool in get_security_tools() if tool.name == "inspect_variable") + context = FunctionInvocationContext( + function=inspect_tool, + arguments={"variable_id": foreign_id, "reason": "foreign id"}, + ) + + async def next_fn() -> None: + context.result = await inspect_tool.invoke(arguments=context.arguments, context=context) + + await middleware.process(context, next_fn) + + payload = json.loads(context.result[0].text) + assert payload["security_label"] is None + assert "error" in payload + assert context.metadata["result_label"].integrity == IntegrityLabel.TRUSTED assert context.metadata["result_label"].confidentiality == ConfidentialityLabel.PRIVATE + assert middleware.get_context_label().integrity == IntegrityLabel.TRUSTED @pytest.mark.asyncio async def test_multiple_calls_accumulate_variables(self, middleware_auto_hide, mock_function): @@ -2869,6 +2933,36 @@ class TestQuarantinedLLM: via source_integrity="untrusted", not by quarantined_llm itself. """ + async def test_quarantined_llm_publishes_combined_confidentiality(self) -> None: + """Quarantine output is UNTRUSTED at the highest input confidentiality.""" + middleware = LabelTrackingFunctionMiddleware() + variable_id = middleware.get_variable_store().store( + "identity secret", + ContentLabel( + integrity=IntegrityLabel.TRUSTED, + confidentiality=ConfidentialityLabel.USER_IDENTITY, + ), + ) + quarantine_tool = next(tool for tool in middleware.get_security_tools() if tool.name == "quarantined_llm") + context = FunctionInvocationContext( + function=quarantine_tool, + arguments={"prompt": "Summarize", "variable_ids": [variable_id]}, + ) + + async def next_fn() -> None: + context.result = await quarantine_tool.invoke(arguments=context.arguments, context=context) + + await middleware.process(context, next_fn) + + assert context.metadata["result_label"].integrity == IntegrityLabel.UNTRUSTED + assert context.metadata["result_label"].confidentiality == ConfidentialityLabel.USER_IDENTITY + assert middleware.get_context_label().integrity == IntegrityLabel.TRUSTED + assert middleware.get_context_label().confidentiality == ConfidentialityLabel.USER_IDENTITY + hidden_reference = json.loads(context.result[0].text) + _, hidden_label = middleware.get_variable_store().retrieve(hidden_reference["variable_id"]) + assert hidden_label.integrity == IntegrityLabel.UNTRUSTED + assert hidden_label.confidentiality == ConfidentialityLabel.USER_IDENTITY + @pytest.mark.asyncio async def test_quarantined_llm_returns_response(self): """Test that quarantined_llm returns a plain response dict.""" @@ -3357,6 +3451,31 @@ async def next_fn(): parsed = json.loads(item.text) # type: ignore[arg-type] # pyrefly: ignore[bad-argument-type] assert parsed.get("type") == "variable_reference" + async def test_fully_hidden_result_updates_confidentiality_without_integrity_taint( + self, middleware, mock_function + ) -> None: + """Hidden content affects cumulative confidentiality but not integrity.""" + context = FunctionInvocationContext(function=mock_function, arguments=mock_function.args_schema()) + + async def next_fn() -> None: + context.result = [ + Content.from_text( + "hidden identity data", + additional_properties={ + "security_label": { + "integrity": "untrusted", + "confidentiality": "user_identity", + } + }, + ) + ] + + await middleware.process(context, next_fn) + + assert context.result[0].additional_properties["_variable_reference"] is True + assert middleware.get_context_label().integrity == IntegrityLabel.TRUSTED + assert middleware.get_context_label().confidentiality == ConfidentialityLabel.USER_IDENTITY + @pytest.mark.asyncio async def test_items_without_labels_use_fallback(self, middleware, mock_function): """Test that items without embedded labels use the fallback (call) label.""" @@ -3576,6 +3695,86 @@ async def next_fn(): # Tier 2 (source_integrity=trusted) wins over tier 3 (untrusted input) assert label.integrity == IntegrityLabel.TRUSTED + @pytest.mark.parametrize( + "confidentiality", + [ConfidentialityLabel.PRIVATE, ConfidentialityLabel.USER_IDENTITY], + ) + async def test_hidden_input_confidentiality_propagates_to_transform_result( + self, + middleware: LabelTrackingFunctionMiddleware, + confidentiality: ConfidentialityLabel, + ) -> None: + """A trusted transformer cannot implicitly declassify hidden input.""" + + class Args(BaseModel): + value: str + + async def transform(value: str) -> str: + return value.upper() + + function = FunctionTool( + fn=transform, + name="trusted_transformer", + description="Transform hidden input", + args_schema=Args, + additional_properties={"source_integrity": "trusted"}, + ) + variable_id = middleware.get_variable_store().store( + "secret", + ContentLabel(integrity=IntegrityLabel.UNTRUSTED, confidentiality=confidentiality), + ) + context = FunctionInvocationContext( + function=function, + arguments=Args(value=f"[{variable_id}]"), + ) + + async def next_fn() -> None: + current_args = context.arguments + assert isinstance(current_args, dict) + context.result = [Content.from_text(current_args["value"].upper())] + + await middleware.process(context, next_fn) + + assert context.metadata["result_label"].integrity == IntegrityLabel.TRUSTED + assert context.metadata["result_label"].confidentiality == confidentiality + assert middleware.get_context_label().confidentiality == confidentiality + + async def test_embedded_public_label_cannot_declassify_hidden_input(self, middleware) -> None: + """Embedded labels choose integrity without lowering input confidentiality.""" + + class Args(BaseModel): + value: str + + async def transform(value: str) -> str: + return value + + function = FunctionTool( + fn=transform, + name="embedded_label_transformer", + description="Transform hidden input", + args_schema=Args, + additional_properties={"source_integrity": "untrusted"}, + ) + variable_id = middleware.get_variable_store().store( + "private payload", + ContentLabel(integrity=IntegrityLabel.UNTRUSTED, confidentiality=ConfidentialityLabel.PRIVATE), + ) + context = FunctionInvocationContext(function=function, arguments=Args(value=f"[{variable_id}]")) + + async def next_fn() -> None: + context.result = [ + Content.from_text( + "transformed", + additional_properties={"security_label": {"integrity": "trusted", "confidentiality": "public"}}, + ) + ] + + await middleware.process(context, next_fn) + + assert context.metadata["result_label"].integrity == IntegrityLabel.TRUSTED + assert context.metadata["result_label"].confidentiality == ConfidentialityLabel.PRIVATE + assert middleware.get_context_label().confidentiality == ConfidentialityLabel.PRIVATE + @pytest.mark.asyncio async def test_embedded_labels_override_source_integrity(self, middleware): """Test that embedded labels (tier 1) override source_integrity (tier 2). From c894ae48c4c5ba711fd9b4ca3bc8f9fe93e8981f Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Tue, 8 Sep 2026 09:06:13 +0200 Subject: [PATCH 2/5] Python: extract quarantine payload predicate Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../packages/core/agent_framework/security.py | 27 ++++++++++--------- 1 file changed, 14 insertions(+), 13 deletions(-) diff --git a/python/packages/core/agent_framework/security.py b/python/packages/core/agent_framework/security.py index 55ffa5feac..456d70410e 100644 --- a/python/packages/core/agent_framework/security.py +++ b/python/packages/core/agent_framework/security.py @@ -90,6 +90,18 @@ def _get_additional_properties(obj: Any) -> dict[str, Any]: return cast(dict[str, Any], props) if isinstance(props, dict) else {} +def _is_quarantine_payload(payload: dict[str, Any]) -> bool: + """Return whether a mapping has the internal quarantined LLM result shape.""" + return ( + payload.get("quarantined") is True + and "response" in payload + and isinstance(payload.get("security_label"), MutableMapping) + and isinstance(payload.get("metadata"), MutableMapping) + and isinstance(payload.get("variables_processed"), list) + and isinstance(payload.get("content_summary"), list) + ) + + # ============================================================================= # Core Security Primitives # ============================================================================= @@ -1267,20 +1279,9 @@ def _extract_primary_tool_content(expanded_content: Any) -> Any: primary response they would have seen without variable indirection. Ordinary mappings or JSON text that happen to contain a ``response`` key remain intact. """ - - def is_quarantine_payload(payload: dict[str, Any]) -> bool: - return ( - payload.get("quarantined") is True - and "response" in payload - and isinstance(payload.get("security_label"), MutableMapping) - and isinstance(payload.get("metadata"), MutableMapping) - and isinstance(payload.get("variables_processed"), list) - and isinstance(payload.get("content_summary"), list) - ) - if isinstance(expanded_content, dict): content_map = cast(dict[str, Any], expanded_content) - if is_quarantine_payload(content_map): + if _is_quarantine_payload(content_map): return content_map["response"] return content_map @@ -1291,7 +1292,7 @@ def is_quarantine_payload(payload: dict[str, Any]) -> bool: parsed = json.loads(stripped) if isinstance(parsed, dict): parsed_map = cast(dict[str, Any], parsed) - if is_quarantine_payload(parsed_map): + if _is_quarantine_payload(parsed_map): return parsed_map["response"] return expanded_content From c730173f60098c6aa44f4829bcc73230e593a5d7 Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Tue, 8 Sep 2026 09:26:02 +0200 Subject: [PATCH 3/5] Python: integrate confidentiality propagation with argument labels Resolve the stacked PR2/PR3 boundary so hidden arguments are expanded once, retain their combined label for policy, and preserve their confidentiality on transformed results.\n\nCo-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- python/packages/core/agent_framework/security.py | 6 ++---- python/packages/core/tests/test_security.py | 6 +++--- 2 files changed, 5 insertions(+), 7 deletions(-) diff --git a/python/packages/core/agent_framework/security.py b/python/packages/core/agent_framework/security.py index 456d70410e..904be0c2df 100644 --- a/python/packages/core/agent_framework/security.py +++ b/python/packages/core/agent_framework/security.py @@ -1548,8 +1548,8 @@ async def process( confidentiality=result_confidentiality, metadata={"source": "source_integrity", "function_name": function_name}, ) - elif input_labels: - combined = combine_labels(*input_labels) + elif argument_labels: + combined = combine_labels(*argument_labels) fallback_label = ContentLabel( integrity=combined.integrity, confidentiality=result_confidentiality, @@ -1562,8 +1562,6 @@ async def process( metadata={"source": "default", "function_name": function_name}, ) - resolved_labels = self._expand_variable_references_in_context(context) - argument_label = combine_labels(*resolved_labels) if resolved_labels else ContentLabel() context_label = self._context_label context.metadata["context_label"] = context_label context.metadata["argument_label"] = argument_label diff --git a/python/packages/core/tests/test_security.py b/python/packages/core/tests/test_security.py index fb64e96d98..54524159f5 100644 --- a/python/packages/core/tests/test_security.py +++ b/python/packages/core/tests/test_security.py @@ -4801,7 +4801,7 @@ async def test_private_hidden_argument_is_blocked_from_public_sink(self) -> None assert context.metadata["effective_invocation_label"].confidentiality == ConfidentialityLabel.PRIVATE assert policy.get_audit_log()[-1]["subtype"] == "max_allowed_confidentiality" - async def test_argument_labels_do_not_rewrite_result_labels(self) -> None: + async def test_argument_labels_preserve_result_confidentiality(self) -> None: tracker = LabelTrackingFunctionMiddleware() policy = PolicyEnforcementFunctionMiddleware() variable_id = tracker.get_variable_store().store( @@ -4822,8 +4822,8 @@ async def execute(_context: FunctionInvocationContext) -> list[Content]: await FunctionMiddlewarePipeline(tracker, policy).execute(context, execute) assert context.metadata["result_label"].integrity == IntegrityLabel.TRUSTED - assert context.metadata["result_label"].confidentiality == ConfidentialityLabel.PUBLIC - assert tracker.get_context_label().confidentiality == ConfidentialityLabel.PUBLIC + assert context.metadata["result_label"].confidentiality == ConfidentialityLabel.PRIVATE + assert tracker.get_context_label().confidentiality == ConfidentialityLabel.PRIVATE async def test_policy_approval_allows_exact_resolved_invocation(self) -> None: tracker = LabelTrackingFunctionMiddleware() From 384489e14ab5d86f8645d2c8818cc162d7a59996 Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Wed, 9 Sep 2026 10:10:34 +0200 Subject: [PATCH 4/5] Python: harden FIDES result label authority Use middleware provenance for quarantine unwrapping, require explicit label fields, fail unknown quarantine inputs closed, and scope trusted inspect errors to clean parser-marked invocations.\n\nCo-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../packages/core/agent_framework/security.py | 164 ++++++++----- python/packages/core/tests/test_security.py | 227 +++++++++++++++++- 2 files changed, 331 insertions(+), 60 deletions(-) diff --git a/python/packages/core/agent_framework/security.py b/python/packages/core/agent_framework/security.py index 904be0c2df..29f8df7cda 100644 --- a/python/packages/core/agent_framework/security.py +++ b/python/packages/core/agent_framework/security.py @@ -75,6 +75,9 @@ _WHOLE_VAR_REF_RE = re.compile(rf"\s*(?:\[\s*(?P{_BRACKETED_VAR_ID})\s*\]|(?P{_BARE_VAR_ID}))\s*") _BARE_REFERENCE_WARNING = "Expanded a bare variable reference in a tool argument; models should use [var_] instead." _UNRESOLVED = object() +_AUTHORITATIVE_CONFIDENTIALITY = "_security_label_authoritative_confidentiality" +_INSPECT_VARIABLE_ERROR = "_inspect_variable_error" +_INTERNAL_RESULT_MARKER = object() # Tools that consume variable IDs literally (as opaque references) and therefore # must NOT have ``var_xxx`` arguments expanded to stored content before execution. @@ -90,15 +93,21 @@ def _get_additional_properties(obj: Any) -> dict[str, Any]: return cast(dict[str, Any], props) if isinstance(props, dict) else {} -def _is_quarantine_payload(payload: dict[str, Any]) -> bool: - """Return whether a mapping has the internal quarantined LLM result shape.""" - return ( - payload.get("quarantined") is True - and "response" in payload - and isinstance(payload.get("security_label"), MutableMapping) - and isinstance(payload.get("metadata"), MutableMapping) - and isinstance(payload.get("variables_processed"), list) - and isinstance(payload.get("content_summary"), list) +def _parse_content_label(label_data: MutableMapping[str, Any], *, source: str) -> ContentLabel: + """Parse explicit label fields while ignoring malformed optional metadata.""" + integrity = label_data.get("integrity") + confidentiality = label_data.get("confidentiality") + if not isinstance(integrity, str) or not isinstance(confidentiality, str): + raise ValueError(f"{source} security label requires integrity and confidentiality") + + metadata = label_data.get("metadata") + if metadata is not None and not isinstance(metadata, dict): + logger.warning("Ignoring malformed metadata from %s security label", source) + metadata = {} + return ContentLabel( + integrity=IntegrityLabel(integrity), + confidentiality=ConfidentialityLabel(confidentiality), + metadata=cast(dict[str, Any], metadata) if isinstance(metadata, dict) else None, ) @@ -1271,17 +1280,18 @@ def _update_context_label(self, new_content_label: ContentLabel) -> None: ) @staticmethod - def _extract_primary_tool_content(expanded_content: Any) -> Any: - """Return the tool-visible content for an expanded variable payload. + def _extract_primary_tool_content(expanded_content: Any, *, from_quarantined_llm: bool) -> Any: + """Return the primary response from a proven quarantined LLM result. - Hidden ``quarantined_llm`` results are stored as rich payloads containing - the response plus security metadata. Tool arguments should receive only the - primary response they would have seen without variable indirection. Ordinary - mappings or JSON text that happen to contain a ``response`` key remain intact. + Producer metadata is owned by the middleware and distinguishes internal + quarantine wrappers from ordinary, potentially attacker-controlled JSON. """ + if not from_quarantined_llm: + return expanded_content + if isinstance(expanded_content, dict): content_map = cast(dict[str, Any], expanded_content) - if _is_quarantine_payload(content_map): + if "response" in content_map: return content_map["response"] return content_map @@ -1292,7 +1302,7 @@ def _extract_primary_tool_content(expanded_content: Any) -> Any: parsed = json.loads(stripped) if isinstance(parsed, dict): parsed_map = cast(dict[str, Any], parsed) - if _is_quarantine_payload(parsed_map): + if "response" in parsed_map: return parsed_map["response"] return expanded_content @@ -1308,7 +1318,11 @@ def _lookup_variable(self, variable_id: str, labels: list[ContentLabel]) -> Any: except KeyError: return _UNRESOLVED labels.append(stored_label) - return self._extract_primary_tool_content(stored_content) + metadata = self.get_variable_metadata(variable_id) + return self._extract_primary_tool_content( + stored_content, + from_quarantined_llm=metadata is not None and metadata.get("function_name") == "quarantined_llm", + ) def _resolve_string(self, value: str, labels: list[ContentLabel]) -> Any: if not _EMBEDDED_VAR_REF_RE.search(value): @@ -1395,14 +1409,18 @@ def _extract_labels_recursive(value: Any) -> None: if "security_label" in value_dict: label_data = value_dict["security_label"] if isinstance(label_data, ContentLabel): - labels.append(label_data) + labels.append( + _parse_content_label( + cast(MutableMapping[str, Any], label_data.to_dict()), + source="input", + ) + ) elif isinstance(label_data, dict): - with contextlib.suppress(Exception): # nosec B110 - best-effort label extraction - labels.append(ContentLabel.from_dict(cast(dict[str, Any], label_data))) - # Fall back to "label" for backward compatibility + with contextlib.suppress(TypeError, ValueError): + labels.append(_parse_content_label(cast(dict[str, Any], label_data), source="input")) elif "label" in value_dict and isinstance(value_dict.get("label"), dict): - with contextlib.suppress(Exception): # nosec B110 - best-effort label extraction - labels.append(ContentLabel.from_dict(cast(dict[str, Any], value_dict["label"]))) + with contextlib.suppress(TypeError, ValueError): + labels.append(_parse_content_label(cast(dict[str, Any], value_dict["label"]), source="input")) # Recurse into dict values for v in value_dict.values(): _extract_labels_recursive(v) @@ -1712,7 +1730,7 @@ def _process_result_with_embedded_labels( visible_item_labels: list[ContentLabel] = [] for item in items: - item_label = self._extract_content_label(item, fallback_label) + item_label = self._extract_content_label(item, fallback_label, function_name) item_labels.append(item_label) if self._should_hide(item_label, function_name): @@ -1732,38 +1750,58 @@ def _extract_content_label( self, item: Content, fallback_label: ContentLabel, + function_name: str, ) -> ContentLabel: - """Extract the security label for a single Content item. - - Checks (in order): - 1. ``additional_properties.security_label`` (explicit label) - 2. Falls back to ``fallback_label`` + """Extract and constrain the security label for one result item. Args: item: The Content item to inspect. - fallback_label: The label to use if no embedded label is found. + fallback_label: The label inherited from the invocation. + function_name: The name of the tool that produced the item. Returns: The resolved ContentLabel for this item. """ additional_props = _get_additional_properties(item) + authoritative_marker = additional_props.pop(_AUTHORITATIVE_CONFIDENTIALITY, None) + inspect_error_marker = additional_props.pop(_INSPECT_VARIABLE_ERROR, None) + authoritative_confidentiality = ( + function_name == "quarantined_llm" and authoritative_marker is _INTERNAL_RESULT_MARKER + ) + inspect_error = function_name == "inspect_variable" and inspect_error_marker is _INTERNAL_RESULT_MARKER - # Embedded labels remain authoritative for integrity, but cannot - # declassify data inherited from the tool's inputs. label_data = additional_props.get("security_label") if label_data and isinstance(label_data, dict): try: - embedded_label = ContentLabel.from_dict(cast(dict[str, Any], label_data)) + embedded_label = _parse_content_label( + cast(dict[str, Any], label_data), + source="embedded", + ) combined_label = combine_labels(fallback_label, embedded_label) return ContentLabel( integrity=embedded_label.integrity, - confidentiality=combined_label.confidentiality, + confidentiality=( + embedded_label.confidentiality + if authoritative_confidentiality + else combined_label.confidentiality + ), metadata=combined_label.metadata, ) - except Exception as e: - logger.warning(f"Failed to parse security_label from Content: {e}") + except (TypeError, ValueError) as exc: + logger.warning("Failed to parse security_label from Content: %s", exc) + + if inspect_error: + integrity = ( + IntegrityLabel.TRUSTED + if fallback_label.metadata.get("source") == "default" + else fallback_label.integrity + ) + return ContentLabel( + integrity=integrity, + confidentiality=fallback_label.confidentiality, + metadata=fallback_label.metadata, + ) - # No embedded label — use fallback return fallback_label def _hide_item( @@ -2914,9 +2952,17 @@ def _quarantined_llm_result_parser(result: Any) -> list[Content]: label_data = cast(dict[str, Any], result).get("security_label") if not isinstance(label_data, MutableMapping): return contents + typed_label_data = cast(MutableMapping[str, Any], label_data) + confidentiality = typed_label_data.get("confidentiality") + if not isinstance(confidentiality, str): + logger.warning("quarantined_llm result label is missing confidentiality") + return contents try: - parsed_label = ContentLabel.from_dict(cast(MutableMapping[str, Any], label_data)) + parsed_label = _parse_content_label( + typed_label_data, + source="quarantined_llm result", + ) except (TypeError, ValueError) as exc: logger.warning("Failed to parse quarantined_llm result label: %s", exc) return contents @@ -2929,6 +2975,7 @@ def _quarantined_llm_result_parser(result: Any) -> list[Content]: first = contents[0] props = first.additional_properties or {} props["security_label"] = quarantine_label.to_dict() + props[_AUTHORITATIVE_CONFIDENTIALITY] = _INTERNAL_RESULT_MARKER first.additional_properties = props return contents @@ -3013,6 +3060,10 @@ async def quarantined_llm( variable_store = middleware.get_variable_store() if middleware else _global_variable_store labels: list[ContentLabel] = [] + unknown_input_label = ContentLabel( + integrity=IntegrityLabel.UNTRUSTED, + confidentiality=ConfidentialityLabel.PRIVATE, + ) retrieved_content: dict[str, Any] = {} # Retrieve content from variable_ids @@ -3025,7 +3076,7 @@ async def quarantined_llm( except KeyError: logger.warning("A requested quarantine variable was not found in the current security scope") # Still add untrusted label for unknown variables - labels.append(ContentLabel(integrity=IntegrityLabel.UNTRUSTED)) + labels.append(unknown_input_label) # Parse labels and content from labelled_data labelled_data_content: dict[str, Any] = {} @@ -3044,19 +3095,24 @@ async def quarantined_llm( try: label_data = value_dict[label_key] if isinstance(label_data, dict): - label = ContentLabel.from_dict(cast(dict[str, Any], label_data)) + label = _parse_content_label( + cast(dict[str, Any], label_data), + source="quarantine input", + ) elif isinstance(label_data, ContentLabel): label = label_data else: - label = ContentLabel(integrity=IntegrityLabel.UNTRUSTED) + label = unknown_input_label labels.append(label) - except Exception as e: + except (TypeError, ValueError) as e: logger.warning("Failed to parse a quarantine data security label; using UNTRUSTED") logger.debug("Quarantine label parse failure for %s: %s", key, e) - labels.append(ContentLabel(integrity=IntegrityLabel.UNTRUSTED)) + labels.append(unknown_input_label) else: # No label provided, default to UNTRUSTED - labels.append(ContentLabel(integrity=IntegrityLabel.UNTRUSTED)) + labels.append(unknown_input_label) + else: + labels.append(unknown_input_label) # Combine all labels (most restrictive) combined_label = combine_labels(*labels) if labels else ContentLabel(integrity=IntegrityLabel.UNTRUSTED) @@ -3197,25 +3253,21 @@ def _inspect_variable_result_parser(result: Any) -> list[Content]: downgrading the real label. This parser stamps the inspected label back onto the produced Content so - ``LabelTrackingFunctionMiddleware`` propagates it faithfully. Missing-variable - errors are labeled as trusted tool output so probing an absent or foreign id - does not falsely taint cumulative integrity. + ``LabelTrackingFunctionMiddleware`` propagates it faithfully. Error results + remain unstamped so their integrity inherits from the invocation fallback. """ contents = FunctionTool.parse_result(result) if not contents: return contents + first = contents[0] + props = first.additional_properties or {} label = cast(dict[str, Any], result).get("security_label") if isinstance(result, dict) else None - if label is None and isinstance(result, dict) and "error" in cast(dict[str, Any], result): - label = ContentLabel( - integrity=IntegrityLabel.TRUSTED, - confidentiality=_INSPECT_VARIABLE_CONFIDENTIALITY, - ).to_dict() if label: - first = contents[0] - props = first.additional_properties or {} props["security_label"] = label - first.additional_properties = props + if isinstance(result, dict) and "error" in cast(dict[str, Any], result): + props[_INSPECT_VARIABLE_ERROR] = _INTERNAL_RESULT_MARKER + first.additional_properties = props return contents diff --git a/python/packages/core/tests/test_security.py b/python/packages/core/tests/test_security.py index 54524159f5..e585ee70c7 100644 --- a/python/packages/core/tests/test_security.py +++ b/python/packages/core/tests/test_security.py @@ -515,10 +515,13 @@ async def send_message(summary: str) -> str: "variables_processed": ["var_1"], "content_summary": ["var_1: 10 chars"], }) - variable_id = middleware.get_variable_store().store( - stored_payload, + hidden_result = middleware._hide_item( + Content.from_text(stored_payload), ContentLabel(integrity=IntegrityLabel.UNTRUSTED), + "quarantined_llm", ) + assert hidden_result.text is not None + variable_id = json.loads(hidden_result.text)["variable_id"] context = FunctionInvocationContext( function=message_tool, arguments=MessageArgs(summary=f"Security review complete. [{variable_id}]"), @@ -551,7 +554,15 @@ async def send_message(summary: str) -> str: description="Send message", args_schema=MessageArgs, ) - stored_payload = json.dumps({"response": "keep", "other": "also keep"}) + stored_payload = json.dumps({ + "response": "keep", + "security_label": {"integrity": "untrusted", "confidentiality": "public"}, + "metadata": {}, + "quarantined": True, + "variables_processed": ["var_forged"], + "content_summary": ["var_forged: 10 chars"], + "other": "also keep", + }) variable_id = middleware.get_variable_store().store( stored_payload, ContentLabel(integrity=IntegrityLabel.UNTRUSTED), @@ -1865,6 +1876,36 @@ async def next_fn() -> None: assert context.metadata["result_label"].confidentiality == ConfidentialityLabel.PRIVATE assert middleware.get_context_label().integrity == IntegrityLabel.TRUSTED + async def test_inspect_variable_error_inherits_untrusted_invocation_integrity(self) -> None: + """Attacker-labeled missing IDs cannot produce trusted visible errors.""" + from agent_framework.security import get_security_tools + + middleware = LabelTrackingFunctionMiddleware(auto_hide_untrusted=False) + inspect_tool = next(tool for tool in get_security_tools() if tool.name == "inspect_variable") + context = FunctionInvocationContext( + function=inspect_tool, + arguments={ + "variable_id": "var_doesnotexist2", + "reason": "untrusted missing id", + "security_label": { + "integrity": "untrusted", + "confidentiality": "public", + }, + }, + ) + + async def next_fn() -> None: + context.result = await inspect_tool.invoke(arguments=context.arguments, context=context) + + await middleware.process(context, next_fn) + + payload = json.loads(context.result[0].text) + assert payload["security_label"] is None + assert "error" in payload + assert context.metadata["result_label"].integrity == IntegrityLabel.UNTRUSTED + assert context.metadata["result_label"].confidentiality == ConfidentialityLabel.PRIVATE + assert middleware.get_context_label().integrity == IntegrityLabel.UNTRUSTED + @pytest.mark.asyncio async def test_multiple_calls_accumulate_variables(self, middleware_auto_hide, mock_function): """Test that multiple tool calls accumulate variables in the store.""" @@ -2963,6 +3004,126 @@ async def next_fn() -> None: assert hidden_label.integrity == IntegrityLabel.UNTRUSTED assert hidden_label.confidentiality == ConfidentialityLabel.USER_IDENTITY + async def test_quarantined_llm_public_input_remains_public(self) -> None: + """A valid quarantine label overrides the fail-closed PRIVATE fallback.""" + from agent_framework.security import set_quarantine_client + + set_quarantine_client(None) + middleware = LabelTrackingFunctionMiddleware() + quarantine_tool = next(tool for tool in middleware.get_security_tools() if tool.name == "quarantined_llm") + context = FunctionInvocationContext( + function=quarantine_tool, + arguments={ + "prompt": "Summarize", + "labelled_data": { + "data": { + "content": "public information", + "security_label": { + "integrity": "trusted", + "confidentiality": "public", + "metadata": ["malformed"], + }, + } + }, + }, + ) + + async def next_fn() -> None: + context.result = await quarantine_tool.invoke(arguments=context.arguments, context=context) + + await middleware.process(context, next_fn) + + assert context.metadata["result_label"].integrity == IntegrityLabel.UNTRUSTED + assert context.metadata["result_label"].confidentiality == ConfidentialityLabel.PUBLIC + assert middleware.get_context_label().integrity == IntegrityLabel.TRUSTED + assert middleware.get_context_label().confidentiality == ConfidentialityLabel.PUBLIC + + async def test_quarantined_llm_invalid_label_falls_back_to_private(self) -> None: + """A partial parser label cannot override the PRIVATE tool fallback.""" + from agent_framework.security import _quarantined_llm_result_parser + + middleware = LabelTrackingFunctionMiddleware() + quarantine_tool = next(tool for tool in middleware.get_security_tools() if tool.name == "quarantined_llm") + context = FunctionInvocationContext( + function=quarantine_tool, + arguments={"prompt": "Summarize"}, + ) + + async def next_fn() -> None: + context.result = _quarantined_llm_result_parser({ + "response": "partial label", + "security_label": {"integrity": "untrusted"}, + }) + + await middleware.process(context, next_fn) + + assert context.metadata["result_label"].integrity == IntegrityLabel.UNTRUSTED + assert context.metadata["result_label"].confidentiality == ConfidentialityLabel.PRIVATE + hidden_reference = json.loads(context.result[0].text) + _, hidden_label = middleware.get_variable_store().retrieve(hidden_reference["variable_id"]) + assert hidden_label.integrity == IntegrityLabel.UNTRUSTED + assert hidden_label.confidentiality == ConfidentialityLabel.PRIVATE + + async def test_quarantined_llm_partial_input_label_is_private(self) -> None: + """An incomplete quarantine input label fails closed to PRIVATE.""" + from agent_framework.security import set_quarantine_client + + set_quarantine_client(None) + middleware = LabelTrackingFunctionMiddleware() + quarantine_tool = next(tool for tool in middleware.get_security_tools() if tool.name == "quarantined_llm") + context = FunctionInvocationContext( + function=quarantine_tool, + arguments={ + "prompt": "Summarize", + "labelled_data": { + "data": { + "content": "unknown information", + "security_label": {"integrity": "trusted"}, + } + }, + }, + ) + + async def next_fn() -> None: + context.result = await quarantine_tool.invoke(arguments=context.arguments, context=context) + + await middleware.process(context, next_fn) + + assert context.metadata["result_label"].integrity == IntegrityLabel.UNTRUSTED + assert context.metadata["result_label"].confidentiality == ConfidentialityLabel.PRIVATE + + async def test_quarantined_llm_malformed_input_is_private(self) -> None: + """A malformed item makes mixed quarantine input PRIVATE.""" + from agent_framework.security import set_quarantine_client + + set_quarantine_client(None) + middleware = LabelTrackingFunctionMiddleware() + quarantine_tool = next(tool for tool in middleware.get_security_tools() if tool.name == "quarantined_llm") + context = FunctionInvocationContext( + function=quarantine_tool, + arguments={ + "prompt": "Summarize", + "labelled_data": { + "public": { + "content": "public information", + "security_label": { + "integrity": "trusted", + "confidentiality": "public", + }, + }, + "malformed": "unlabeled information", + }, + }, + ) + + async def next_fn() -> None: + context.result = await quarantine_tool.invoke(arguments=context.arguments, context=context) + + await middleware.process(context, next_fn) + + assert context.metadata["result_label"].integrity == IntegrityLabel.UNTRUSTED + assert context.metadata["result_label"].confidentiality == ConfidentialityLabel.PRIVATE + @pytest.mark.asyncio async def test_quarantined_llm_returns_response(self): """Test that quarantined_llm returns a plain response dict.""" @@ -3476,6 +3637,61 @@ async def next_fn() -> None: assert middleware.get_context_label().integrity == IntegrityLabel.TRUSTED assert middleware.get_context_label().confidentiality == ConfidentialityLabel.USER_IDENTITY + async def test_malformed_embedded_metadata_preserves_mandatory_label_fields( + self, middleware, mock_function + ) -> None: + """Malformed optional metadata cannot discard integrity or confidentiality.""" + mock_function.additional_properties = {"source_integrity": "trusted"} + context = FunctionInvocationContext(function=mock_function, arguments=mock_function.args_schema()) + + async def next_fn() -> None: + context.result = [ + Content.from_text( + "hidden identity data", + additional_properties={ + "security_label": { + "integrity": "untrusted", + "confidentiality": "user_identity", + "metadata": ["malformed"], + } + }, + ) + ] + + await middleware.process(context, next_fn) + + result_label = context.metadata["result_label"] + assert result_label.integrity == IntegrityLabel.UNTRUSTED + assert result_label.confidentiality == ConfidentialityLabel.USER_IDENTITY + assert result_label.metadata["source"] == "source_integrity" + assert context.result[0].additional_properties["_variable_reference"] is True + assert middleware.get_context_label().integrity == IntegrityLabel.TRUSTED + assert middleware.get_context_label().confidentiality == ConfidentialityLabel.USER_IDENTITY + + async def test_partial_embedded_label_uses_untrusted_fallback(self, middleware, mock_function) -> None: + """A partial embedded label cannot promote an untrusted fallback.""" + context = FunctionInvocationContext(function=mock_function, arguments=mock_function.args_schema()) + + async def next_fn() -> None: + context.result = [ + Content.from_text( + "untrusted result", + additional_properties={ + "security_label": { + "metadata": {"note": "missing mandatory fields"}, + } + }, + ) + ] + + await middleware.process(context, next_fn) + + result_label = context.metadata["result_label"] + assert result_label.integrity == IntegrityLabel.UNTRUSTED + assert result_label.confidentiality == ConfidentialityLabel.PUBLIC + assert context.result[0].additional_properties["_variable_reference"] is True + assert middleware.get_context_label().integrity == IntegrityLabel.TRUSTED + @pytest.mark.asyncio async def test_items_without_labels_use_fallback(self, middleware, mock_function): """Test that items without embedded labels use the fallback (call) label.""" @@ -3765,7 +3981,10 @@ async def next_fn() -> None: context.result = [ Content.from_text( "transformed", - additional_properties={"security_label": {"integrity": "trusted", "confidentiality": "public"}}, + additional_properties={ + "security_label": {"integrity": "trusted", "confidentiality": "public"}, + "_security_label_authoritative_confidentiality": True, + }, ) ] From d800337793a1876c7b9c37daed055a6c3eb22596 Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Wed, 9 Sep 2026 12:07:13 +0200 Subject: [PATCH 5/5] Python: keep unlabeled quarantine output private Use the existing UNTRUSTED/PRIVATE fail-closed label when quarantined_llm receives no labeled input, with placeholder and configured-client regressions.\n\nCo-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../packages/core/agent_framework/security.py | 2 +- python/packages/core/tests/test_security.py | 71 +++++++++++++++++++ 2 files changed, 72 insertions(+), 1 deletion(-) diff --git a/python/packages/core/agent_framework/security.py b/python/packages/core/agent_framework/security.py index 29f8df7cda..411cb09696 100644 --- a/python/packages/core/agent_framework/security.py +++ b/python/packages/core/agent_framework/security.py @@ -3115,7 +3115,7 @@ async def quarantined_llm( labels.append(unknown_input_label) # Combine all labels (most restrictive) - combined_label = combine_labels(*labels) if labels else ContentLabel(integrity=IntegrityLabel.UNTRUSTED) + combined_label = combine_labels(*labels) if labels else unknown_input_label content_summary: list[str] = [] for var_id, content in retrieved_content.items(): diff --git a/python/packages/core/tests/test_security.py b/python/packages/core/tests/test_security.py index e585ee70c7..32c3938118 100644 --- a/python/packages/core/tests/test_security.py +++ b/python/packages/core/tests/test_security.py @@ -3124,6 +3124,77 @@ async def next_fn() -> None: assert context.metadata["result_label"].integrity == IntegrityLabel.UNTRUSTED assert context.metadata["result_label"].confidentiality == ConfidentialityLabel.PRIVATE + async def test_quarantined_llm_empty_input_placeholder_is_private(self) -> None: + """An unlabeled placeholder response remains PRIVATE.""" + from agent_framework.security import set_quarantine_client + + set_quarantine_client(None) + middleware = LabelTrackingFunctionMiddleware() + quarantine_tool = next(tool for tool in middleware.get_security_tools() if tool.name == "quarantined_llm") + context = FunctionInvocationContext( + function=quarantine_tool, + arguments={ + "prompt": "Summarize sensitive instructions", + "variable_ids": [], + "labelled_data": {}, + }, + ) + + async def next_fn() -> None: + context.result = await quarantine_tool.invoke(arguments=context.arguments, context=context) + + await middleware.process(context, next_fn) + + result_label = context.metadata["result_label"] + assert result_label.integrity == IntegrityLabel.UNTRUSTED + assert result_label.confidentiality == ConfidentialityLabel.PRIVATE + hidden_reference = json.loads(context.result[0].text) + hidden_content, hidden_label = middleware.get_variable_store().retrieve(hidden_reference["variable_id"]) + assert hidden_label.integrity == IntegrityLabel.UNTRUSTED + assert hidden_label.confidentiality == ConfidentialityLabel.PRIVATE + assert "[Quarantined LLM Response]" in json.loads(hidden_content)["response"] + + async def test_quarantined_llm_empty_input_client_response_is_private(self) -> None: + """An unlabeled client response remains PRIVATE.""" + from unittest.mock import AsyncMock, MagicMock + + from agent_framework.security import set_quarantine_client + + mock_response = MagicMock() + mock_response.text = "client-produced sensitive response" + mock_client = MagicMock() + mock_client.get_response = AsyncMock(return_value=mock_response) + set_quarantine_client(mock_client) + + try: + middleware = LabelTrackingFunctionMiddleware() + quarantine_tool = next(tool for tool in middleware.get_security_tools() if tool.name == "quarantined_llm") + context = FunctionInvocationContext( + function=quarantine_tool, + arguments={ + "prompt": "Summarize sensitive instructions", + "variable_ids": [], + "labelled_data": {}, + }, + ) + + async def next_fn() -> None: + context.result = await quarantine_tool.invoke(arguments=context.arguments, context=context) + + await middleware.process(context, next_fn) + + result_label = context.metadata["result_label"] + assert result_label.integrity == IntegrityLabel.UNTRUSTED + assert result_label.confidentiality == ConfidentialityLabel.PRIVATE + hidden_reference = json.loads(context.result[0].text) + hidden_content, hidden_label = middleware.get_variable_store().retrieve(hidden_reference["variable_id"]) + assert hidden_label.integrity == IntegrityLabel.UNTRUSTED + assert hidden_label.confidentiality == ConfidentialityLabel.PRIVATE + assert json.loads(hidden_content)["response"] == "client-produced sensitive response" + mock_client.get_response.assert_awaited_once() + finally: + set_quarantine_client(None) + @pytest.mark.asyncio async def test_quarantined_llm_returns_response(self): """Test that quarantined_llm returns a plain response dict."""