diff --git a/docs/features/FIDES_IMPLEMENTATION_SUMMARY.md b/docs/features/FIDES_IMPLEMENTATION_SUMMARY.md index 100166b7da..6eee1baac4 100644 --- a/docs/features/FIDES_IMPLEMENTATION_SUMMARY.md +++ b/docs/features/FIDES_IMPLEMENTATION_SUMMARY.md @@ -108,10 +108,24 @@ async def fetch_emails(count: int = 5) -> list[Content]: }), additional_properties={ "security_label": { - "integrity": "trusted" if email["is_internal"] else "untrusted", + "integrity": "trusted" if email["internal"] else "untrusted", "confidentiality": "private", } - }, + ), + ) + for email in emails + ] +``` + +These embedded labels are automatically consumed by `LabelTrackingFunctionMiddleware`, which: +- Extracts the `security_label` from `additional_properties` +- Uses the embedded label as the highest-priority source for that item +- Automatically hides UNTRUSTED items in the variable store +- Replaces hidden items with `VariableReferenceContent` in the LLM context +- Preserves TRUSTED items visible to the LLM without tainting the context label + +This enables tools to return mixed-trust data where some items (internal emails) remain visible while untrusted items (external emails) are automatically hidden without manual intervention. + }, ) for email in emails ] @@ -119,6 +133,8 @@ async def fetch_emails(count: int = 5) -> list[Content]: ### 3. Automatic Variable Hiding +This feature automatically hides any UNTRUSTED content returned by tools while keeping the hiding logic transparent to the developer. Developers do not need to manually call `store_untrusted_content()`. This allows the LLM /agent's context to remain clean and secure. Key aspects include: + - **Automatic Detection**: Middleware checks integrity label after each tool call - **Automatic Storage**: UNTRUSTED results/items stored in variable store - **Transparent Replacement**: LLM context receives `VariableReferenceContent` @@ -168,16 +184,6 @@ agent = Agent( ) ``` -### 7. Message-Level Label Tracking (Phase 1) - -Track security labels at the message level: - -```python -labeled_messages = middleware.label_messages(messages) -label = middleware.get_message_label(5) -all_labels = middleware.get_all_message_labels() -``` - ## Security Properties ### Deterministic Defense @@ -323,11 +329,6 @@ cd python/packages/core && ../../.venv/bin/pytest tests/test_security.py -v ✅ `quarantine_chat_client` support for real LLM calls ✅ `SECURITY_TOOL_INSTRUCTIONS` constant -### Phase 1: Message-Level Tracking -✅ `LabeledMessage` class with auto-inference from role -✅ `label_message()`, `get_message_label()`, `label_messages()` methods -✅ `get_all_message_labels()` method - ### Documentation & Testing ✅ Complete FIDES Developer Guide (~1250 lines) ✅ Architecture Decision Record (ADR) diff --git a/python/packages/core/agent_framework/__init__.py b/python/packages/core/agent_framework/__init__.py index 686abf781b..13d7bade00 100644 --- a/python/packages/core/agent_framework/__init__.py +++ b/python/packages/core/agent_framework/__init__.py @@ -130,7 +130,6 @@ from ._tools import ( FunctionInvocationLayer, FunctionTool, ToolTypes, - ai_function, normalize_function_invocation_configuration, tool, ) @@ -430,7 +429,6 @@ __all__ = [ "__version__", "add_usage_details", "agent_middleware", - "ai_function", "annotate_message_groups", "apply_compaction", "chat_middleware", diff --git a/python/packages/core/agent_framework/security.py b/python/packages/core/agent_framework/security.py index aa80b12fbf..6d3b1d0d59 100644 --- a/python/packages/core/agent_framework/security.py +++ b/python/packages/core/agent_framework/security.py @@ -831,15 +831,19 @@ class LabelTrackingFunctionMiddleware(FunctionMiddleware): Examples: .. code-block:: python - from agent_framework import Agent + from agent_framework import Agent, LabelTrackingFunctionMiddleware, tool + + + @tool(additional_properties={"source_integrity": "trusted"}) + async def get_weather(city: str) -> str: + return f"Weather in {city}: 72°F" - from agent_framework.security import LabelTrackingFunctionMiddleware # Create agent with automatic hiding enabled middleware = LabelTrackingFunctionMiddleware( auto_hide_untrusted=True # Enabled by default ) - agent = Agent(client=client, name="assistant", middleware=[middleware]) + agent = Agent(client=client, name="assistant", tools=[get_weather], middleware=[middleware]) # Run agent - untrusted tool results are automatically hidden response = await agent.run(messages=[{"role": "user", "content": "What's the weather?"}]) @@ -880,10 +884,6 @@ class LabelTrackingFunctionMiddleware(FunctionMiddleware): # Metadata about stored variables self._variable_metadata: dict[str, dict[str, Any]] = {} - # Phase 1: Message-level label tracking - # Maps message index to its security label - self._message_labels: dict[int, ContentLabel] = {} - def get_context_label(self) -> ContentLabel: """Get the current context-level security label. @@ -904,79 +904,8 @@ class LabelTrackingFunctionMiddleware(FunctionMiddleware): self._context_label = ContentLabel( integrity=IntegrityLabel.TRUSTED, confidentiality=ConfidentialityLabel.PUBLIC, metadata={"reset": True} ) - # Also reset message labels for new conversation - self._message_labels.clear() logger.info("Context label reset to TRUSTED + PUBLIC") - # ========== Phase 1: Message-Level Label Tracking ========== - - def label_message( - self, - message_index: int, - label: ContentLabel, - source_labels: list[ContentLabel] | None = None, - ) -> None: - """Assign a security label to a message in the conversation. - - Args: - message_index: The index of the message in the conversation. - label: The security label to assign. - source_labels: Optional list of labels that contributed to this message. - """ - self._message_labels[message_index] = label - logger.debug(f"Labeled message {message_index}: {label.integrity.value}/{label.confidentiality.value}") - - def get_message_label(self, message_index: int) -> ContentLabel | None: - """Get the security label of a specific message. - - Args: - message_index: The index of the message. - - Returns: - The message's ContentLabel, or None if not labeled. - """ - return self._message_labels.get(message_index) - - def label_messages(self, messages: list[dict[str, Any]]) -> list[LabeledMessage]: - """Label a list of messages based on their roles and content. - - This method automatically assigns labels to messages: - - user/system messages: TRUSTED - - assistant messages: Inherit from source labels or TRUSTED - - tool messages: UNTRUSTED (external data) - - Args: - messages: List of message dicts with 'role' and 'content'. - - Returns: - List of LabeledMessage objects. - """ - labeled: list[LabeledMessage] = [] - for i, msg in enumerate(messages): - # Check if message already has a label - existing_label = self._message_labels.get(i) - - labeled_msg = LabeledMessage( - role=msg.get("role", "unknown"), - content=msg.get("content", ""), - security_label=existing_label, # Will auto-infer if None - message_index=i, - ) - - # Store the label - self._message_labels[i] = labeled_msg.security_label - labeled.append(labeled_msg) - - return labeled - - def get_all_message_labels(self) -> dict[int, ContentLabel]: - """Get all message labels. - - Returns: - Dictionary mapping message index to ContentLabel. - """ - return dict(self._message_labels) - def _update_context_label(self, new_content_label: ContentLabel) -> None: """Update the context label based on new content added to the context. diff --git a/python/packages/core/tests/test_security.py b/python/packages/core/tests/test_security.py index 931aa074fa..0a638f5883 100644 --- a/python/packages/core/tests/test_security.py +++ b/python/packages/core/tests/test_security.py @@ -1464,61 +1464,6 @@ class TestLabeledMessage: assert labeled.is_trusted() -class TestMiddlewareMessageLabeling: - """Tests for middleware message label tracking.""" - - def test_label_message(self): - """Test labeling a message by index.""" - middleware = LabelTrackingFunctionMiddleware() - - label = ContentLabel(integrity=IntegrityLabel.UNTRUSTED, confidentiality=ConfidentialityLabel.PRIVATE) - middleware.label_message(5, label) - - retrieved = middleware.get_message_label(5) - assert retrieved is not None - assert retrieved.integrity == IntegrityLabel.UNTRUSTED - - def test_get_unlabeled_message_returns_none(self): - """Test that unlabeled messages return None.""" - middleware = LabelTrackingFunctionMiddleware() - - assert middleware.get_message_label(999) is None - - def test_label_messages_batch(self): - """Test batch labeling of messages.""" - middleware = LabelTrackingFunctionMiddleware() - - messages = [ - {"role": "user", "content": "Hello"}, - {"role": "assistant", "content": "Hi there"}, - {"role": "tool", "content": "External data"}, - ] - - labeled = middleware.label_messages(messages) - - assert len(labeled) == 3 - assert labeled[0].security_label.integrity == IntegrityLabel.TRUSTED - assert labeled[1].security_label.integrity == IntegrityLabel.TRUSTED - assert labeled[2].security_label.integrity == IntegrityLabel.UNTRUSTED - - # Check that labels are stored in middleware - all_labels = middleware.get_all_message_labels() - assert len(all_labels) == 3 - - def test_reset_clears_message_labels(self): - """Test that reset_context_label also clears message labels.""" - middleware = LabelTrackingFunctionMiddleware() - - middleware.label_message(0, ContentLabel()) - middleware.label_message(1, ContentLabel()) - - assert len(middleware.get_all_message_labels()) == 2 - - middleware.reset_context_label() - - assert len(middleware.get_all_message_labels()) == 0 - - # ========== Quarantined LLM Tests ========== diff --git a/python/samples/02-agents/security/FIDES_DEVELOPER_GUIDE.md b/python/samples/02-agents/security/FIDES_DEVELOPER_GUIDE.md index 9cf72549dc..3a1fbf82d2 100644 --- a/python/samples/02-agents/security/FIDES_DEVELOPER_GUIDE.md +++ b/python/samples/02-agents/security/FIDES_DEVELOPER_GUIDE.md @@ -495,36 +495,9 @@ The instructions explain: - How to pass `variable_ids` to reference hidden content - Best practices for secure content handling -### 9. Message-Level Label Tracking (Phase 1) +### 9. LabeledMessage Class -The middleware now tracks security labels at the **message level**, not just tool calls: - -```python -from agent_framework.security import LabelTrackingFunctionMiddleware, LabeledMessage - -middleware = LabelTrackingFunctionMiddleware() - -# Label messages in a conversation -messages = [ - {"role": "user", "content": "Hello"}, # Auto-labeled TRUSTED - {"role": "assistant", "content": "Hi there"}, # Auto-labeled TRUSTED (no untrusted sources) - {"role": "tool", "content": "API response"}, # Auto-labeled UNTRUSTED -] - -labeled_messages = middleware.label_messages(messages) -# labeled_messages[0].security_label.integrity == TRUSTED -# labeled_messages[2].security_label.integrity == UNTRUSTED - -# Individual message labeling -middleware.label_message(message_index=5, label=custom_label) -label = middleware.get_message_label(5) - -# Get all message labels -all_labels = middleware.get_all_message_labels() -``` - -**LabeledMessage Class:** -- Automatically infers labels based on message role +**LabeledMessage** automatically infers security labels based on message role: - User/system messages → TRUSTED - Tool messages → UNTRUSTED - Assistant messages → Inherit from source_labels or TRUSTED @@ -1108,18 +1081,6 @@ LabeledMessage.from_dict(data) -> LabeledMessage # Deserialize LabeledMessage.from_message(msg, index) -> LabeledMessage # Wrap standard message ``` -### LabelTrackingFunctionMiddleware Extensions - -```python -middleware = LabelTrackingFunctionMiddleware(...) - -# Message-level label tracking (Phase 1) -middleware.label_message(message_index, label, source_labels=None) # Label a message -middleware.get_message_label(message_index) -> ContentLabel | None # Get message label -middleware.label_messages(messages) -> List[LabeledMessage] # Batch label messages -middleware.get_all_message_labels() -> Dict[int, ContentLabel] # Get all message labels -``` - ### SecureAgentConfig ```python diff --git a/python/samples/02-agents/security/email_security_example.py b/python/samples/02-agents/security/email_security_example.py index 55093f6f10..b8cd0a36d1 100644 --- a/python/samples/02-agents/security/email_security_example.py +++ b/python/samples/02-agents/security/email_security_example.py @@ -278,7 +278,13 @@ async def run_scenarios(agent, config): print("- Injection attempts in emails are NOT followed") print() - response = await agent.run("Please fetch my recent emails and give me a brief summary of each one.") + # Use a shared session so conversation history persists across scenarios. + # Without this, each agent.run() starts a fresh conversation and the LLM + # won't know about the emails fetched in Scenario 1 — it would never + # attempt to call send_email, so the policy enforcer would never trigger. + session = agent.create_session() + + response = await agent.run("Please fetch my recent emails and give me a brief summary of each one.", session=session) print(f"\n📋 Agent Response:\n{'-' * 40}") print(response.text) @@ -295,7 +301,9 @@ async def run_scenarios(agent, config): print("- Agent should explain it cannot send email due to security policy") print() - response = await agent.run("Now please send an email to colleague@company.com summarizing what you found.") + response = await agent.run( + "Now please send an email to colleague@company.com summarizing what you found.", session=session + ) print(f"\n📋 Agent Response:\n{'-' * 40}") print(response.text)