Python: bridge AG-UI request state and session continuity

[moonbox3]

Motivation & Context

AG-UI hosted agents currently have two related session-state gaps. Request Shared State is not available through AgentSession.state, preventing context providers from consuming values such as user and tenant identifiers. State written by context providers is also lost when the next request creates a fresh session, so session-stateful features can contradict the replayed conversation.

This change connects both halves through the existing scoped thread snapshot boundary while keeping client-provided state separate from trusted server-owned provider state.

Description & Review Guide

• What are the major changes? Request Shared State is bridged into AgentSession.state for each run, while configured provider namespaces and reserved middleware state remain server-owned. Scoped AG-UI thread snapshots can now persist private Session Continuation State through the core typed session serialization contract. Corrupt or unsupported continuation data degrades safely without blocking the thread or suppressing RUN_FINISHED.
• What is the impact of these changes? Context providers can consume AG-UI request context without snapshot persistence, and eligible provider working state can survive later requests when scoped snapshots are configured. Private continuation state is neither accepted from nor emitted to AG-UI clients.
• What do you want reviewers to focus on? Please focus on the state-authority boundaries, snapshot lifecycle and isolation semantics, and the endpoint-level regressions for protected namespaces and failed continuation restoration or serialization.

Related Issue

Fixes #5197
Fixes #6920

Contribution Checklist

• The code builds clean without any errors or warnings
• All unit tests pass, and I have added new tests where possible
• The PR follows the Contribution Guidelines
• This PR is linked to an issue and there is no other open PR for this issue (see Related Issue above).
• This is not a breaking change. If it is a breaking change, add the breaking change label (or add "[BREAKING]" to the title prefix, before or after any language prefix) — a workflow keeps the label and title prefix in sync automatically.

[github-actions]

Python Test Coverage Report •

File	Stmts	Miss	Cover	Missing
packages/ag-ui/agent_framework_ag_ui
_agent_run.py	1108	124	88%	183–190, 237–238, 245, 354, 358, 360, 377, 404–405, 518, 532, 536, 540, 543, 548, 553, 562, 565, 572–578, 611, 623, 634, 637, 672, 726–730, 795, 810, 813, 815, 841, 867–869, 927, 929, 931, 934–938, 952, 960–965, 974–975, 1024–1027, 1038, 1046, 1078, 1093, 1107, 1119, 1149, 1153, 1156, 1158, 1198–1200, 1264, 1270–1271, 1276, 1280–1281, 1331–1333, 1345–1347, 1445, 1453, 1470, 1474, 1519, 1578, 1608–1609, 1731, 1877, 1937, 1954, 1974–1975, 1982, 2090, 2118, 2126, 2128, 2131, 2137, 2192, 2195, 2205–2206, 2213, 2259
_snapshots.py	83	5	93%	138, 221, 223, 232–233
TOTAL	44195	5266	88%

Python Unit Test Overview

Tests	Skipped	Failures	Errors	Time
8873	33 💤	0 ❌	0 🔥	2m 23s ⏱️

[Copilot]

Pull request overview

This PR closes two AG-UI hosted-agent state gaps by (1) exposing per-request Shared State through AgentSession.state so context providers can consume request context, and (2) persisting eligible server-owned provider working state across scoped runs via a private AGUIThreadSnapshot.session_state continuation payload (never replayed to clients).

Changes:

• Add optional private session_state to AGUIThreadSnapshot and persist/restore it through AgentSession’s typed serialization contract.
• Bridge request Shared State into AgentSession.state while protecting provider namespaces and reserved middleware buckets from client ownership.
• Expand tests and docs to cover authority boundaries, snapshot lifecycle, failure-safe degradation, and continuity semantics.

Reviewed changes

Copilot reviewed 6 out of 6 changed files in this pull request and generated 2 comments.

Show a summary per file

File	Description
python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py	Implements session-state bridging, protected namespaces, and snapshot persistence/restoration of private continuation state.
python/packages/ag-ui/agent_framework_ag_ui/_snapshots.py	Extends the snapshot model with session_state (private continuation) while keeping the store protocol intact.
python/packages/ag-ui/tests/ag_ui/test_endpoint.py	Adds endpoint-level tests covering request-state bridging, protected namespaces, continuity, and failure paths.
python/packages/ag-ui/tests/ag_ui/test_snapshots.py	Updates snapshot model/store tests for the new session_state field and defensive copying behavior.
python/packages/ag-ui/README.md	Documents state authority boundaries and continuity behavior for scoped snapshots and hosted endpoints.

[github-actions]

Automated Code Review

Reviewers: 5 | Confidence: 90%

✓ Correctness

The PR correctly implements AG-UI request state bridging and session continuation. State authority boundaries are cleanly enforced: continuation is restored first from trusted storage, request state overlays only non-protected keys, tool approval uses its own store, and serialization correctly excludes shared-state keys, tool approval state, and history provider state. Error handling around restoration and serialization is graceful with appropriate logging. The get_final_response() addition ensures provider after_run hooks execute for interrupted flows. The InMemoryAGUIThreadSnapshotStore properly deep-copies session_state. No correctness issues found.

✓ Security Reliability

This PR implements a well-designed state-authority boundary between untrusted AG-UI request state and trusted server-owned Session Continuation State. Private continuation is correctly excluded from client hydration (line 1616-1620 of _agent_run.py only emits messages/state/interrupt, never session_state). Protected namespaces prevent client injection of provider state, approval state, and middleware state. Failed restoration and serialization degrade gracefully without blocking threads. The InMemoryAGUIThreadSnapshotStore correctly deep-copies session_state on save and get. No high-severity security or reliability issues were found.

✓ Test Coverage

The test coverage for this PR is excellent. Eleven new integration tests cover the full spectrum of the new session state bridging and continuation features: happy paths (bridging, restoration, persistence), error paths (corrupt state, serialization failures, save failures, run failures), security boundaries (protected namespaces, approval isolation), isolation (scope/thread boundaries), and edge cases (collisions, history exclusion, disabled snapshots). Snapshot model tests are appropriately updated. The only untested branch is the defensive make_json_safe non-dict warning path (line 1731), which is practically unreachable since make_json_safe on a dict returns a dict. No blocking issues found.

✓ Failure Modes

The implementation demonstrates comprehensive failure-mode handling. Corrupt continuation state degrades gracefully via exception handling in _restore_session_continuation_state. Serialization failures in _safe_serialize_session_continuation_state are caught and return None, preserving RUN_FINISHED emission. The snapshot store save already had exception handling that logs and continues. Protected namespace enforcement prevents client injection into server-owned state. The InMemoryAGUIThreadSnapshotStore uses deepcopy for defensive isolation. Private session_state is excluded from the hydration path, preventing leakage to clients. No concrete silent-failure, lost-error, or partial-write paths were identified that could lead to data loss or thread bricking.

✗ Design Approach

I found one design issue in the new request-state bridge. The implementation reuses the same nested objects when copying request Shared State into AgentSession.state, so a context provider that mutates session.state can silently turn untrusted per-run request context into persisted replayable AG-UI state. That contradicts the new README contract that request values are only per-run context and should not cross the continuation boundary except as the client supplied them.

Flagged Issues

• python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py — request state is shallow-copied into both flow.current_state and session.state, so provider mutations to nested request objects leak back into replayable snapshot state even though the README documents request values as untrusted per-run context.

Suggestions

• The warning branch at _agent_run.py:1731-1734 (where make_json_safe returns a non-dict) has no dedicated test and is effectively unreachable defensive code since make_json_safe on a dict always returns a dict. Consider either removing the branch (relying on the outer except) or adding a brief unit test that mocks make_json_safe to return a non-dict, verifying the warning is logged and None is returned.

---

Automated review by moonbox3's agents

[github-actions]

Flagged issue

python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py — request state is shallow-copied into both flow.current_state and session.state, so provider mutations to nested request objects leak back into replayable snapshot state even though the README documents request values as untrusted per-run context.

---

Source: automated DevFlow PR review

[moonbox3]

Addressed the DevFlow flagged aliasing issue in 00d44f5. Bridged Shared State values are now deep-copied before exposure through AgentSession.state, and an endpoint-level regression verifies that provider mutations cannot alter replayable Shared State.