Python: updated import naming and comment from review (#5421)

* updated import naming and comment from review

* Add approval replay None call-id test

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

---------

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

python/packages/core/AGENTS.md

@@ -7,6 +7,7 @@ The foundation package containing all core abstractions, types, and built-in Ope
 ```
 agent_framework/
 ├── __init__.py          # Public API exports
+├── security.py          # Public security primitives, middleware, and tools
 ├── _agents.py           # Agent implementations
 ├── _clients.py          # Chat client base classes and protocols
 ├── _types.py            # Core types (Message, ChatResponse, Content, etc.)

python/packages/core/agent_framework/__init__.py

@@ -100,25 +100,6 @@ from ._middleware import (
     chat_middleware,
     function_middleware,
 )
-from ._security import (
-    SECURITY_TOOL_INSTRUCTIONS,
-    ConfidentialityLabel,
-    ContentLabel,
-    ContentVariableStore,
-    IntegrityLabel,
-    LabeledMessage,
-    LabelTrackingFunctionMiddleware,
-    PolicyEnforcementFunctionMiddleware,
-    SecureAgentConfig,
-    VariableReferenceContent,
-    check_confidentiality_allowed,
-    combine_labels,
-    get_quarantine_client,
-    get_security_tools,
-    quarantined_llm,
-    set_quarantine_client,
-    store_untrusted_content,
-)
 from ._sessions import (
     AgentSession,
     ContextProvider,
@@ -289,7 +270,6 @@ __all__ = [
     "GROUP_INDEX_KEY",
     "GROUP_KIND_KEY",
     "GROUP_TOKEN_COUNT_KEY",
-    "SECURITY_TOOL_INSTRUCTIONS",
     "SKIP_PARSING",
     "SUMMARIZED_BY_SUMMARY_ID_KEY",
     "SUMMARY_OF_GROUP_IDS_KEY",
@@ -328,10 +308,7 @@ __all__ = [
     "CheckpointStorage",
     "CompactionProvider",
     "CompactionStrategy",
-    "ConfidentialityLabel",
     "Content",
-    "ContentLabel",
-    "ContentVariableStore",
     "ContextProvider",
     "ContinuationToken",
     "ConversationSplit",
@@ -375,9 +352,6 @@ __all__ = [
     "InMemoryCheckpointStorage",
     "InMemoryHistoryProvider",
     "InProcRunnerContext",
-    "IntegrityLabel",
-    "LabelTrackingFunctionMiddleware",
-    "LabeledMessage",
     "LocalEvaluator",
     "MCPStdioTool",
     "MCPStreamableHTTPTool",
@@ -389,7 +363,6 @@ __all__ = [
     "MiddlewareTypes",
     "OuterFinalT",
     "OuterUpdateT",
-    "PolicyEnforcementFunctionMiddleware",
     "RawAgent",
     "ReleaseCandidateFeature",
     "ResponseStream",
@@ -399,7 +372,6 @@ __all__ = [
     "Runner",
     "RunnerContext",
     "SecretString",
-    "SecureAgentConfig",
     "SelectiveToolCallCompactionStrategy",
     "SessionContext",
     "SingleEdgeGroup",
@@ -436,7 +408,6 @@ __all__ = [
     "UsageDetails",
     "UserInputRequiredException",
     "ValidationTypeEnum",
-    "VariableReferenceContent",
     "Workflow",
     "WorkflowAgent",
     "WorkflowBuilder",
@@ -463,8 +434,6 @@ __all__ = [
     "annotate_message_groups",
     "apply_compaction",
     "chat_middleware",
-    "check_confidentiality_allowed",
-    "combine_labels",
     "create_edge_runner",
     "detect_media_type_from_base64",
     "evaluate_agent",
@@ -472,9 +441,7 @@ __all__ = [
     "evaluator",
     "executor",
     "function_middleware",
-    "get_quarantine_client",
     "get_run_context",
-    "get_security_tools",
     "handler",
     "included_messages",
     "included_token_count",
@@ -487,13 +454,10 @@ __all__ = [
     "normalize_tools",
     "prepend_agent_framework_to_user_agent",
     "prepend_instructions_to_messages",
-    "quarantined_llm",
     "register_state_type",
     "resolve_agent_id",
     "response_handler",
-    "set_quarantine_client",
     "step",
-    "store_untrusted_content",
     "tool",
     "tool_call_args_match",
     "tool_called_check",

python/packages/core/agent_framework/_tools.py

@@ -1917,21 +1917,16 @@ def _replace_approval_contents_with_results(
         Content,
     )
 
-    # Build a map of call_id -> actual result for replacing placeholders
+    # Match results back to approvals by actual call_id instead of relying on
+    # approval/result iteration order.
     result_by_call_id: dict[str, Content] = {}
-    for resp in fcc_todo.values():
-        if resp.approved and resp.function_call is not None and resp.function_call.call_id is not None:
-            # Map the call_id from the function_call to be replaced
-            call_id = resp.function_call.call_id
-            if call_id not in result_by_call_id and approved_function_results:
-                idx = len(result_by_call_id)
-                if idx < len(approved_function_results):
-                    result_by_call_id[call_id] = approved_function_results[idx]
+    for approved_result in approved_function_results:
+        if approved_result.call_id is not None and approved_result.call_id not in result_by_call_id:
+            result_by_call_id[approved_result.call_id] = approved_result
 
     # Track which call_ids had their placeholders replaced
     placeholders_replaced: set[str] = set()
 
-    result_idx = 0
     for msg in messages:
         # First pass - collect existing function call IDs to avoid duplicates
         existing_call_ids = {
@@ -1970,9 +1965,9 @@ def _replace_approval_contents_with_results(
                     else:
                         # No placeholder - replace approval response with result directly
                         # This handles the original approval_mode="always_require" case
-                        if result_idx < len(approved_function_results):
-                            msg.contents[content_idx] = approved_function_results[result_idx]
-                            result_idx += 1
+                        replacement_result = result_by_call_id.get(call_id)
+                        if replacement_result is not None:
+                            msg.contents[content_idx] = replacement_result
                             msg.role = "tool"
                 else:
                     # Create a "not approved" result for rejected calls

python/packages/core/agent_framework/security.py

@@ -2,7 +2,7 @@
 
 """Security infrastructure for prompt injection defense.
 
-This module provides information-flow control-basedsecurity mechanisms to defend against prompt injection attacks
+This module provides information-flow control-based security mechanisms to defend against prompt injection attacks
 by tracking integrity and confidentiality of content throughout agent execution.
 
 It includes:
@@ -12,6 +12,8 @@ It includes:
 - SecureAgentConfig as a context provider for easy setup
 """
 
+from __future__ import annotations
+
 import asyncio
 import contextlib
 import json
@@ -85,6 +87,7 @@ class IntegrityLabel(str, Enum):
     UNTRUSTED = "untrusted"
 
     def __str__(self) -> str:
+        """Return the string value of the integrity label."""
         return self.value
 
 
@@ -103,6 +106,7 @@ class ConfidentialityLabel(str, Enum):
     USER_IDENTITY = "user_identity"
 
     def __str__(self) -> str:
+        """Return the string value of the confidentiality label."""
         return self.value
 
 
@@ -118,7 +122,7 @@ class ContentLabel(SerializationMixin):
     Examples:
         .. code-block:: python
 
-            from agent_framework import ContentLabel, IntegrityLabel, ConfidentialityLabel
+            from agent_framework.security import ContentLabel, IntegrityLabel, ConfidentialityLabel
 
             # Create a label for trusted public content
             label = ContentLabel(integrity=IntegrityLabel.TRUSTED, confidentiality=ConfidentialityLabel.PUBLIC)
@@ -161,6 +165,7 @@ class ContentLabel(SerializationMixin):
         return self.confidentiality == ConfidentialityLabel.PUBLIC
 
     def __repr__(self) -> str:
+        """Return a debug representation of the content label."""
         return f"ContentLabel(integrity={self.integrity}, confidentiality={self.confidentiality})"
 
     def to_dict(self, *, exclude: set[str] | None = None, exclude_none: bool = True) -> dict[str, Any]:
@@ -180,7 +185,7 @@ class ContentLabel(SerializationMixin):
         /,
         *,
         dependencies: MutableMapping[str, Any] | None = None,
-    ) -> "ContentLabel":
+    ) -> ContentLabel:
         """Create ContentLabel from dictionary."""
         del dependencies
         return cls(
@@ -207,7 +212,7 @@ def combine_labels(*labels: ContentLabel) -> ContentLabel:
     Examples:
         .. code-block:: python
 
-            from agent_framework import ContentLabel, IntegrityLabel, ConfidentialityLabel, combine_labels
+            from agent_framework.security import ContentLabel, IntegrityLabel, ConfidentialityLabel, combine_labels
 
             label1 = ContentLabel(IntegrityLabel.TRUSTED, ConfidentialityLabel.PUBLIC)
             label2 = ContentLabel(IntegrityLabel.UNTRUSTED, ConfidentialityLabel.PRIVATE)
@@ -268,7 +273,7 @@ def check_confidentiality_allowed(
     Examples:
         .. code-block:: python
 
-            from agent_framework import ContentLabel, ConfidentialityLabel, check_confidentiality_allowed
+            from agent_framework.security import ContentLabel, ConfidentialityLabel, check_confidentiality_allowed
 
             # PUBLIC data can be written anywhere
             public_label = ContentLabel(confidentiality=ConfidentialityLabel.PUBLIC)
@@ -310,7 +315,7 @@ class ContentVariableStore:
     Examples:
         .. code-block:: python
 
-            from agent_framework import ContentVariableStore, ContentLabel, IntegrityLabel
+            from agent_framework.security import ContentVariableStore, ContentLabel, IntegrityLabel
 
             store = ContentVariableStore()
 
@@ -403,7 +408,7 @@ class VariableReferenceContent:
     Examples:
         .. code-block:: python
 
-            from agent_framework import VariableReferenceContent, ContentLabel, IntegrityLabel
+            from agent_framework.security import VariableReferenceContent, ContentLabel, IntegrityLabel
 
             label = ContentLabel(integrity=IntegrityLabel.UNTRUSTED)
             ref = VariableReferenceContent(variable_id="var_abc123", label=label, description="External API response")
@@ -428,6 +433,7 @@ class VariableReferenceContent:
         self.type: str = "variable_reference"
 
     def __repr__(self) -> str:
+        """Return a debug representation of the variable reference."""
         desc = f", description='{self.description}'" if self.description else ""
         return f"VariableReferenceContent(variable_id='{self.variable_id}'{desc})"
 
@@ -455,7 +461,7 @@ class VariableReferenceContent:
         return result
 
     @classmethod
-    def from_dict(cls, data: dict[str, Any]) -> "VariableReferenceContent":
+    def from_dict(cls, data: dict[str, Any]) -> VariableReferenceContent:
         """Create VariableReferenceContent from dictionary."""
         # Accept both "security_label" (preferred) and "label" (legacy) keys
         label_data = data.get("security_label") or data.get("label")
@@ -490,7 +496,7 @@ class LabeledMessage(Message):
     Examples:
         .. code-block:: python
 
-            from agent_framework import LabeledMessage, ContentLabel, IntegrityLabel
+            from agent_framework.security import LabeledMessage, ContentLabel, IntegrityLabel
 
             # User message is always TRUSTED
             user_msg = LabeledMessage(
@@ -591,6 +597,7 @@ class LabeledMessage(Message):
         return self.security_label.is_trusted()
 
     def __repr__(self) -> str:
+        """Return a debug representation of the labeled message."""
         return (
             f"LabeledMessage(role='{self.role}', "
             f"label={self.security_label.integrity.value}/{self.security_label.confidentiality.value})"
@@ -619,7 +626,7 @@ class LabeledMessage(Message):
         /,
         *,
         dependencies: MutableMapping[str, Any] | None = None,
-    ) -> "LabeledMessage":
+    ) -> LabeledMessage:
         """Create LabeledMessage from dictionary."""
         del dependencies
         source_labels: list[ContentLabel] | None = None
@@ -636,7 +643,7 @@ class LabeledMessage(Message):
         )
 
     @classmethod
-    def from_message(cls, message: dict[str, Any], index: int | None = None) -> "LabeledMessage":
+    def from_message(cls, message: dict[str, Any], index: int | None = None) -> LabeledMessage:
         """Create a LabeledMessage from a standard message dict.
 
         This is a convenience method to wrap existing messages with labels.
@@ -824,7 +831,9 @@ class LabelTrackingFunctionMiddleware(FunctionMiddleware):
     Examples:
         .. code-block:: python
 
-            from agent_framework import Agent, LabelTrackingFunctionMiddleware
+            from agent_framework import Agent
+
+            from agent_framework.security import LabelTrackingFunctionMiddleware
 
             # Create agent with automatic hiding enabled
             middleware = LabelTrackingFunctionMiddleware(
@@ -1605,7 +1614,9 @@ class PolicyEnforcementFunctionMiddleware(FunctionMiddleware):
     Examples:
         .. code-block:: python
 
-            from agent_framework import Agent, PolicyEnforcementFunctionMiddleware
+            from agent_framework import Agent
+
+            from agent_framework.security import PolicyEnforcementFunctionMiddleware
 
             # Create policy enforcement middleware
             policy = PolicyEnforcementFunctionMiddleware(allow_untrusted_tools={"search_web", "get_news"})
@@ -2000,7 +2011,9 @@ class SecureAgentConfig(ContextProvider):
     Examples:
         .. code-block:: python
 
-            from agent_framework import Agent, SecureAgentConfig
+            from agent_framework import Agent
+
+            from agent_framework.security import SecureAgentConfig
 
             # Create security configuration (also a context provider)
             security = SecureAgentConfig(
@@ -2029,7 +2042,7 @@ class SecureAgentConfig(ContextProvider):
         approval_on_violation: bool = False,
         enable_audit_log: bool = True,
         enable_policy_enforcement: bool = True,
-        quarantine_chat_client: "SupportsChatGetResponse | None" = None,
+        quarantine_chat_client: SupportsChatGetResponse | None = None,
         source_id: str | None = None,
     ) -> None:
         """Initialize secure agent configuration.
@@ -2162,7 +2175,7 @@ class SecureAgentConfig(ContextProvider):
         """
         return self.label_tracker.list_variables()
 
-    def get_quarantine_client(self) -> "SupportsChatGetResponse | None":
+    def get_quarantine_client(self) -> SupportsChatGetResponse | None:
         """Get the quarantine chat client.
 
         Returns:
@@ -2179,10 +2192,10 @@ class SecureAgentConfig(ContextProvider):
 _global_variable_store = ContentVariableStore()
 
 # Global quarantine chat client (set via set_quarantine_client or SecureAgentConfig)
-_quarantine_chat_client: "SupportsChatGetResponse | None" = None
+_quarantine_chat_client: SupportsChatGetResponse | None = None
 
 
-def set_quarantine_client(client: "SupportsChatGetResponse | None") -> None:
+def set_quarantine_client(client: SupportsChatGetResponse | None) -> None:
     """Set the global quarantine chat client.
 
     This client will be used by quarantined_llm to make actual LLM calls
@@ -2196,7 +2209,7 @@ def set_quarantine_client(client: "SupportsChatGetResponse | None") -> None:
         .. code-block:: python
 
             from agent_framework.openai import OpenAIChatClient
-            from agent_framework import set_quarantine_client
+            from agent_framework.security import set_quarantine_client
             from azure.identity import AzureCliCredential
 
             # Create a dedicated client for quarantine operations
@@ -2215,7 +2228,7 @@ def set_quarantine_client(client: "SupportsChatGetResponse | None") -> None:
         logger.info("Quarantine chat client cleared")
 
 
-def get_quarantine_client() -> "SupportsChatGetResponse | None":
+def get_quarantine_client() -> SupportsChatGetResponse | None:
     """Get the current quarantine chat client.
 
     Returns:
@@ -2672,7 +2685,7 @@ def store_untrusted_content(
     Examples:
         .. code-block:: python
 
-            from agent_framework import store_untrusted_content, ContentLabel, IntegrityLabel
+            from agent_framework.security import store_untrusted_content, ContentLabel, IntegrityLabel
 
             # Store external API response
             external_data = get_external_api_response()
@@ -2731,7 +2744,9 @@ def get_security_tools() -> list[FunctionTool]:
     Examples:
         .. code-block:: python
 
-            from agent_framework import Agent, get_security_tools
+            from agent_framework import Agent
+
+            from agent_framework.security import get_security_tools
 
             agent = Agent(
                 chat_client=client,

python/packages/core/tests/core/test_function_invocation_logic.py

@@ -37,6 +37,18 @@ def _group_id(message: Message) -> str | None:
     return value if isinstance(value, str) else None
 
 
+def _build_approved_tool_roundtrip(
+    *,
+    call_id: str,
+    approval_id: str,
+    tool_name: str,
+) -> tuple[Content, Content, Content]:
+    function_call = Content.from_function_call(call_id=call_id, name=tool_name, arguments="{}")
+    approval_request = Content.from_function_approval_request(id=approval_id, function_call=function_call)
+    approval_response = approval_request.to_function_approval_response(approved=True)
+    return function_call, approval_request, approval_response
+
+
 async def test_base_client_with_function_calling(chat_client_base: SupportsChatGetResponse):
     exec_counter = 0
 
@@ -2008,6 +2020,108 @@ def test_is_hosted_tool_approval_without_server_label():
     assert _is_hosted_tool_approval("not a content") is False
 
 
+def test_replace_approval_contents_with_results_uses_result_call_ids_without_placeholders() -> None:
+    from agent_framework._tools import _collect_approval_responses, _replace_approval_contents_with_results
+
+    call_one, request_one, response_one = _build_approved_tool_roundtrip(
+        call_id="call_1", approval_id="approval_1", tool_name="first_tool"
+    )
+    call_two, request_two, response_two = _build_approved_tool_roundtrip(
+        call_id="call_2", approval_id="approval_2", tool_name="second_tool"
+    )
+
+    messages = [
+        Message(role="assistant", contents=[call_one, request_one, call_two, request_two]),
+        Message(role="user", contents=[response_one, response_two]),
+    ]
+
+    _replace_approval_contents_with_results(
+        messages,
+        _collect_approval_responses(messages),
+        [
+            Content.from_function_result(call_id="call_2", result="second result"),
+            Content.from_function_result(call_id="call_1", result="first result"),
+        ],
+    )
+
+    assert len(messages) == 2
+    assert messages[0].contents == [call_one, call_two]
+    assert messages[1].role == "tool"
+    assert [(content.call_id, content.result) for content in messages[1].contents] == [
+        ("call_1", "first result"),
+        ("call_2", "second result"),
+    ]
+
+
+def test_replace_approval_contents_with_results_uses_result_call_ids_for_placeholders() -> None:
+    from agent_framework._tools import _collect_approval_responses, _replace_approval_contents_with_results
+
+    call_one, request_one, response_one = _build_approved_tool_roundtrip(
+        call_id="call_1", approval_id="approval_1", tool_name="first_tool"
+    )
+    call_two, request_two, response_two = _build_approved_tool_roundtrip(
+        call_id="call_2", approval_id="approval_2", tool_name="second_tool"
+    )
+
+    messages = [
+        Message(role="assistant", contents=[call_one, request_one, call_two, request_two]),
+        Message(
+            role="tool",
+            contents=[
+                Content.from_function_result(call_id="call_1", result="[APPROVAL_PENDING] first placeholder"),
+                Content.from_function_result(call_id="call_2", result="[APPROVAL_PENDING] second placeholder"),
+            ],
+        ),
+        Message(role="user", contents=[response_one, response_two]),
+    ]
+
+    _replace_approval_contents_with_results(
+        messages,
+        _collect_approval_responses(messages),
+        [
+            Content.from_function_result(call_id="call_2", result="second result"),
+            Content.from_function_result(call_id="call_1", result="first result"),
+        ],
+    )
+
+    assert len(messages) == 2
+    assert messages[0].contents == [call_one, call_two]
+    assert [(content.call_id, content.result) for content in messages[1].contents] == [
+        ("call_1", "first result"),
+        ("call_2", "second result"),
+    ]
+
+
+def test_replace_approval_contents_with_results_skips_results_without_call_id() -> None:
+    from agent_framework._tools import _collect_approval_responses, _replace_approval_contents_with_results
+
+    call_one, request_one, response_one = _build_approved_tool_roundtrip(
+        call_id="call_1", approval_id="approval_1", tool_name="first_tool"
+    )
+
+    messages = [
+        Message(role="assistant", contents=[call_one, request_one]),
+        Message(
+            role="tool",
+            contents=[Content.from_function_result(call_id="call_1", result="[APPROVAL_PENDING] placeholder")],
+        ),
+        Message(role="user", contents=[response_one]),
+    ]
+
+    _replace_approval_contents_with_results(
+        messages,
+        _collect_approval_responses(messages),
+        [
+            Content.from_function_result(call_id=None, result="ignored result"),
+            Content.from_function_result(call_id="call_1", result="first result"),
+        ],
+    )
+
+    assert len(messages) == 2
+    assert messages[0].contents == [call_one]
+    assert [(content.call_id, content.result) for content in messages[1].contents] == [("call_1", "first result")]
+
+
 async def test_mixed_local_and_hosted_approval_flow(chat_client_base: SupportsChatGetResponse):
     """Test that mixed local + hosted MCP approvals are handled correctly.
 

python/packages/core/tests/test_security.py

@@ -7,13 +7,15 @@ import json
 import pytest
 from pydantic import BaseModel
 
-from agent_framework import (
+from agent_framework import ExperimentalFeature, FunctionInvocationContext, FunctionMiddleware
+from agent_framework._middleware import FunctionMiddlewarePipeline, MiddlewareTermination
+from agent_framework._tools import FunctionTool, _auto_invoke_function, normalize_function_invocation_configuration
+from agent_framework._types import Content
+from agent_framework.security import (
     ConfidentialityLabel,
     ContentLabel,
     ContentVariableStore,
-    ExperimentalFeature,
-    FunctionInvocationContext,
-    FunctionMiddleware,
+    InspectVariableInput,
     IntegrityLabel,
     LabeledMessage,
     LabelTrackingFunctionMiddleware,
@@ -23,10 +25,6 @@ from agent_framework import (
     combine_labels,
     store_untrusted_content,
 )
-from agent_framework._middleware import FunctionMiddlewarePipeline, MiddlewareTermination
-from agent_framework._security import InspectVariableInput
-from agent_framework._tools import FunctionTool, _auto_invoke_function, normalize_function_invocation_configuration
-from agent_framework._types import Content
 
 
 class TestContentLabel:
@@ -840,7 +838,7 @@ class TestAutomaticHiding:
         context = FunctionInvocationContext(function=mock_function, arguments=args)
 
         async def next_fn():
-            from agent_framework._security import get_current_middleware
+            from agent_framework.security import get_current_middleware
 
             # Should be able to access middleware from thread-local
             current = get_current_middleware()
@@ -893,7 +891,7 @@ class TestSecureAgentConfig:
 
     def test_create_config_defaults(self):
         """Test creating config with default values."""
-        from agent_framework import SecureAgentConfig
+        from agent_framework.security import SecureAgentConfig
 
         config = SecureAgentConfig()
 
@@ -905,7 +903,7 @@ class TestSecureAgentConfig:
 
     def test_create_config_with_options(self):
         """Test creating config with custom options."""
-        from agent_framework import SecureAgentConfig
+        from agent_framework.security import SecureAgentConfig
 
         config = SecureAgentConfig(
             auto_hide_untrusted=True,
@@ -925,7 +923,7 @@ class TestSecureAgentConfig:
 
     def test_get_tools_returns_security_tools(self):
         """Test that get_tools returns quarantined_llm and inspect_variable."""
-        from agent_framework import SecureAgentConfig
+        from agent_framework.security import SecureAgentConfig
 
         config = SecureAgentConfig()
         tools = config.get_tools()
@@ -937,7 +935,7 @@ class TestSecureAgentConfig:
 
     def test_get_instructions_returns_string(self):
         """Test that get_instructions returns instruction text."""
-        from agent_framework import SECURITY_TOOL_INSTRUCTIONS, SecureAgentConfig
+        from agent_framework.security import SECURITY_TOOL_INSTRUCTIONS, SecureAgentConfig
 
         config = SecureAgentConfig()
         instructions = config.get_instructions()
@@ -950,7 +948,7 @@ class TestSecureAgentConfig:
 
     def test_inspect_variable_uses_generic_approval_mode(self):
         """Test that inspect_variable does not require approval (context tainting handles security)."""
-        from agent_framework import get_security_tools
+        from agent_framework.security import get_security_tools
 
         inspect_variable = next(tool for tool in get_security_tools() if tool.name == "inspect_variable")
         assert inspect_variable.approval_mode == "never_require"
@@ -962,7 +960,7 @@ class TestGetSecurityTools:
 
     def test_get_security_tools_from_module(self):
         """Test importing get_security_tools from agent_framework."""
-        from agent_framework import get_security_tools
+        from agent_framework.security import get_security_tools
 
         tools = get_security_tools()
         assert len(tools) == 2
@@ -995,7 +993,7 @@ class TestQuarantinedLLMWithVariableIds:
     @pytest.mark.asyncio
     async def test_quarantined_llm_with_single_variable_id(self, middleware_with_store):
         """Test quarantined_llm retrieves content from variable store."""
-        from agent_framework import quarantined_llm
+        from agent_framework.security import quarantined_llm
 
         # Store a variable
         store = middleware_with_store.get_variable_store()
@@ -1013,7 +1011,7 @@ class TestQuarantinedLLMWithVariableIds:
     @pytest.mark.asyncio
     async def test_quarantined_llm_with_multiple_variable_ids(self, middleware_with_store):
         """Test quarantined_llm retrieves multiple variables."""
-        from agent_framework import quarantined_llm
+        from agent_framework.security import quarantined_llm
 
         # Store multiple variables
         store = middleware_with_store.get_variable_store()
@@ -1033,7 +1031,7 @@ class TestQuarantinedLLMWithVariableIds:
     @pytest.mark.asyncio
     async def test_quarantined_llm_with_unknown_variable_id(self, middleware_with_store):
         """Test quarantined_llm handles unknown variable IDs gracefully."""
-        from agent_framework import quarantined_llm
+        from agent_framework.security import quarantined_llm
 
         # Call with non-existent variable ID
         result = await quarantined_llm(prompt="Process this", variable_ids=["var_nonexistent"])
@@ -1046,7 +1044,7 @@ class TestQuarantinedLLMWithVariableIds:
     @pytest.mark.asyncio
     async def test_quarantined_llm_without_variable_ids(self, middleware_with_store):
         """Test quarantined_llm works with labelled_data instead of variable_ids."""
-        from agent_framework import quarantined_llm
+        from agent_framework.security import quarantined_llm
 
         result = await quarantined_llm(
             prompt="Process this data",
@@ -1064,7 +1062,7 @@ class TestQuarantinedLLMWithVariableIds:
     @pytest.mark.asyncio
     async def test_quarantined_llm_with_legacy_label_key(self, middleware_with_store):
         """Test quarantined_llm accepts legacy 'label' key for backward compatibility."""
-        from agent_framework import quarantined_llm
+        from agent_framework.security import quarantined_llm
 
         result = await quarantined_llm(
             prompt="Process this data",
@@ -1085,7 +1083,7 @@ class TestMiddlewareSetCurrent:
 
     def test_set_and_clear_current(self):
         """Test setting and clearing thread-local middleware reference."""
-        from agent_framework._security import get_current_middleware
+        from agent_framework.security import get_current_middleware
 
         # Initially no middleware
         assert get_current_middleware() is None
@@ -1103,7 +1101,7 @@ class TestMiddlewareSetCurrent:
 
     def test_set_current_overwrites_previous(self):
         """Test that setting current overwrites previous middleware."""
-        from agent_framework._security import get_current_middleware
+        from agent_framework.security import get_current_middleware
 
         middleware1 = LabelTrackingFunctionMiddleware()
         middleware2 = LabelTrackingFunctionMiddleware()
@@ -1375,7 +1373,7 @@ class TestLabeledMessage:
 
     def test_create_user_message_defaults_to_trusted(self):
         """Test that user messages are TRUSTED by default."""
-        from agent_framework import LabeledMessage
+        from agent_framework.security import LabeledMessage
 
         msg = LabeledMessage(role="user", content="Hello!")
         assert msg.role == "user"
@@ -1384,14 +1382,14 @@ class TestLabeledMessage:
 
     def test_create_system_message_defaults_to_trusted(self):
         """Test that system messages are TRUSTED by default."""
-        from agent_framework import LabeledMessage
+        from agent_framework.security import LabeledMessage
 
         msg = LabeledMessage(role="system", content="You are an assistant.")
         assert msg.security_label.integrity == IntegrityLabel.TRUSTED
 
     def test_create_tool_message_defaults_to_untrusted(self):
         """Test that tool messages are UNTRUSTED by default."""
-        from agent_framework import LabeledMessage
+        from agent_framework.security import LabeledMessage
 
         msg = LabeledMessage(role="tool", content="External API result")
         assert msg.security_label.integrity == IntegrityLabel.UNTRUSTED
@@ -1399,14 +1397,14 @@ class TestLabeledMessage:
 
     def test_create_assistant_message_no_sources(self):
         """Test assistant message without sources defaults to TRUSTED."""
-        from agent_framework import LabeledMessage
+        from agent_framework.security import LabeledMessage
 
         msg = LabeledMessage(role="assistant", content="I'll help you.")
         assert msg.security_label.integrity == IntegrityLabel.TRUSTED
 
     def test_create_assistant_message_with_untrusted_source(self):
         """Test assistant message inherits UNTRUSTED from sources."""
-        from agent_framework import LabeledMessage
+        from agent_framework.security import LabeledMessage
 
         untrusted_source = ContentLabel(integrity=IntegrityLabel.UNTRUSTED)
         msg = LabeledMessage(role="assistant", content="Based on the data...", source_labels=[untrusted_source])
@@ -1414,7 +1412,7 @@ class TestLabeledMessage:
 
     def test_explicit_label_overrides_inference(self):
         """Test that explicit label overrides role-based inference."""
-        from agent_framework import LabeledMessage
+        from agent_framework.security import LabeledMessage
 
         explicit_label = ContentLabel(integrity=IntegrityLabel.UNTRUSTED, confidentiality=ConfidentialityLabel.PRIVATE)
         msg = LabeledMessage(
@@ -1427,7 +1425,7 @@ class TestLabeledMessage:
 
     def test_message_serialization(self):
         """Test LabeledMessage serialization to dict."""
-        from agent_framework import LabeledMessage
+        from agent_framework.security import LabeledMessage
 
         msg = LabeledMessage(role="user", content="Hello", message_index=5, metadata={"key": "value"})
 
@@ -1439,7 +1437,7 @@ class TestLabeledMessage:
 
     def test_message_deserialization(self):
         """Test LabeledMessage deserialization from dict."""
-        from agent_framework import LabeledMessage
+        from agent_framework.security import LabeledMessage
 
         data = {
             "role": "tool",
@@ -1455,7 +1453,7 @@ class TestLabeledMessage:
 
     def test_from_message_convenience_method(self):
         """Test creating LabeledMessage from a standard message dict."""
-        from agent_framework import LabeledMessage
+        from agent_framework.security import LabeledMessage
 
         standard_msg = {"role": "user", "content": "What's the weather?"}
         labeled = LabeledMessage.from_message(standard_msg, index=0)
@@ -1534,8 +1532,7 @@ class TestQuarantinedLLM:
     @pytest.mark.asyncio
     async def test_quarantined_llm_returns_response(self):
         """Test that quarantined_llm returns a plain response dict."""
-        from agent_framework import LabelTrackingFunctionMiddleware, quarantined_llm
-        from agent_framework._security import _current_middleware
+        from agent_framework.security import LabelTrackingFunctionMiddleware, _current_middleware, quarantined_llm
 
         middleware = LabelTrackingFunctionMiddleware()
 
@@ -1560,8 +1557,7 @@ class TestQuarantinedLLM:
     @pytest.mark.asyncio
     async def test_quarantined_llm_trusted_input(self):
         """Test quarantined_llm with TRUSTED input returns response directly."""
-        from agent_framework import LabelTrackingFunctionMiddleware, quarantined_llm
-        from agent_framework._security import _current_middleware
+        from agent_framework.security import LabelTrackingFunctionMiddleware, _current_middleware, quarantined_llm
 
         middleware = LabelTrackingFunctionMiddleware()
 
@@ -1587,8 +1583,7 @@ class TestQuarantinedLLM:
     @pytest.mark.asyncio
     async def test_quarantined_llm_multiple_variables(self):
         """Test that quarantined_llm handles multiple variables correctly."""
-        from agent_framework import LabelTrackingFunctionMiddleware, quarantined_llm
-        from agent_framework._security import _current_middleware
+        from agent_framework.security import LabelTrackingFunctionMiddleware, _current_middleware, quarantined_llm
 
         middleware = LabelTrackingFunctionMiddleware()
 
@@ -1608,7 +1603,7 @@ class TestQuarantinedLLM:
 
     def test_quarantined_llm_declares_source_integrity(self):
         """Test that quarantined_llm declares source_integrity='untrusted'."""
-        from agent_framework import get_security_tools
+        from agent_framework.security import get_security_tools
 
         q_llm = next(tool for tool in get_security_tools() if tool.name == "quarantined_llm")
         assert q_llm.additional_properties.get("source_integrity") == "untrusted"
@@ -1620,7 +1615,7 @@ class TestQuarantineClient:
 
     def test_set_and_get_quarantine_client(self):
         """Test setting and getting the quarantine client."""
-        from agent_framework import get_quarantine_client, set_quarantine_client
+        from agent_framework.security import get_quarantine_client, set_quarantine_client
 
         # Initially should be None (or whatever state it's in)
         # Clear it first
@@ -1643,7 +1638,7 @@ class TestQuarantineClient:
 
     def test_secure_agent_config_sets_quarantine_client(self):
         """Test that SecureAgentConfig sets the quarantine client."""
-        from agent_framework import SecureAgentConfig, get_quarantine_client, set_quarantine_client
+        from agent_framework.security import SecureAgentConfig, get_quarantine_client, set_quarantine_client
 
         # Clear any existing client
         set_quarantine_client(None)
@@ -1669,7 +1664,7 @@ class TestQuarantineClient:
 
     def test_secure_agent_config_without_quarantine_client(self):
         """Test SecureAgentConfig without quarantine client doesn't set one."""
-        from agent_framework import SecureAgentConfig, get_quarantine_client, set_quarantine_client
+        from agent_framework.security import SecureAgentConfig, get_quarantine_client, set_quarantine_client
 
         # Clear any existing client
         set_quarantine_client(None)
@@ -1688,14 +1683,14 @@ class TestQuarantineClient:
         """Test that quarantined_llm uses real client when available."""
         from unittest.mock import AsyncMock, MagicMock
 
-        from agent_framework import (
+        from agent_framework.security import (
             ContentLabel,
             IntegrityLabel,
             LabelTrackingFunctionMiddleware,
+            _current_middleware,
             quarantined_llm,
             set_quarantine_client,
         )
-        from agent_framework._security import _current_middleware
 
         # Clear any existing client
         set_quarantine_client(None)
@@ -1747,14 +1742,14 @@ class TestQuarantineClient:
     @pytest.mark.asyncio
     async def test_quarantined_llm_fallback_without_client(self):
         """Test that quarantined_llm falls back to placeholder without client."""
-        from agent_framework import (
+        from agent_framework.security import (
             ContentLabel,
             IntegrityLabel,
             LabelTrackingFunctionMiddleware,
+            _current_middleware,
             quarantined_llm,
             set_quarantine_client,
         )
-        from agent_framework._security import _current_middleware
 
         # Clear the client
         set_quarantine_client(None)
@@ -1785,14 +1780,14 @@ class TestQuarantineClient:
         """Test that quarantined_llm handles client errors gracefully."""
         from unittest.mock import AsyncMock, MagicMock
 
-        from agent_framework import (
+        from agent_framework.security import (
             ContentLabel,
             IntegrityLabel,
             LabelTrackingFunctionMiddleware,
+            _current_middleware,
             quarantined_llm,
             set_quarantine_client,
         )
-        from agent_framework._security import _current_middleware
 
         # Create a mock client that raises an error
         mock_client = MagicMock()
@@ -1822,14 +1817,14 @@ class TestQuarantineClient:
         """Test that quarantined_llm builds messages correctly with content."""
         from unittest.mock import AsyncMock, MagicMock
 
-        from agent_framework import (
+        from agent_framework.security import (
             ContentLabel,
             IntegrityLabel,
             LabelTrackingFunctionMiddleware,
+            _current_middleware,
             quarantined_llm,
             set_quarantine_client,
         )
-        from agent_framework._security import _current_middleware
 
         mock_response = MagicMock()
         mock_response.text = "Summary"
@@ -2517,63 +2512,63 @@ class TestCheckConfidentialityAllowed:
 
     def test_public_to_public_allowed(self):
         """Test PUBLIC data can be written to PUBLIC destination."""
-        from agent_framework import check_confidentiality_allowed
+        from agent_framework.security import check_confidentiality_allowed
 
         public_label = ContentLabel(confidentiality=ConfidentialityLabel.PUBLIC)
         assert check_confidentiality_allowed(public_label, ConfidentialityLabel.PUBLIC) is True
 
     def test_public_to_private_allowed(self):
         """Test PUBLIC data can be written to PRIVATE destination."""
-        from agent_framework import check_confidentiality_allowed
+        from agent_framework.security import check_confidentiality_allowed
 
         public_label = ContentLabel(confidentiality=ConfidentialityLabel.PUBLIC)
         assert check_confidentiality_allowed(public_label, ConfidentialityLabel.PRIVATE) is True
 
     def test_public_to_user_identity_allowed(self):
         """Test PUBLIC data can be written to USER_IDENTITY destination."""
-        from agent_framework import check_confidentiality_allowed
+        from agent_framework.security import check_confidentiality_allowed
 
         public_label = ContentLabel(confidentiality=ConfidentialityLabel.PUBLIC)
         assert check_confidentiality_allowed(public_label, ConfidentialityLabel.USER_IDENTITY) is True
 
     def test_private_to_public_blocked(self):
         """Test PRIVATE data cannot be written to PUBLIC destination."""
-        from agent_framework import check_confidentiality_allowed
+        from agent_framework.security import check_confidentiality_allowed
 
         private_label = ContentLabel(confidentiality=ConfidentialityLabel.PRIVATE)
         assert check_confidentiality_allowed(private_label, ConfidentialityLabel.PUBLIC) is False
 
     def test_private_to_private_allowed(self):
         """Test PRIVATE data can be written to PRIVATE destination."""
-        from agent_framework import check_confidentiality_allowed
+        from agent_framework.security import check_confidentiality_allowed
 
         private_label = ContentLabel(confidentiality=ConfidentialityLabel.PRIVATE)
         assert check_confidentiality_allowed(private_label, ConfidentialityLabel.PRIVATE) is True
 
     def test_private_to_user_identity_allowed(self):
         """Test PRIVATE data can be written to USER_IDENTITY destination."""
-        from agent_framework import check_confidentiality_allowed
+        from agent_framework.security import check_confidentiality_allowed
 
         private_label = ContentLabel(confidentiality=ConfidentialityLabel.PRIVATE)
         assert check_confidentiality_allowed(private_label, ConfidentialityLabel.USER_IDENTITY) is True
 
     def test_user_identity_to_public_blocked(self):
         """Test USER_IDENTITY data cannot be written to PUBLIC destination."""
-        from agent_framework import check_confidentiality_allowed
+        from agent_framework.security import check_confidentiality_allowed
 
         ui_label = ContentLabel(confidentiality=ConfidentialityLabel.USER_IDENTITY)
         assert check_confidentiality_allowed(ui_label, ConfidentialityLabel.PUBLIC) is False
 
     def test_user_identity_to_private_blocked(self):
         """Test USER_IDENTITY data cannot be written to PRIVATE destination."""
-        from agent_framework import check_confidentiality_allowed
+        from agent_framework.security import check_confidentiality_allowed
 
         ui_label = ContentLabel(confidentiality=ConfidentialityLabel.USER_IDENTITY)
         assert check_confidentiality_allowed(ui_label, ConfidentialityLabel.PRIVATE) is False
 
     def test_user_identity_to_user_identity_allowed(self):
         """Test USER_IDENTITY data can be written to USER_IDENTITY destination."""
-        from agent_framework import check_confidentiality_allowed
+        from agent_framework.security import check_confidentiality_allowed
 
         ui_label = ContentLabel(confidentiality=ConfidentialityLabel.USER_IDENTITY)
         assert check_confidentiality_allowed(ui_label, ConfidentialityLabel.USER_IDENTITY) is True