diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py b/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py index b029d3cb18..28419f9324 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py @@ -8,6 +8,7 @@ import json import logging import uuid +from collections import OrderedDict from collections.abc import AsyncIterable, Awaitable from typing import TYPE_CHECKING, Any, TypedDict, cast @@ -555,12 +556,55 @@ def _pending_approval_alias_keys( return {_pending_approval_key(thread_id, alias) for alias in aliases} +def _remove_pending_approval(registry: dict[PendingApprovalKey, PendingApprovalEntry], key: PendingApprovalKey) -> None: + """Remove one pending approval and every alias that references the same entry.""" + entry = registry.pop(key, None) + if entry is None or isinstance(entry, str): + return + + for alias_key, alias_entry in list(registry.items()): + if alias_entry is entry: + registry.pop(alias_key, None) + + +def _register_pending_approval( + registry: dict[PendingApprovalKey, PendingApprovalEntry], + thread_ids: list[str], + name: str, + arguments: str | None, + *, + request_id: str, + interrupt_id: str | None, +) -> None: + """Register one pending approval under each distinct thread identity.""" + keys = list( + dict.fromkeys( + _pending_approval_key(thread_id, approval_id) + for thread_id in thread_ids + for approval_id in (request_id, interrupt_id) + if approval_id + ) + ) + for key in keys: + _remove_pending_approval(registry, key) + + entry = _make_pending_approval_entry(name, arguments, request_id=request_id, interrupt_id=interrupt_id) + for key in keys: + registry[key] = entry + + def _consume_pending_approval_entry( pending_approvals: dict[PendingApprovalKey, PendingApprovalEntry], thread_id: str, entry: PendingApprovalEntry, *ids: str | None, ) -> None: + if not isinstance(entry, str): + for key, candidate in list(pending_approvals.items()): + if candidate is entry: + pending_approvals.pop(key, None) + return + for alias_key in _pending_approval_alias_keys(thread_id, entry, *ids): pending_approvals.pop(alias_key, None) @@ -795,13 +839,11 @@ def _evict_oldest_approvals(registry: dict[PendingApprovalKey, PendingApprovalEn 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: + if len(registry) <= max_size or not isinstance(registry, OrderedDict): return - try: - while len(registry) > max_size: - registry.popitem(last=False) # type: ignore[call-arg] - except (TypeError, KeyError): - pass + while len(registry) > max_size: + oldest_key = next(iter(registry)) + _remove_pending_approval(registry, oldest_key) async def _resolve_approval_responses( @@ -1213,6 +1255,7 @@ async def run_agent_stream( """ # Parse IDs thread_id = input_data.get("thread_id") or input_data.get("threadId") or str(uuid.uuid4()) + client_thread_id = thread_id run_id = input_data.get("run_id") or input_data.get("runId") or str(uuid.uuid4()) snapshot_scope = cast(str | None, input_data.get(_SNAPSHOT_SCOPE_INPUT_KEY)) @@ -1484,15 +1527,14 @@ async def run_agent_stream( if content_type == "function_approval_request" and pending_approvals is not None: if content.id and content.function_call and content.function_call.name: canonical_interrupt_id = content.function_call.call_id or content.id - pending_entry = _make_pending_approval_entry( + _register_pending_approval( + pending_approvals, + [str(client_thread_id), str(thread_id)], content.function_call.name, canonical_function_arguments(content.function_call), request_id=str(content.id), interrupt_id=str(canonical_interrupt_id), ) - pending_approvals[_pending_approval_key(thread_id, str(content.id))] = pending_entry - if canonical_interrupt_id: - pending_approvals[_pending_approval_key(thread_id, str(canonical_interrupt_id))] = pending_entry # Evict oldest entries if the registry exceeds a safe bound (LRU) _evict_oldest_approvals(pending_approvals, max_size=10_000) else: diff --git a/python/packages/ag-ui/tests/ag_ui/test_agent_wrapper_comprehensive.py b/python/packages/ag-ui/tests/ag_ui/test_agent_wrapper_comprehensive.py index 8e6bf7e3c6..533026bca0 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_agent_wrapper_comprehensive.py +++ b/python/packages/ag-ui/tests/ag_ui/test_agent_wrapper_comprehensive.py @@ -1087,6 +1087,95 @@ async def stream_fn_post_approval( assert call_count == 0, "Replay of consumed approval should not execute the tool" +@pytest.mark.parametrize( + "resume_thread_id", + [ + pytest.param("client-thread", id="client-thread"), + pytest.param("provider-conversation", id="provider-conversation"), + ], +) +async def test_approval_resolves_with_client_or_provider_thread_id( + streaming_chat_client_stub: Any, + resume_thread_id: str, +) -> None: + """A stateful provider approval remains resolvable by either advertised thread identity.""" + from agent_framework import tool + from agent_framework.ag_ui import AgentFrameworkAgent + + execution_count = 0 + + @tool( + name="sensitive_action", + description="A sensitive action requiring approval", + approval_mode="always_require", + ) + def sensitive_action() -> str: + nonlocal execution_count + execution_count += 1 + return "executed" + + async def approval_stream( + messages: MutableSequence[Message], options: ChatOptions, **kwargs: Any + ) -> AsyncIterator[ChatResponseUpdate]: + yield ChatResponseUpdate( + contents=[ + Content.from_function_call( + name="sensitive_action", + call_id="call_sensitive", + arguments="{}", + ) + ], + conversation_id="provider-conversation", + ) + + wrapper = AgentFrameworkAgent( + agent=Agent( + client=streaming_chat_client_stub(approval_stream), + name="test_agent", + instructions="Test", + tools=[sensitive_action], + ) + ) + + async for _ in wrapper.run({"thread_id": "client-thread", "messages": [{"role": "user", "content": "do it"}]}): + pass + + assert ("client-thread", "call_sensitive") in wrapper._pending_approvals + assert ("provider-conversation", "call_sensitive") in wrapper._pending_approvals + + async def completion_stream( + messages: MutableSequence[Message], options: ChatOptions, **kwargs: Any + ) -> AsyncIterator[ChatResponseUpdate]: + yield ChatResponseUpdate(contents=[Content.from_text(text="Done")]) + + wrapper.agent = Agent( + client=streaming_chat_client_stub(completion_stream), + name="test_agent", + instructions="Test", + tools=[sensitive_action], + ) + + def approval_input(thread_id: str) -> dict[str, Any]: + return { + "thread_id": thread_id, + "messages": [], + "resume": [{"interruptId": "call_sensitive", "status": "resolved", "payload": {"accepted": True}}], + } + + async for _ in wrapper.run(approval_input(resume_thread_id)): + pass + + assert execution_count == 1 + assert ("client-thread", "call_sensitive") not in wrapper._pending_approvals + assert ("provider-conversation", "call_sensitive") not in wrapper._pending_approvals + + replay_thread_id = "provider-conversation" if resume_thread_id == "client-thread" else "client-thread" + async for _ in wrapper.run(approval_input(replay_thread_id)): + pass + + assert execution_count == 1 + + async def test_approval_function_name_mismatch_is_blocked(streaming_chat_client_stub): """Test that an approval response with a mismatched function name is rejected.""" from agent_framework import tool