mirror of
https://github.com/microsoft/agent-framework.git
synced 2026-06-16 21:04:09 +08:00
Python: Validate approval responses against server-side pending request registry (#4548)
* Validate approval responses against server-side pending request registry * improvements * pin GHCP sdk version to non-breaking for now * Pin CHCP sdk to LKG. * really fix GHCP sdk pkg version * Fix HITL approval validation security gaps and memory leak - Validate rejected approval responses against pending_approvals registry, not just approved ones. Fabricated rejections without a prior request are now stripped from messages before reaching the LLM. - Bound _pending_approvals with OrderedDict + LRU eviction (max 10k) to prevent unbounded memory growth from abandoned approval requests. - Skip registration when function_call.name is None/empty; log warning when content.id or function_call is missing at registration time. - Document pending_approvals parameter in run_agent_stream docstring. - Add test for fabricated rejection attack scenario. - Assert pending approval entry is preserved after function name mismatch. - Pre-populate pending_approvals in rejection test for correct validation. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Apply pre-commit auto-fixes --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
This commit is contained in:
committed by
GitHub
Unverified
parent
2f2495e196
commit
18e433fc6d
@@ -2,6 +2,7 @@
|
||||
|
||||
"""AgentFrameworkAgent wrapper for AG-UI protocol."""
|
||||
|
||||
from collections import OrderedDict
|
||||
from collections.abc import AsyncGenerator
|
||||
from typing import Any, cast
|
||||
|
||||
@@ -101,6 +102,14 @@ class AgentFrameworkAgent:
|
||||
require_confirmation=require_confirmation,
|
||||
)
|
||||
|
||||
# Server-side registry of pending approval requests.
|
||||
# Keys are "{thread_id}:{request_id}", values are the function name.
|
||||
# Populated when approval requests are emitted; consumed when responses arrive.
|
||||
# Prevents bypass, function name spoofing, and replay attacks.
|
||||
# Bounded to prevent unbounded growth from abandoned approval requests.
|
||||
self._pending_approvals: OrderedDict[str, str] = OrderedDict()
|
||||
self._pending_approvals_max_size: int = 10_000
|
||||
|
||||
async def run(
|
||||
self,
|
||||
input_data: dict[str, Any],
|
||||
@@ -113,5 +122,7 @@ class AgentFrameworkAgent:
|
||||
Yields:
|
||||
AG-UI events
|
||||
"""
|
||||
async for event in run_agent_stream(input_data, self.agent, self.config):
|
||||
async for event in run_agent_stream(
|
||||
input_data, self.agent, self.config, pending_approvals=self._pending_approvals
|
||||
):
|
||||
yield event
|
||||
|
||||
@@ -369,11 +369,28 @@ def _handle_step_based_approval(messages: list[Any]) -> list[BaseEvent]:
|
||||
return events
|
||||
|
||||
|
||||
def _evict_oldest_approvals(registry: dict[str, str], max_size: int = 10_000) -> None:
|
||||
"""Evict the oldest entries from the pending-approvals registry (LRU).
|
||||
|
||||
Only effective when *registry* is an ``OrderedDict``; plain dicts are
|
||||
left untouched because insertion-order eviction is unreliable for them.
|
||||
"""
|
||||
if len(registry) <= max_size:
|
||||
return
|
||||
try:
|
||||
while len(registry) > max_size:
|
||||
registry.popitem(last=False) # type: ignore[call-arg]
|
||||
except (TypeError, KeyError):
|
||||
pass
|
||||
|
||||
|
||||
async def _resolve_approval_responses(
|
||||
messages: list[Any],
|
||||
tools: list[Any],
|
||||
agent: SupportsAgentRun,
|
||||
run_kwargs: dict[str, Any],
|
||||
pending_approvals: dict[str, str] | None = None,
|
||||
thread_id: str = "",
|
||||
) -> None:
|
||||
"""Execute approved function calls and replace approval content with results.
|
||||
|
||||
@@ -385,6 +402,11 @@ async def _resolve_approval_responses(
|
||||
tools: List of available tools
|
||||
agent: The agent instance (to get client and config)
|
||||
run_kwargs: Kwargs for tool execution
|
||||
pending_approvals: Server-side registry of pending approval requests.
|
||||
Keys are ``{thread_id}:{request_id}``, values are function names.
|
||||
When provided, every approval response is validated against this
|
||||
registry to prevent bypass, function name spoofing, and replay.
|
||||
thread_id: The conversation thread ID used to scope registry keys.
|
||||
"""
|
||||
fcc_todo = _collect_approval_responses(messages)
|
||||
if not fcc_todo:
|
||||
@@ -392,6 +414,59 @@ async def _resolve_approval_responses(
|
||||
|
||||
approved_responses = [resp for resp in fcc_todo.values() if resp.approved]
|
||||
rejected_responses = [resp for resp in fcc_todo.values() if not resp.approved]
|
||||
|
||||
# Validate every approval response (approved AND rejected) against the
|
||||
# pending approvals registry. Invalid responses are stripped from messages
|
||||
# entirely — not converted to rejection results, which would inject
|
||||
# attacker-controlled content into the LLM conversation.
|
||||
if pending_approvals is not None and (approved_responses or rejected_responses):
|
||||
validated: list[Any] = []
|
||||
validated_rejected: list[Any] = []
|
||||
invalid_ids: set[str] = set()
|
||||
for resp in approved_responses + rejected_responses:
|
||||
resp_id = resp.id or ""
|
||||
resp_name = resp.function_call.name if resp.function_call else None
|
||||
registry_key = f"{thread_id}:{resp_id}"
|
||||
|
||||
if registry_key not in pending_approvals:
|
||||
logger.warning(
|
||||
"Rejected approval response id=%s: no matching pending approval request",
|
||||
resp_id,
|
||||
)
|
||||
invalid_ids.add(resp_id)
|
||||
continue
|
||||
|
||||
pending_name = pending_approvals[registry_key]
|
||||
if resp_name != pending_name:
|
||||
logger.warning(
|
||||
"Rejected approval response id=%s: function name mismatch (response=%s, pending=%s)",
|
||||
resp_id,
|
||||
resp_name,
|
||||
pending_name,
|
||||
)
|
||||
invalid_ids.add(resp_id)
|
||||
continue
|
||||
|
||||
# Valid — consume entry to prevent replay
|
||||
del pending_approvals[registry_key]
|
||||
if resp.approved:
|
||||
validated.append(resp)
|
||||
else:
|
||||
validated_rejected.append(resp)
|
||||
|
||||
# Strip invalid approval responses from messages and fcc_todo so
|
||||
# _replace_approval_contents_with_results never sees them.
|
||||
if invalid_ids:
|
||||
for inv_id in invalid_ids:
|
||||
fcc_todo.pop(inv_id, None)
|
||||
for msg in messages:
|
||||
msg.contents = [
|
||||
c for c in msg.contents if not (c.type == "function_approval_response" and c.id in invalid_ids)
|
||||
]
|
||||
|
||||
approved_responses = validated
|
||||
rejected_responses = validated_rejected
|
||||
|
||||
approved_function_results: list[Any] = []
|
||||
|
||||
# Execute approved tool calls
|
||||
@@ -597,6 +672,7 @@ async def run_agent_stream(
|
||||
input_data: dict[str, Any],
|
||||
agent: SupportsAgentRun,
|
||||
config: AgentConfig,
|
||||
pending_approvals: dict[str, str] | None = None,
|
||||
) -> AsyncGenerator[BaseEvent]:
|
||||
"""Run agent and yield AG-UI events.
|
||||
|
||||
@@ -607,6 +683,10 @@ async def run_agent_stream(
|
||||
input_data: AG-UI request data with messages, state, tools, etc.
|
||||
agent: The Agent Framework agent to run
|
||||
config: Agent configuration
|
||||
pending_approvals: Optional server-side registry of pending approval
|
||||
requests. Keys are ``{thread_id}:{request_id}``, values are
|
||||
function names. When provided, approval responses are validated
|
||||
against this registry to prevent bypass, spoofing, and replay.
|
||||
|
||||
Yields:
|
||||
AG-UI events
|
||||
@@ -707,7 +787,7 @@ async def run_agent_stream(
|
||||
# Resolve approval responses (execute approved tools, replace approvals with results)
|
||||
# This must happen before running the agent so it sees the tool results
|
||||
tools_for_execution = tools if tools is not None else server_tools
|
||||
await _resolve_approval_responses(messages, tools_for_execution, agent, run_kwargs)
|
||||
await _resolve_approval_responses(messages, tools_for_execution, agent, run_kwargs, pending_approvals, thread_id)
|
||||
|
||||
# Defense-in-depth: replace approval payloads in snapshot with actual tool results
|
||||
# so CopilotKit does not re-send stale approval content on subsequent turns.
|
||||
@@ -782,6 +862,20 @@ async def run_agent_stream(
|
||||
for content in update.contents:
|
||||
content_type = getattr(content, "type", None)
|
||||
logger.debug(f"Processing content type={content_type}, message_id={flow.message_id}")
|
||||
|
||||
# Register pending approval requests so we can validate responses later
|
||||
if content_type == "function_approval_request" and pending_approvals is not None:
|
||||
if content.id and content.function_call and content.function_call.name:
|
||||
pending_approvals[f"{thread_id}:{content.id}"] = content.function_call.name
|
||||
# Evict oldest entries if the registry exceeds a safe bound (LRU)
|
||||
_evict_oldest_approvals(pending_approvals, max_size=10_000)
|
||||
else:
|
||||
logger.warning(
|
||||
"Approval request not registered: missing id=%s, function_call=%s, or function name",
|
||||
getattr(content, "id", None),
|
||||
getattr(content, "function_call", None),
|
||||
)
|
||||
|
||||
for event in _emit_content(
|
||||
content,
|
||||
flow,
|
||||
|
||||
Reference in New Issue
Block a user