Python: Fix AG-UI approval thread aliases

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:

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