Python: information-flow control prompt injection defense (#5331)

* Python: Information-flow control based prompt injection defense (#5024) * fides integration * documentation * documentation * documentation * human-approval on policy violation * numenous hyena 'works' * IFC based implementation * minor edits in documentation * rebasing the branch and running the email example * Add security tests for IFC middleware * Fix Role.TOOL NameError in approval handling * tiered labelling scheme * 3 tier labelling scheme in middleware * Adapt security middleware to list[Content] tool results * Refactor SecureAgentConfig as context provider and address Copilot review comments * Update FIDES docs to reflect context provider pattern and update code for ContextProvider rename * Fix security examples: use OpenAIChatClient instead of non-existent AzureOpenAIChatClient * Address PR review: consolidate security modules, remove ContentLineage, update docs * remove unrelated files * remove comment from _tools.py and rename decision file * Fix CI failures: Bandit B110, broken md links, hosted approval passthrough * apply template to decision doc 0024 * minor fixes to decision doc 0024 --------- Co-authored-by: Aashish <t-akolluri@microsoft.com> * Python: follow up FIDES security flow (#5330) * Python: follow up FIDES security flow Refine the secure approval path, mark the security classes with the FIDES experimental feature label, and clean up the related docs/tests. Also fix workspace-level validation regressions uncovered while running the full Python check suite. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Python: remove FIDES GitHub MCP sample Drop the GitHub MCP security sample from the FIDES follow-up branch while keeping the remaining security docs and samples intact. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Address PR review: fix paths and update FIDES implementation (#5352) * 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: Address PR 5331 comments and track sesssion while calling Agent in email_security_example (#5446) * Address PR review: fix paths and update FIDES implementation * Address PR comments and add session tracking in email example in samples * Fix session creation and resolve merge conflict in docstring example * Resolve merge conflict in docstring example * Python: add test for empty-message pruning in approval result replacement (#5617) Adds test coverage for the second-pass logic in `_replace_approval_contents_with_results` that removes messages whose `contents` list becomes empty after first-pass content removal. Addresses review comment on PR #5331: https://github.com/microsoft/agent-framework/pull/5331#discussion_r3129039445 Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --------- Co-authored-by: shrutitople <shruti.tople@gmail.com> Co-authored-by: Aashish <t-akolluri@microsoft.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
2026-06-16 21:04:09 +08:00 · 2026-05-05 20:08:08 +02:00
parent 806075ae61
commit ddfbdf5c7a
15 changed files with 7978 additions and 31 deletions
@@ -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.)
@@ -48,6 +48,7 @@ class ExperimentalFeature(str, Enum):

    EVALS = "EVALS"
    FILE_HISTORY = "FILE_HISTORY"
+    FIDES = "FIDES"
    FUNCTIONAL_WORKFLOWS = "FUNCTIONAL_WORKFLOWS"
    HARNESS = "HARNESS"
    SKILLS = "SKILLS"
@@ -1448,6 +1448,8 @@ async def _auto_invoke_function(
    # non-declaration-only functions.

    tool: FunctionTool | None = None
+    approval_response: Content | None = None
+
    if function_call_content.type == "function_call":
        tool = tool_map.get(function_call_content.name)  # type: ignore[arg-type]
        # Tool should exist because _try_execute_function_calls validates this
@@ -1462,14 +1464,20 @@ async def _auto_invoke_function(
    else:
        # Note: Unapproved tools (approved=False) are handled in _replace_approval_contents_with_results
        # and never reach this function, so we only handle approved=True cases here.
-        inner_call = function_call_content.function_call  # type: ignore[attr-defined]
-        if inner_call.type != "function_call":  # type: ignore[union-attr]
+        approved_function_call = function_call_content.function_call  # type: ignore[attr-defined]
+        if (
+            approved_function_call is None
+            or approved_function_call.type != "function_call"
+            or approved_function_call.name is None
+        ):
            return function_call_content
-        tool = tool_map.get(inner_call.name)  # type: ignore[attr-defined, union-attr, arg-type]
+        tool = tool_map.get(approved_function_call.name)
        if tool is None:
            # we assume it is a hosted tool
            return function_call_content
-        function_call_content = inner_call  # type: ignore[assignment]
+
+        approval_response = function_call_content
+        function_call_content = approved_function_call

    parsed_args: dict[str, Any] = dict(function_call_content.parse_arguments() or {})

@@ -1546,32 +1554,56 @@ async def _auto_invoke_function(
        kwargs=runtime_kwargs.copy(),
    )

+    call_id = function_call_content.call_id
+    if call_id is None:
+        raise KeyError(f'Function "{function_call_content.name}" is missing call_id.')
+
+    # Always pass call_id to middleware for policy violation approval flow
+    middleware_context.metadata["call_id"] = call_id
+
+    # Pass through the original approval response so middleware can decide whether
+    # this replay corresponds to a middleware-specific approval flow.
+    if approval_response is not None:
+        middleware_context.metadata["approval_response"] = approval_response
+
    async def final_function_handler(context_obj: Any) -> Any:
        return await tool.invoke(
            arguments=context_obj.arguments,
            context=context_obj,
-            tool_call_id=function_call_content.call_id,
+            tool_call_id=call_id,
        )

    from ._middleware import MiddlewareTermination

    # MiddlewareTermination bubbles up to signal loop termination
    try:
-        function_result = await middleware_pipeline.execute(middleware_context, final_function_handler)
-        return Content.from_function_result(
-            call_id=function_call_content.call_id,  # type: ignore[arg-type]
-            result=function_result,
-            additional_properties=function_call_content.additional_properties,
+        function_result = await middleware_pipeline.execute(
+            context=middleware_context,
+            final_handler=final_function_handler,
        )
+
+        # Pass through function_approval_request directly (e.g., from security middleware)
+        if isinstance(function_result, Content) and function_result.type == "function_approval_request":
+            return function_result
+
+        return Content.from_function_result(call_id=call_id, result=function_result)
    except MiddlewareTermination as term_exc:
        # Re-raise to signal loop termination, but first capture any result set by middleware
        if middleware_context.result is not None:
-            # Store result in exception for caller to extract
-            term_exc.result = Content.from_function_result(
-                call_id=function_call_content.call_id,  # type: ignore[arg-type]
-                result=middleware_context.result,
-                additional_properties=function_call_content.additional_properties,
-            )
+            # Pass through function_approval_request directly (e.g., from security policy middleware)
+            # so the approval flow in _handle_function_call_results activates correctly.
+            if (
+                isinstance(middleware_context.result, Content)
+                and middleware_context.result.type == "function_approval_request"
+            ):
+                term_exc.result = middleware_context.result
+            else:
+                # Store result in exception for caller to extract
+                term_exc.result = Content.from_function_result(
+                    call_id=call_id,
+                    result=middleware_context.result,
+                    additional_properties=function_call_content.additional_properties,
+                )
        raise
    except UserInputRequiredException:
        raise
@@ -1877,12 +1909,24 @@ def _replace_approval_contents_with_results(
    fcc_todo: dict[str, Content],
    approved_function_results: list[Content],
 ) -> None:
-    """Replace approval request/response contents with function call/result contents in-place."""
+    """Replace approval request/response contents with function call/result contents in-place.
+
+    Also replaces placeholder tool results (marked with [APPROVAL_PENDING]) with actual results.
+    """
    from ._types import (
        Content,
    )

-    result_idx = 0
+    # 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 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()
+
    for msg in messages:
        # First pass - collect existing function call IDs to avoid duplicates
        existing_call_ids = {
@@ -1900,22 +1944,31 @@ def _replace_approval_contents_with_results(
                if _is_hosted_tool_approval(content):
                    continue
                # Don't add the function call if it already exists (would create duplicate)
-                if content.function_call.call_id in existing_call_ids:  # type: ignore[attr-defined, union-attr, operator]
+                if content.function_call is not None and content.function_call.call_id in existing_call_ids:
                    # Just mark for removal - the function call already exists
                    contents_to_remove.append(content_idx)
-                else:
+                elif content.function_call is not None:
                    # Put back the function call content only if it doesn't exist
-                    msg.contents[content_idx] = content.function_call  # type: ignore[attr-defined, assignment]
+                    msg.contents[content_idx] = content.function_call
            elif content.type == "function_approval_response":
                # Skip hosted tool approvals — they must pass through to the API unchanged
                if _is_hosted_tool_approval(content):
                    continue
-                if content.approved and content.id in fcc_todo:  # type: ignore[attr-defined]
-                    # Replace with the corresponding result
-                    if result_idx < len(approved_function_results):
-                        msg.contents[content_idx] = approved_function_results[result_idx]
-                        result_idx += 1
-                        msg.role = "tool"
+                if content.function_call is None or content.function_call.call_id is None:
+                    continue
+                call_id = content.function_call.call_id
+                if content.approved and content.id in fcc_todo:
+                    # Check if we already replaced a placeholder for this call_id
+                    if call_id in placeholders_replaced:
+                        # Placeholder was replaced - just remove the approval response
+                        contents_to_remove.append(content_idx)
+                    else:
+                        # No placeholder - replace approval response with result directly
+                        # This handles the original approval_mode="always_require" case
+                        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
                    # Use function_call.call_id (the function's ID), not content.id (approval's ID)
@@ -1924,11 +1977,31 @@ def _replace_approval_contents_with_results(
                        result="Error: Tool call invocation was rejected by user.",
                    )
                    msg.role = "tool"
+            elif content.type == "function_result":
+                # Check if this is a placeholder result that should be replaced
+                if (
+                    hasattr(content, "result")
+                    and isinstance(content.result, str)
+                    and "[APPROVAL_PENDING]" in content.result
+                    and content.call_id in result_by_call_id
+                ):
+                    # Replace placeholder with actual result
+                    msg.contents[content_idx] = result_by_call_id[content.call_id]
+                    placeholders_replaced.add(content.call_id)

-        # Remove approval requests that were duplicates (in reverse order to preserve indices)
+        # Remove contents marked for removal (in reverse order to preserve indices)
        for idx in reversed(contents_to_remove):
            msg.contents.pop(idx)

+    # Second pass: Remove messages that are now empty after content removal
+    # We need to iterate in reverse to safely remove by index
+    messages_to_remove: list[int] = []
+    for msg_idx, msg in enumerate(messages):
+        if not msg.contents:
+            messages_to_remove.append(msg_idx)
+    for msg_idx in reversed(messages_to_remove):
+        messages.pop(msg_idx)
+

 def _get_result_hooks_from_stream(stream: Any) -> list[Callable[[Any], Any]]:
    inner_stream = getattr(stream, "_inner_stream", None)
@@ -2595,3 +2668,7 @@ class FunctionInvocationLayer(Generic[OptionsCoT]):
            return ChatResponse.from_updates(updates, output_format_type=response_format)

        return ResponseStream(_stream(), finalizer=_finalize)
+
+
+# Alias for the @tool decorator, used by security tools and samples
+ai_function = tool
@@ -2121,7 +2121,7 @@ def _get_response_attributes(
        finish_reason = (
            getattr(response.raw_representation, "finish_reason", None) if response.raw_representation else None
        )
-    if finish_reason:
+    if isinstance(finish_reason, str) and finish_reason:
        attributes[OtelAttr.FINISH_REASONS] = json.dumps([finish_reason])
    if model := getattr(response, "model", None):
        attributes[OtelAttr.RESPONSE_MODEL] = model
@@ -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,162 @@ 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")]
+
+
+def test_replace_approval_contents_with_results_prunes_emptied_messages() -> None:
+    """Messages whose contents are fully consumed during the first pass should be removed.
+
+    When approval responses are paired with placeholder results, the responses are marked
+    for removal in the first pass. If a message contained only such responses, it ends up
+    with an empty `contents` list and the second pass should drop it from `messages`.
+    """
+    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"),
+            ],
+        ),
+        # This user message holds only approval_responses whose placeholders are replaced
+        # in the tool message above, so every content here is marked for removal and the
+        # message itself becomes empty -> it must be pruned by the second pass.
+        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_1", result="first result"),
+            Content.from_function_result(call_id="call_2", result="second result"),
+        ],
+    )
+
+    # The now-empty user message should have been pruned, leaving just the assistant
+    # message and the tool message with the resolved results.
+    assert len(messages) == 2
+    assert messages[0].role == "assistant"
+    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"),
+    ]
+    # Sanity-check: no leftover empty messages.
+    assert all(msg.contents for msg in messages)
+
+
 async def test_mixed_local_and_hosted_approval_flow(chat_client_base: SupportsChatGetResponse):
    """Test that mixed local + hosted MCP approvals are handled correctly.

@@ -744,6 +744,15 @@ class AgentFrameworkExecutor:
                                            )
                                            continue

+                                        # Extract policy_violation info if present (from security middleware)
+                                        policy_violation_data = content_dict.get("policy_violation")
+                                        approval_additional_props: dict[str, Any] | None = None
+                                        if isinstance(policy_violation_data, dict):
+                                            approval_additional_props = {
+                                                "policy_violation": True,
+                                                **policy_violation_data,
+                                            }
+
                                        # Reconstruct function_call from server-stored data
                                        function_call = Content.from_function_call(
                                            call_id=stored_fc["call_id"],
@@ -756,14 +765,16 @@ class AgentFrameworkExecutor:
                                            approved,
                                            id=request_id,
                                            function_call=function_call,
+                                            additional_properties=approval_additional_props,
                                        )
                                        contents.append(approval_response)
                                        logger.info(
                                            "Validated FunctionApprovalResponseContent: id=%s, "
-                                            "approved=%s, function=%s",
+                                            "approved=%s, function=%s, policy_violation=%s",
                                            request_id,
                                            approved,
                                            stored_fc["name"],
+                                            approval_additional_props is not None,
                                        )
                                    except ImportError:
                                        logger.warning(
@@ -1747,7 +1747,7 @@ class MessageMapper:
                # Fallback to direct access if parse_arguments doesn't exist
                arguments = getattr(content.function_call, "arguments", {})

-        return {
+        result = {
            "type": "response.function_approval.requested",
            "request_id": getattr(content, "id", "unknown"),
            "function_call": {
@@ -1760,6 +1760,17 @@ class MessageMapper:
            "sequence_number": self._next_sequence(context),
        }

+        # Include policy violation details if present (from security middleware)
+        additional_props = cast(dict[str, Any] | None, getattr(content, "additional_properties", None))
+        if additional_props and isinstance(additional_props, dict) and additional_props.get("policy_violation"):
+            result["policy_violation"] = {
+                "reason": additional_props.get("reason", "Policy violation detected"),
+                "violation_type": additional_props.get("violation_type"),
+                "context_label": additional_props.get("context_label"),
+            }
+
+        return result
+
    async def _map_approval_response_content(self, content: Any, context: dict[str, Any]) -> dict[str, Any]:
        """Map FunctionApprovalResponseContent to custom event."""
        return {