Python: [Feature]: Make the AG-UI Approval State store pluggable

[antsok]

Description

Summary

AG-UI thread snapshots are persistable through a pluggable AGUIThreadSnapshotStore protocol (with an app-owned Redis/DB/file implementation passed via add_agent_framework_fastapi_endpoint(..., snapshot_store=...)). AG-UI Approval State has no such extension point: it is a single concrete, process-local InMemoryAGUIApprovalStateStore that is hard-wired into the agent and cannot be replaced. This makes durable, multi-replica human-in-the-loop (HITL) approval impossible on the AG-UI FastAPI host without forking or monkey-patching framework internals.

This request asks for the same pluggability the snapshot store already has: a public AGUIApprovalStateStore protocol plus a constructor/endpoint parameter to inject a custom implementation, with InMemoryAGUIApprovalStateStore remaining the default.

Current behavior (agent-framework-ag-ui 1.0.0rc8 / core 1.11.0, main)

The Approval State store is the only one of its kind and is not swappable:

• InMemoryAGUIApprovalStateStore is the single implementation — there is no protocol/ABC for it (_approval_state.py:32). It holds two bounded OrderedDicts: pending_approvals ((thread_id, interrupt_id) -> entry) and tool_approval_states (thread_id -> serialized ToolApprovalMiddleware state).
• It is hard-coded in AgentFrameworkAgent.__init__ with no parameter to override it:

  # agent_framework_ag_ui/_agent.py:116
  self._approval_state_store = InMemoryAGUIApprovalStateStore()

• The endpoint helper exposes snapshot_store= but no approval_state_store=:

  # agent_framework_ag_ui/_endpoint.py:81
  def add_agent_framework_fastapi_endpoint(..., snapshot_store: AGUIThreadSnapshotStore | None = None, ...):

• Every internal parameter is typed to the concrete class, not an interface — so even calling run_agent_stream directly gives no clean seam:

  # agent_framework_ag_ui/_agent_run.py:1654
  approval_state_store: InMemoryAGUIApprovalStateStore | None = None,
  # (also :580, :595, :614, :750, :786)

• InMemoryAGUIApprovalStateStore is not exported from agent_framework_ag_ui/__init__ (internal only), whereas AGUIThreadSnapshotStore + InMemoryAGUIThreadSnapshotStore are public.

For contrast, the snapshot store is a proper, exported, @runtime_checkable protocol:

# agent_framework_ag_ui/_snapshots.py:54
@runtime_checkable
class AGUIThreadSnapshotStore(Protocol): ...

Why it matters

InMemoryAGUIApprovalStateStore persists approval state between AG-UI requests only because it lives on the process-lifetime agent object — each request builds a fresh AgentSession (_agent_run.py:1821) and the middleware bookkeeping is copied in/out via _restore_tool_approval_state/_save_tool_approval_state (:578/:593). That is exactly the fix from #6947 (for #6910). It works for a single process, but:

• Multi-replica (e.g. Azure Container Apps / K8s with >1 replica): the request that raises an approval and the request that resumes it can land on different replicas with different in-memory stores → pending_approvals/tool_approval_states are missing on the resume replica → the approval is lost.
• Restart / scale-to-zero: in-flight approvals are dropped.

The docs acknowledge this is out of scope for the default store and push it onto the application ("not a distributed durability mechanism … choose deployment and storage architecture that matches your availability and worker topology requirements", Security Considerations for AG-UI), but there is no hook to actually supply that storage architecture for approvals — unlike snapshots.

A subtle asymmetry sharpens the case: the pending-approval interrupt already rides the durable snapshot (stored_snapshot.interrupt, _agent_run.py:1705), which is pluggable/durable. But the harness ToolApprovalMiddleware bookkeeping (tool_approval_states — queued sibling approvals, the auto-approved-tools set, hidden never-require siblings) lives only in the process-local approval store. So on a durable snapshot backend, cross-replica resume recovers the interrupt but not the approval bookkeeping — a partially-broken HITL that's hard to diagnose.

Proposed solution

Mirror the snapshot-store design:

1. Introduce a public, @runtime_checkable AGUIApprovalStateStore protocol in agent_framework_ag_ui and export it (alongside InMemoryAGUIApprovalStateStore). Because a durable backend can't expose raw OrderedDict attributes, the protocol should be method-based rather than attribute-based (today the code reaches into .tool_approval_states / .pending_approvals directly). A minimal shape covering what run_agent_stream currently does:

   @runtime_checkable
   class AGUIApprovalStateStore(Protocol):
       # tool-approval middleware state (per storage thread key)
       def get_tool_approval_state(self, thread_id: str) -> dict[str, Any] | None: ...
       def set_tool_approval_state(self, thread_id: str, state: dict[str, Any]) -> None: ...
       def clear_tool_approval_state(self, thread_id: str) -> None: ...
       # pending approval registry (keyed by (thread_id, interrupt_id) / alias keys)
       def register_pending_approval(self, key: tuple[str, str], entry: Any) -> None: ...
       def get_pending_approval(self, key: tuple[str, str]) -> Any | None: ...
       def pop_pending_approval(self, key: tuple[str, str]) -> None: ...
       def iter_pending_approvals(self, thread_id: str) -> Iterable[tuple[tuple[str, str], Any]]: ...

   (Async variants — async def — would be preferable so Redis/DB backends don't block the event loop; the snapshot store is already async.)
2. Accept it where the snapshot store is accepted:
   • AgentFrameworkAgent.__init__(..., approval_state_store: AGUIApprovalStateStore | None = None) — default InMemoryAGUIApprovalStateStore() (unchanged behavior when omitted).
   • add_agent_framework_fastapi_endpoint(..., approval_state_store: AGUIApprovalStateStore | None = None).
3. Widen the internal type hints in _agent_run.py from InMemoryAGUIApprovalStateStore to the protocol.

Backward compatibility

Fully backward compatible: omitting the parameter keeps today's InMemoryAGUIApprovalStateStore, byte-for-byte. This is purely an additive extension point.

Alternatives considered

• Subclass AgentFrameworkAgent and overwrite self._approval_state_store after super().__init__() — works at runtime (it's an instance attribute) but relies on a private field, a concrete type hint, and the store's raw OrderedDict shape; brittle across releases.
• Handle durability outside the framework (sticky sessions / single replica) — viable but constrains deployment topology and doesn't survive restarts; the framework already rejected this reasoning for snapshots by making them pluggable.

Related

• .NET: Python: [Bug]: AG-UI host loses tool calls when parallel calls require approval #6910 (closed, fixed by Python: Clear AG-UI queued approvals on cancel #6947) — introduced InMemoryAGUIApprovalStateStore and the restore/save round-trip this builds on.
• .NET: Python: [Bug]: AG-UI host builds a fresh AgentSession per request — session-stateful harness features silently lose their state between runs #6920 (open) — the broader "fresh AgentSession per request loses session-stateful harness state" root cause; a durable approval store is one piece of the multi-replica story.
• Snapshot store precedent: AGUIThreadSnapshotStore protocol + snapshot_store= endpoint parameter.

Environment

• agent-framework-ag-ui 1.0.0rc8, agent-framework-core 1.11.0 (git main, rev 68136ee), Python 3.12, AG-UI FastAPI host.

Code Sample

Language/SDK

No response

[rpelevin]

This is the right extension point. One change I would make before freezing the protocol: a distributed implementation cannot safely expose get and pop as separate operations. Approval resume is single-use authority, so two replicas can both read the same pending entry before either removes it and both proceed.

I would make consumption one atomic method, for example:

await claim_pending_approval(scope, thread_id, interrupt_id, decision_id, expected_version) -> claimed | already_consumed | not_found | version_mismatch

Aliases should be indexes to one canonical pending record, not independently consumable entries. The tool-approval middleware state also needs versioned compare-and-set semantics. A resume should bind to one joined generation of snapshot interrupt, pending approval, and tool_approval_state; if any part is missing, stale, expired, or from another scope, it should fail closed rather than reconstruct partial authority.

The regression anchors I would add are:

• two replicas race the same approval and exactly one claim succeeds;
• the durable snapshot survives but pending or middleware state is missing, so no tool executes;
• request-id and call-id aliases converge on the same consumed record;
• the process stops after claim but before a durable result, so the next replica reconciles rather than blindly re-executing;
• a stale writer cannot overwrite newer queued approval state.

The existing scope-resolver pattern looks reusable for authorization, but thread ID alone should never authorize store access.

Boundary: architecture and conformance feedback only. I have not reproduced the issue, run the AG-UI host, or validated an implementation.