agent/agent-framework
2
0
Fork 0
mirror of https://github.com/microsoft/agent-framework.git synced 2026-06-16 21:04:09 +08:00
Code Issues Packages Projects Releases Wiki Activity

Python: Information-flow control based prompt injection defense (#5024)

Browse Source 
* 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>
This commit is contained in:
shrutitople
2026-04-16 12:15:15 +01:00
committed by eavanvalkenburg
Unverified
parent 6582926af5
commit 8a08776a32
13 changed files with 9112 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/packages/core/agent_framework/__init__.py
+38
View File
@@ -109,6 +109,25 @@ from ._sessions import (
SessionContext,
register_state_type,
)
from ._security import (
ContentLabel,
IntegrityLabel,
ConfidentialityLabel,
ContentVariableStore,
LabeledMessage,
LabelTrackingFunctionMiddleware,
PolicyEnforcementFunctionMiddleware,
SECURITY_TOOL_INSTRUCTIONS,
SecureAgentConfig,
VariableReferenceContent,
check_confidentiality_allowed,
combine_labels,
get_quarantine_client,
get_security_tools,
quarantined_llm,
set_quarantine_client,
store_untrusted_content,
)
from ._settings import SecretString, load_settings
from ._skills import (
Skill,
@@ -130,6 +149,7 @@ from ._tools import (
FunctionInvocationLayer,
FunctionTool,
ToolTypes,
ai_function,
normalize_function_invocation_configuration,
tool,
)
@@ -307,7 +327,10 @@ __all__ = [
"CheckpointStorage",
"CompactionProvider",
"CompactionStrategy",
"ConfidentialityLabel",
"Content",
"ContentLabel",
"ContentVariableStore",
"ContextProvider",
"ContinuationToken",
"ConversationSplit",
@@ -351,6 +374,9 @@ __all__ = [
"InMemoryCheckpointStorage",
"InMemoryHistoryProvider",
"InProcRunnerContext",
"IntegrityLabel",
"LabelTrackingFunctionMiddleware",
"LabeledMessage",
"LocalEvaluator",
"MCPStdioTool",
"MCPStreamableHTTPTool",
@@ -358,6 +384,7 @@ __all__ = [
"Message",
"MiddlewareException",
"MiddlewareTermination",
"PolicyEnforcementFunctionMiddleware",
"MiddlewareType",
"MiddlewareTypes",
"OuterFinalT",
@@ -370,6 +397,8 @@ __all__ = [
"RunContext",
"Runner",
"RunnerContext",
"SECURITY_TOOL_INSTRUCTIONS",
"SecureAgentConfig",
"SecretString",
"SelectiveToolCallCompactionStrategy",
"SessionContext",
@@ -407,6 +436,7 @@ __all__ = [
"UsageDetails",
"UserInputRequiredException",
"ValidationTypeEnum",
"VariableReferenceContent",
"Workflow",
"WorkflowAgent",
"WorkflowBuilder",
@@ -428,10 +458,13 @@ __all__ = [
"WorkflowViz",
"__version__",
"add_usage_details",
"ai_function",
"agent_middleware",
"annotate_message_groups",
"apply_compaction",
"chat_middleware",
"check_confidentiality_allowed",
"combine_labels",
"create_edge_runner",
"detect_media_type_from_base64",
"evaluate_agent",
@@ -439,7 +472,9 @@ __all__ = [
"evaluator",
"executor",
"function_middleware",
"get_quarantine_client",
"get_run_context",
"get_security_tools",
"handler",
"included_messages",
"included_token_count",
@@ -452,10 +487,13 @@ __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/_security.py
+2777
View File
File diff suppressed because it is too large Load Diff
python/packages/core/agent_framework/_tools.py
+89 -15
View File
@@ -1448,6 +1448,9 @@ async def _auto_invoke_function(
# non-declaration-only functions.
tool: FunctionTool | None = None
# Track if this is a re-invocation after policy violation approval
policy_approval_granted = False
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
@@ -1469,7 +1472,14 @@ async def _auto_invoke_function(
if tool is None:
# we assume it is a hosted tool
return function_call_content
function_call_content = inner_call # type: ignore[assignment]
# Check if this is an approval for a policy violation
# The additional_properties may contain {"policy_violation": True, ...} or just truthy value
approval_props = getattr(function_call_content, "additional_properties", None) or {}
if approval_props.get("policy_violation"):
policy_approval_granted = True
function_call_content = function_call_content.function_call
parsed_args: dict[str, Any] = dict(function_call_content.parse_arguments() or {})
@@ -1545,6 +1555,13 @@ async def _auto_invoke_function(
session=invocation_session,
kwargs=runtime_kwargs.copy(),
)
# Always pass call_id to middleware for policy violation approval flow
middleware_context.metadata["call_id"] = function_call_content.call_id
# Pass policy approval flag to middleware via metadata (for re-invocation after approval)
if policy_approval_granted:
middleware_context.metadata["policy_approval_granted"] = True
async def final_function_handler(context_obj: Any) -> Any:
return await tool.invoke(
@@ -1557,12 +1574,21 @@ async def _auto_invoke_function(
# 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
result_content = Content.from_function_result(
call_id=function_call_content.call_id,
result=function_result,
)
return result_content
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:
@@ -1877,11 +1903,28 @@ 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,
)
# Build a map of call_id -> actual result for replacing placeholders
result_by_call_id: dict[str, Contents] = {}
for resp in fcc_todo.values():
if resp.approved:
# 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]
# 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
@@ -1905,17 +1948,24 @@ def _replace_approval_contents_with_results(
contents_to_remove.append(content_idx)
else:
# 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"
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
if result_idx < len(approved_function_results):
msg.contents[content_idx] = approved_function_results[result_idx]
result_idx += 1
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,10 +1974,30 @@ 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 = []
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]]:
@@ -2595,3 +2665,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
python/packages/core/tests/test_security.py
+2649
View File
File diff suppressed because it is too large Load Diff
python/packages/devui/agent_framework_devui/_executor.py
+9 -1
View File
@@ -744,6 +744,12 @@ class AgentFrameworkExecutor:
)
continue
# Extract policy_violation info if present (from security middleware)
policy_violation_data = content_dict.get("policy_violation")
additional_props: dict[str, Any] | None = None
if isinstance(policy_violation_data, dict):
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 +762,16 @@ class AgentFrameworkExecutor:
approved,
id=request_id,
function_call=function_call,
additional_properties=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"],
additional_props is not None,
)
except ImportError:
logger.warning(
python/packages/devui/agent_framework_devui/_mapper.py
+13 -1
View File
@@ -1744,7 +1744,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": {
@@ -1756,6 +1756,18 @@ class MessageMapper:
"output_index": context["output_index"],
"sequence_number": self._next_sequence(context),
}
# Include policy violation details if present (from security middleware)
additional_props = getattr(content, "additional_properties", None)
if additional_props and isinstance(additional_props, dict):
if 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."""