Python: bridge AG-UI request state and session continuity

python/packages/ag-ui/AGENTS.md

@@ -15,7 +15,8 @@ AG-UI protocol integration for building agent UIs with the AG-UI standard.
 ## Types
 
 - **`AGUIRequest`** / **`AGUIChatOptions`** - Request types
-- **`AGUIThreadSnapshot`** / **`AGUIThreadSnapshotStore`** - Replayable thread snapshot model and scoped async store protocol
+- **`AGUIThreadSnapshot`** / **`AGUIThreadSnapshotStore`** - Thread snapshot model with client-replayable data,
+  private Session Continuation State, and a scoped async store protocol
 - **`availableInterrupts` / `resume`** - Optional canonical AG-UI `Interrupt` and `ResumeEntry` protocol data
 - **`AgentState`** / **`RunMetadata`** - State management types
 - **`PredictStateConfig`** - Configuration for state prediction

python/packages/ag-ui/README.md

@@ -337,6 +337,28 @@ Endpoint configuration requires `snapshot_scope_resolver` whenever a snapshot st
 the store is already set on a pre-wrapped `AgentFrameworkAgent` or `AgentFrameworkWorkflow`. The resolver returns
 the application-defined Snapshot Scope used with the AG-UI Thread id as the storage key.
 
+For hosted agents, request Shared State is also available through `AgentSession.state` during that run, whether or
+not snapshot persistence is configured. Request values are untrusted per-run context: they overlay ordinary restored
+values, are not passed through typed session restoration, and are excluded from private Session Continuation State.
+Keys owned by configured context providers or reserved for approval and message-injection middleware are not copied
+into `AgentSession.state`; their server-owned values take precedence over client Shared State.
+
+When scoped snapshots are configured, each category has one State Authority:
+
+| State category | State Authority |
+| --- | --- |
+| Conversation history | AG-UI Thread Snapshot messages |
+| AG-UI Shared State and request context | The current AG-UI request and replayable snapshot state |
+| Approval State | The Approval State Store |
+| Other server-produced provider working state | Private Session Continuation State |
+
+Session Continuation State is stored atomically in the optional `AGUIThreadSnapshot.session_state` field and restored
+through the core `AgentSession` typed serialization contract. It is never accepted from an AG-UI request or emitted
+during hydration. Deleting a scoped thread snapshot resets its replayable and private state together, and clearing a
+Snapshot Scope removes all such records in that scope. Missing or empty request Shared State is not a reset command.
+If private continuation cannot be restored or serialized, the endpoint logs the failure and continues without that
+continuation so stale or unsupported provider state cannot permanently block the thread or suppress `RUN_FINISHED`.
+
 AG-UI Thread ids identify AG-UI Threads; they do not authorize snapshot access. Do not treat a thread id as a bearer
 credential or tenant boundary. Production applications must authenticate and authorize every AG-UI endpoint request
 and choose a Snapshot Scope that represents the app's real access boundary, such as an authenticated user, tenant,
@@ -348,14 +370,22 @@ It is not an authentication, tenant authorization, or distributed durability mec
 responsible for endpoint authentication, tenant authorization, and deployment/storage architecture that matches their
 availability and worker topology requirements.
 
-Stored snapshots are untrusted application data with confidentiality impact. They may contain sensitive user text,
-model output, tool results, function arguments, UI payloads, Shared State, and interrupt data. The built-in
-`InMemoryAGUIThreadSnapshotStore` is in-memory only, process-local, bounded, latest-only, and not durable production
-storage. It is cleared on process restart and is not shared across workers.
+Snapshot storage is treated as trusted server-side storage because private continuation is eligible for typed core
+restoration; applications are responsible for providing its integrity protection. Snapshots also have confidentiality
+impact: they may contain sensitive user text, model output, tool results, function arguments, UI payloads, Shared State,
+interrupt data, and private provider working state. The built-in `InMemoryAGUIThreadSnapshotStore` is in-memory only,
+process-local, bounded, latest-only, and not durable production storage. It is cleared on process restart and is not
+shared across workers.
 
 No file-backed AG-UI snapshot store is provided by the package. Applications that need durable persistence should
 provide an app-owned implementation of the `AGUIThreadSnapshotStore` protocol and own storage hardening, including
-encryption, access control, retention, audit, data residency, and deletion behavior.
+encryption, integrity protection, access control, retention, audit, data residency, and deletion behavior. Existing
+custom stores remain source-compatible because `session_state` is optional, but they provide Session State Continuity
+only when they round-trip that field unchanged with the rest of the snapshot.
+
+The supported consistency model is one active run per `(Snapshot Scope, threadId)`. Concurrent writes to the same
+scoped thread remain last-writer-wins. Applications that require stronger consistency must serialize those runs using
+coordination appropriate to their deployment; a process-local lock does not provide distributed consistency.
 
 ## Architecture
 

python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py

@@ -30,6 +30,9 @@
 from agent_framework import (
     AgentSession,
     Content,
+    HistoryProvider,
+    InMemoryHistoryProvider,
+    MESSAGE_INJECTION_PENDING_MESSAGES_STATE_KEY,
     Message,
     SupportsAgentRun,
 )
@@ -581,6 +584,7 @@ def _restore_tool_approval_state(
     thread_id: str,
 ) -> None:
     """Restore only core tool-approval state into the per-run AgentSession."""
+    session.state.pop(_TOOL_APPROVAL_STATE_KEY, None)
     if approval_state_store is None:
         return
     stored_state = approval_state_store.tool_approval_states.get(thread_id)
@@ -1624,16 +1628,22 @@ async def _save_thread_snapshot(
     messages: list[dict[str, Any]],
     state: dict[str, Any] | None,
     interrupt: list[dict[str, Any]] | None,
+    session_state: dict[str, Any] | None,
 ) -> None:
-    """Save the latest replayable AG-UI Thread Snapshot when persistence is configured."""
+    """Save the latest AG-UI Thread Snapshot when persistence is configured."""
     if config.snapshot_store is None or scope is None:
         return
 
     try:
         await config.snapshot_store.save(
             scope=scope,
             thread_id=thread_id,
-            snapshot=AGUIThreadSnapshot(messages=messages, state=state, interrupt=interrupt),
+            snapshot=AGUIThreadSnapshot(
+                messages=messages,
+                state=state,
+                interrupt=interrupt,
+                session_state=session_state,
+            ),
         )
     except Exception:
         # The run itself already streamed successfully; a transient store failure
@@ -1646,6 +1656,90 @@ async def _save_thread_snapshot(
         )
 
 
+def _restore_session_continuation_state(session: AgentSession, snapshot: AGUIThreadSnapshot | None) -> None:
+    """Restore typed private state from trusted snapshot storage."""
+    if snapshot is None or snapshot.session_state is None:
+        return
+    try:
+        restored = AgentSession.from_dict(
+            {
+                "type": "session",
+                "session_id": session.session_id,
+                "state": snapshot.session_state,
+            }
+        )
+    except Exception:
+        logger.exception(
+            "Failed to restore AG-UI Session Continuation State for session_id=%s; continuing without it.",
+            session.session_id,
+        )
+        return
+    session.state.update(restored.state)
+
+
+def _request_state_protected_keys(agent: SupportsAgentRun) -> set[str]:
+    """Return session-state namespaces that client Shared State cannot own."""
+    context_providers = cast(list[Any], getattr(agent, "context_providers", []))
+    return {
+        _TOOL_APPROVAL_STATE_KEY,
+        InMemoryHistoryProvider.DEFAULT_SOURCE_ID,
+        MESSAGE_INJECTION_PENDING_MESSAGES_STATE_KEY,
+        *(provider.source_id for provider in context_providers),
+    }
+
+
+def _serialize_session_continuation_state(
+    session: AgentSession,
+    agent: SupportsAgentRun,
+    *,
+    shared_state_keys: set[str],
+) -> dict[str, Any] | None:
+    """Serialize server-owned state while preserving each AG-UI State Authority."""
+    context_providers = cast(list[Any], getattr(agent, "context_providers", []))
+    excluded_keys = {
+        *shared_state_keys,
+        _TOOL_APPROVAL_STATE_KEY,
+        *(provider.source_id for provider in context_providers if isinstance(provider, HistoryProvider)),
+    }
+    continuation_state = {key: value for key, value in session.state.items() if key not in excluded_keys}
+    if not continuation_state:
+        return None
+
+    serialized_session = AgentSession(session_id=session.session_id)
+    serialized_session.state.update(continuation_state)
+    return cast(dict[str, Any], serialized_session.to_dict()["state"])
+
+
+def _safe_serialize_session_continuation_state(
+    session: AgentSession,
+    agent: SupportsAgentRun,
+    *,
+    shared_state_keys: set[str],
+) -> dict[str, Any] | None:
+    """Return JSON-safe continuation state without failing a completed run."""
+    try:
+        serialized_state = _serialize_session_continuation_state(
+            session,
+            agent,
+            shared_state_keys=shared_state_keys,
+        )
+        if serialized_state is None:
+            return None
+        safe_state = make_json_safe(serialized_state)
+        if isinstance(safe_state, dict):
+            return cast(dict[str, Any], safe_state)
+        logger.warning(
+            "Ignoring AG-UI Session Continuation State with unsupported serialized type: %s",
+            type(safe_state).__name__,
+        )
+    except Exception:
+        logger.exception(
+            "Failed to serialize AG-UI Session Continuation State for session_id=%s; saving snapshot without it.",
+            session.session_id,
+        )
+    return None
+
+
 async def run_agent_stream(
     input_data: dict[str, Any],
     agent: SupportsAgentRun,
@@ -1821,6 +1915,15 @@ async def run_agent_stream(
         session = AgentSession(session_id=thread_id, service_session_id=supplied_thread_id)
     else:
         session = AgentSession(session_id=thread_id)
+    _restore_session_continuation_state(session, stored_snapshot)
+    protected_session_state_keys = _request_state_protected_keys(agent)
+    session.state.update(
+        {
+            key: copy.deepcopy(value)
+            for key, value in flow.current_state.items()
+            if key not in protected_session_state_keys
+        }
+    )
     _restore_tool_approval_state(session, approval_state_store, approval_thread_id)
 
     # Inject metadata for AG-UI orchestration (Feature #2: Azure-safe truncation)
@@ -1894,6 +1997,11 @@ async def run_agent_stream(
             messages=persisted_messages,
             state=cast(dict[str, Any], make_json_safe(flow.current_state)) if flow.current_state else None,
             interrupt=None,
+            session_state=_safe_serialize_session_continuation_state(
+                session,
+                agent,
+                shared_state_keys=set(flow.current_state).difference(protected_session_state_keys),
+            ),
         )
         _save_tool_approval_state(session, approval_state_store, approval_thread_id)
         yield _build_run_finished_event(run_id=run_id, thread_id=thread_id)
@@ -2000,6 +2108,9 @@ async def run_agent_stream(
         if flow.waiting_for_approval:
             break
 
+    if flow.waiting_for_approval and isinstance(stream, ResponseStream):
+        await stream.get_final_response()
+
     # If no updates at all, still emit RunStarted
     if not run_started_emitted:
         yield RunStartedEvent(run_id=run_id, thread_id=thread_id)
@@ -2188,6 +2299,11 @@ async def run_agent_stream(
         messages=persisted_messages,
         state=latest_state_snapshot,
         interrupt=flow.interrupts or None,
+        session_state=_safe_serialize_session_continuation_state(
+            session,
+            agent,
+            shared_state_keys=set(flow.current_state).difference(protected_session_state_keys),
+        ),
     )
     _save_tool_approval_state(session, approval_state_store, approval_thread_id)
     yield _build_run_finished_event(run_id=run_id, thread_id=thread_id, interrupts=flow.interrupts)

python/packages/ag-ui/agent_framework_ag_ui/_snapshots.py

@@ -32,22 +32,25 @@
 
 @dataclass(slots=True)
 class AGUIThreadSnapshot:
-    """Replayable AG-UI Thread state.
+    """AG-UI Thread state.
 
-    AG-UI Thread Snapshots intentionally contain only data that can be replayed
-    to a UI: message snapshots, optional Shared State, and optional interruption
-    state. They do not include raw events, request metadata, auth claims,
+    AG-UI Thread Snapshots contain client-replayable message, Shared State, and
+    interruption data plus optional private Session Continuation State. Private
+    continuation is trusted server data and must never be replayed to the client.
+    Snapshots do not include raw events, request metadata, auth claims,
     diagnostics, traces, or provider responses.
 
     Attributes:
         messages: Replayable AG-UI message snapshots.
         state: Optional AG-UI Shared State snapshot.
         interrupt: Optional interruption state from ``RUN_FINISHED.outcome.interrupts``.
+        session_state: Optional private serialized ``AgentSession.state`` payload.
     """
 
     messages: list[dict[str, Any]] = field(default_factory=list)
     state: dict[str, Any] | None = None
     interrupt: list[dict[str, Any]] | None = None
+    session_state: dict[str, Any] | None = None
 
 
 @runtime_checkable