microsoft / microsoft/agent-framework

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

Open
#7,082 2 comments 0 reactions 1 assignee View on GitHub

Nobody has claimed this yet.

ag-ui harness python
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

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.