microsoft / microsoft/agent-framework
Python: [Feature]: Make the AG-UI Approval State store pluggable
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 13.6k
- Forks
- 2.3k
- Avg merge
- 2d 45m
- Merged PRs (30d)
- 358
Description
### 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 `OrderedDict`s: `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:
```python
# agent_framework_ag_ui/_agent.py:116
self._approval_state_store = InMemoryAGUIApprovalStateStore()
```
- The endpoint helper exposes `snapshot_store=` but **no** `approval_state_store=`:
```python
# 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:
```python
# 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:
```python
# 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:
```python
@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
- #6910 (closed, fixed by #6947) — introduced `InMemoryAGUIApprovalStateStore` and the restore/save round-trip this builds on.
- #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
```markdown
```
### Language/SDK
_No response_
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Assessment
This issue has not been assessed yet.