Skip to content

Python: Add opt-in AG-UI thread snapshot persistence and hydration#6471

Merged
moonbox3 merged 14 commits into
microsoft:mainfrom
moonbox3:ag-ui-state-persistence
Jun 12, 2026
Merged

Python: Add opt-in AG-UI thread snapshot persistence and hydration#6471
moonbox3 merged 14 commits into
microsoft:mainfrom
moonbox3:ag-ui-state-persistence

Conversation

@moonbox3

Copy link
Copy Markdown
Contributor

Motivation and Context

AG-UI endpoints currently keep no server-side record of a thread. If the user refreshes the page or reconnects, the UI loses the conversation, any shared state, and any pending interrupt, and the backend has no way to replay them. Clients that resend full history also become the source of truth, which allows tampered transcripts to reach the agent.

This change adds opt-in AG-UI Thread Snapshot persistence so a thread can be rehydrated after refresh or reconnect, for both agent and workflow endpoints, while keeping the server-side history authoritative.

Description

Persistence is disabled by default and activates only when both pieces are configured on the endpoint:

  • snapshot_store: an implementation of the new AGUIThreadSnapshotStore protocol
  • snapshot_scope_resolver: resolves an application-defined Snapshot Scope per request; required because an AG-UI thread id identifies a thread but does not authorize access to it. The scope is part of the storage key and must represent the app's authorization boundary (e.g. tenant or user)

What gets stored is intentionally limited to replayable UI data per (scope, thread_id), latest snapshot only: message snapshots, optional shared state, and optional RUN_FINISHED.interrupt metadata. No raw events, request metadata, auth claims, or provider responses.

Behavior added:

  • Hydrate Requests (empty messages, no resume payload) replay the stored state snapshot, messages snapshot, and any pending interrupt without invoking the wrapped agent or workflow
  • Normal agent turns load the stored snapshot and reconstruct history server-side: stored prior history is authoritative, the incoming new user turn is appended, full-history clients are deduplicated, and tampered history is ignored
  • Workflow runs capture replayable output through a snapshot builder that folds emitted state/messages snapshots and synthesizes messages from text and tool events
  • Interrupted runs persist their interrupt metadata so pending approvals survive a refresh; runs that end in RUN_ERROR do not overwrite the last good snapshot
  • InMemoryAGUIThreadSnapshotStore ships as a bounded, process-local store for development, demos, and tests; durable stores are app-owned implementations of the protocol

The README documents the security model: snapshot persistence is opt-in, the scope resolver is mandatory when a store is configured, and snapshot confidentiality follows the store the app provides.

New public exports: AGUIThreadSnapshot, AGUIThreadSnapshotStore, InMemoryAGUIThreadSnapshotStore, SnapshotScope, SnapshotScopeResolver, AGUIThreadID, DEFAULT_MAX_THREAD_SNAPSHOTS.

Contribution Checklist

  • The code builds clean without any errors or warnings
  • The PR follows the Contribution Guidelines
  • All unit tests pass, and I have added new tests where possible
  • Is this a breaking change? If yes, add "[BREAKING]" prefix to the title of the PR.

moonbox3 added 8 commits June 10, 2026 17:01
Key decisions:\n- Introduce an AGUIThreadSnapshot model limited to replayable messages, optional Shared State, and optional interrupt state.\n- Define AGUIThreadSnapshotStore as an async protocol keyed by explicit Snapshot Scope and AG-UI Thread id.\n- Add InMemoryAGUIThreadSnapshotStore as memory-only, latest-only, bounded local/demo/test storage; no file-backed store is introduced.\n- Require snapshot_scope_resolver whenever an endpoint is configured with a snapshot store, including pre-wrapped runners, so thread ids are not authorization boundaries.\n\nFiles changed:\n- packages/ag-ui/agent_framework_ag_ui/_snapshots.py\n- packages/ag-ui/agent_framework_ag_ui/__init__.py\n- packages/ag-ui/agent_framework_ag_ui/_agent.py\n- packages/ag-ui/agent_framework_ag_ui/_workflow.py\n- packages/ag-ui/agent_framework_ag_ui/_endpoint.py\n- packages/core/agent_framework/ag_ui/__init__.py\n- packages/core/agent_framework/ag_ui/__init__.pyi\n- packages/ag-ui/tests/ag_ui/test_snapshots.py\n- packages/ag-ui/tests/ag_ui/test_endpoint.py\n- packages/ag-ui/tests/ag_ui/test_public_exports.py\n- packages/ag-ui/AGENTS.md\n\nVerification:\n- uv run pytest packages/ag-ui/tests/ag_ui/test_snapshots.py packages/ag-ui/tests/ag_ui/test_public_exports.py packages/ag-ui/tests/ag_ui/test_endpoint.py::test_endpoint_requires_snapshot_scope_resolver_when_store_configured packages/ag-ui/tests/ag_ui/test_endpoint.py::test_endpoint_accepts_snapshot_store_with_scope_resolver -q\n- uv run pytest packages/ag-ui/tests/ag_ui/test_endpoint.py::test_endpoint_requires_snapshot_scope_resolver_when_store_configured packages/ag-ui/tests/ag_ui/test_endpoint.py::test_endpoint_requires_snapshot_scope_resolver_when_wrapped_runner_has_store packages/ag-ui/tests/ag_ui/test_endpoint.py::test_endpoint_accepts_snapshot_store_with_scope_resolver -q\n- uv run poe syntax -P ag-ui -C\n- uv run poe pyright -P ag-ui\n- uv run poe syntax -P core -C\n- uv run poe pyright -P core\n- uv run poe typing -P ag-ui\n- uv run poe typing -P core\n- uv run poe test -P ag-ui\n- uv run poe check -P ag-ui\n- git diff --check\n- git diff --cached --check\n\nBlockers / next iteration:\n- No blockers. Next slice can use the store contract to capture and hydrate agent snapshots.\n- uv repeatedly refreshed azure-ai-projects in uv.lock during local runs; reverted the generated lockfile churn because this change does not alter dependencies.\n- The poe-check commit hook was skipped after manual verification because it reformatted unrelated core MCP files outside this task.
Key decisions:
- Resolve Snapshot Scope per endpoint request and pass it to the AG-UI runner only when snapshot storage is active.
- Treat empty messages with no resume payload as an agent Hydrate Request when a scoped snapshot store is configured, replaying stored Shared State and message snapshots without invoking the wrapped agent.
- Save the latest replayable agent message snapshot and Shared State at normal completion under Snapshot Scope plus AG-UI Thread id; no durable or file-backed store is introduced.

Files changed:
- packages/ag-ui/agent_framework_ag_ui/_agent_run.py
- packages/ag-ui/agent_framework_ag_ui/_endpoint.py
- packages/ag-ui/agent_framework_ag_ui/_snapshots.py
- packages/ag-ui/tests/ag_ui/test_endpoint.py

Verification:
- uv run pytest packages/ag-ui/tests/ag_ui/test_endpoint.py::test_agent_endpoint_hydrates_stored_thread_snapshot_without_invoking_agent -q
- uv run pytest packages/ag-ui/tests/ag_ui/test_endpoint.py::test_agent_endpoint_hydrates_stored_thread_snapshot_without_invoking_agent packages/ag-ui/tests/ag_ui/test_endpoint.py::test_agent_endpoint_hydrates_snapshots_by_scope_and_thread -q
- uv run pytest packages/ag-ui/tests/ag_ui/test_endpoint.py::test_endpoint_empty_messages packages/ag-ui/tests/ag_ui/test_endpoint.py::test_agent_endpoint_hydrates_stored_thread_snapshot_without_invoking_agent packages/ag-ui/tests/ag_ui/test_endpoint.py::test_agent_endpoint_hydrates_snapshots_by_scope_and_thread -q
- uv run poe syntax -P ag-ui -C
- uv run poe pyright -P ag-ui
- uv run poe typing -P ag-ui
- uv run poe test -P ag-ui
- uv run poe check -P ag-ui
- git diff --check
- git diff --cached --check

Blockers / next iteration:
- No blockers. Next slice can reconstruct normal new-user agent turns from stored snapshots.
- uv repeatedly refreshed azure-ai-projects in uv.lock during local runs; reverted the generated lockfile churn because this change does not alter dependencies.
- The poe-check commit hook was skipped after manual verification because it refreshed unrelated uv.lock dependency resolution.
Key decisions:
- Load scoped thread snapshots for non-hydrate agent requests only when snapshot storage is active and no resume payload is present.
- Rebuild prior AG-UI history from stored snapshot messages, preserving the incoming new user suffix and treating stored snapshot content as authoritative over conflicting prior client history.
- Merge stored Shared State with request state overrides before schema defaults and existing state-context injection.

Files changed:
- packages/ag-ui/agent_framework_ag_ui/_agent_run.py
- packages/ag-ui/tests/ag_ui/test_endpoint.py

Verification:
- uv run pytest packages/ag-ui/tests/ag_ui/test_endpoint.py::test_agent_endpoint_prepends_stored_snapshot_for_new_user_turn -q
- uv run pytest packages/ag-ui/tests/ag_ui/test_endpoint.py::test_agent_endpoint_deduplicates_full_history_and_merges_fresh_state -q
- uv run pytest packages/ag-ui/tests/ag_ui/test_endpoint.py::test_endpoint_empty_messages packages/ag-ui/tests/ag_ui/test_endpoint.py::test_agent_endpoint_hydrates_stored_thread_snapshot_without_invoking_agent packages/ag-ui/tests/ag_ui/test_endpoint.py::test_agent_endpoint_hydrates_snapshots_by_scope_and_thread packages/ag-ui/tests/ag_ui/test_endpoint.py::test_agent_endpoint_prepends_stored_snapshot_for_new_user_turn packages/ag-ui/tests/ag_ui/test_endpoint.py::test_agent_endpoint_deduplicates_full_history_and_merges_fresh_state -q
- uv run pytest packages/ag-ui/tests/ag_ui/test_endpoint.py -q
- uv run poe syntax -P ag-ui -C
- uv run poe pyright -P ag-ui
- uv run poe test -P ag-ui
- uv run poe check -P ag-ui
- uv run poe typing -P ag-ui
- git diff --check
- git diff --cached --check

Blockers / next iteration:
- No blockers. Next slice can enable workflow AG-UI Thread Snapshot persistence and hydration.
- uv repeatedly refreshed azure-ai-projects in uv.lock during local runs; reverted the generated lockfile churn because this change does not alter dependencies.
- The poe-check commit hook was skipped after manual verification because it refreshes unrelated uv.lock dependency resolution.
Key decisions:
- Handle workflow Hydrate Requests before resolving or invoking the wrapped workflow when snapshot storage and Snapshot Scope are active.
- Capture only replayable workflow protocol data: workflow-emitted state snapshots, workflow-emitted message snapshots, and synthesized messages from text/tool output.
- Keep workflow snapshot capture inactive without configured persistence, and skip saving snapshots when the workflow stream emits RUN_ERROR.

Files changed:
- packages/ag-ui/agent_framework_ag_ui/_workflow.py
- packages/ag-ui/tests/ag_ui/test_endpoint.py

Verification:
- uv run pytest packages/ag-ui/tests/ag_ui/test_endpoint.py::test_workflow_endpoint_hydrates_emitted_snapshots_without_invoking_workflow packages/ag-ui/tests/ag_ui/test_endpoint.py::test_workflow_endpoint_hydrates_synthesized_text_and_tool_snapshot -q
- uv run pytest packages/ag-ui/tests/ag_ui/test_endpoint.py -q
- uv run pytest packages/ag-ui/tests/ag_ui/golden/test_scenario_workflow.py -q
- uv run poe syntax -P ag-ui -C
- uv run poe pyright -P ag-ui
- uv run poe test -P ag-ui
- uv run poe typing -P ag-ui
- uv run poe check -P ag-ui
- git diff --check
- git diff --cached --check

Blockers / next iteration:
- No blockers. Next slice can preserve interruption state and protect snapshots on errors across agent and workflow endpoints.
- uv repeatedly refreshed azure-ai-projects in uv.lock during local runs; reverted the generated lockfile churn because this change does not alter dependencies.
- The poe-check commit hook was skipped after manual verification because it refreshes unrelated uv.lock dependency resolution.
Key decisions:
- Capture workflow RUN_FINISHED interrupt metadata in replayable AG-UI Thread Snapshots so Hydrate Requests can restore pending workflow actions without invoking or resuming the workflow.
- Keep failed agent and workflow runs from replacing the last good snapshot; RUN_ERROR streams leave the previous snapshot available for hydration.
- Verify interruption hydration through endpoint-level AG-UI streams for both agent and workflow wrappers, including Shared State replay and no wrapped runner invocation.

Files changed:
- packages/ag-ui/agent_framework_ag_ui/_workflow.py
- packages/ag-ui/tests/ag_ui/test_endpoint.py

Verification:
- uv run pytest packages/ag-ui/tests/ag_ui/test_endpoint.py::test_workflow_endpoint_hydrates_interrupted_thread_without_invoking_workflow -q
- uv run pytest packages/ag-ui/tests/ag_ui/test_endpoint.py::test_agent_endpoint_hydrates_interrupted_thread_without_invoking_agent packages/ag-ui/tests/ag_ui/test_endpoint.py::test_agent_endpoint_run_error_does_not_overwrite_previous_snapshot packages/ag-ui/tests/ag_ui/test_endpoint.py::test_workflow_endpoint_hydrates_interrupted_thread_without_invoking_workflow packages/ag-ui/tests/ag_ui/test_endpoint.py::test_workflow_endpoint_run_error_does_not_overwrite_previous_snapshot -q
- uv run pytest packages/ag-ui/tests/ag_ui/test_endpoint.py -q
- uv run pytest packages/ag-ui/tests/ag_ui/golden/test_scenario_workflow.py -q
- uv run poe syntax -P ag-ui -C
- uv run poe pyright -P ag-ui
- uv run poe test -P ag-ui
- uv run poe typing -P ag-ui
- uv run poe check -P ag-ui
- git diff --check
- git diff --cached --check

Blockers / next iteration:
- No blockers. Next slice can document AG-UI Thread Snapshot security and usage.
- uv repeatedly refreshed azure-ai-projects in uv.lock during local runs; reverted the generated lockfile churn because this change does not alter dependencies.
- The poe-check commit hook was skipped after manual verification because it refreshes unrelated uv.lock dependency resolution.
Key decisions:
- Document AG-UI Thread Snapshot persistence as opt-in and disabled unless a snapshot_store is configured.
- Place Snapshot Scope guidance next to endpoint authentication guidance, making clear that AG-UI Thread ids identify threads but do not authorize snapshot access.
- Describe built-in storage as in-memory only, process-local, latest-only, and not durable production storage; durable stores remain app-owned implementations of AGUIThreadSnapshotStore.
- Call out snapshot confidentiality impact and that no file-backed AG-UI snapshot store is provided.

Files changed:
- packages/ag-ui/README.md

Verification:
- uv run python scripts/check_md_code_blocks.py packages/ag-ui/README.md --no-glob
- git diff --check
- git diff --cached --check
- commit hook without SKIP ran changed-package lint/format and AG-UI README markdown-code-lint successfully before stopping because uv.lock was modified
- uv run poe markdown-code-lint (failed due existing unrelated packages/mistral/README.md missing agent_framework_mistral import resolution; changed AG-UI README blocks passed)

Blockers / next iteration:
- No blockers. Local issue/PRD planning artifacts remain uncommitted.
- uv refreshed azure-ai-projects in uv.lock during markdown lint and the commit hook; reverted the generated lockfile churn because this documentation change does not alter dependencies.
- The poe-check commit hook was skipped after manual verification because it refreshes unrelated uv.lock dependency resolution.
- Persist the completed confirm_changes turn with interrupt=None so hydration
  no longer replays a stale pending interrupt after the user responds; resume
  requests prepend stored history so the persisted thread is not truncated.
- Defer endpoint default_state application to the runners when snapshot
  persistence is active, filling only keys missing from both the stored
  snapshot state and the request state so defaults never reset persisted
  Shared State.
- Always fold the turn's output into the persisted messages snapshot even when
  the outbound MESSAGES_SNAPSHOT event is suppressed for predictive tools
  without confirmation.
- Load the stored snapshot on workflow follow-up turns, reconstruct full
  thread history into the run input, and seed the snapshot builder with merged
  state so saving a new turn no longer replaces prior history.
- Move snapshot message reconstruction helpers to _run_common for reuse by the
  workflow runner; load stored agent snapshots on resume turns for state merge.
- Add endpoint regression tests for all four scenarios.
@moonbox3 moonbox3 self-assigned this Jun 11, 2026
Copilot AI review requested due to automatic review settings June 11, 2026 08:35
@moonbox3 moonbox3 added the documentation Improvements or additions to documentation label Jun 11, 2026
@github-actions

github-actions Bot commented Jun 11, 2026

Copy link
Copy Markdown
Contributor

Python Test Coverage

Python Test Coverage Report •
FileStmtsMissCoverMissing
packages/ag-ui/agent_framework_ag_ui
   _agent.py44295%68, 125
   _agent_run.py5945191%166–173, 212–213, 220, 329, 333, 335, 352, 379–380, 448–452, 577–579, 591–593, 691, 699, 758, 788–789, 926, 969, 986, 1003–1004, 1011, 1102, 1127, 1135, 1137, 1140, 1146, 1201, 1204, 1214–1215, 1222, 1268
   _endpoint.py103496%60, 176, 208–209
   _run_common.py3832892%413–414, 420–425, 735–736, 741, 743–745, 754, 762, 774, 776–779, 781–783, 810–811, 828, 839
   _snapshots.py68198%133
   _workflow.py1991691%53, 111, 116, 155, 194–195, 277, 320, 322, 324–326, 328, 346–348
packages/core/agent_framework
   _tools.py11388992%222–223, 400, 402, 415, 440–442, 450, 468, 482, 489, 496, 519, 521, 528, 536, 665, 699–701, 704–706, 708, 714, 765–767, 792, 818, 822, 860–862, 866, 888, 1031–1032, 1036, 1072, 1084, 1091–1094, 1115, 1119, 1123, 1137–1139, 1489, 1581, 1609, 1631, 1639, 1730, 1737–1738, 1797, 1801, 1847, 1908–1909, 1956, 1970, 1973, 1986, 1989, 2012, 2019, 2029, 2033, 2112, 2165, 2187, 2243, 2322, 2519, 2585–2586, 2738, 2743, 2745–2746, 2815, 2820, 2827
TOTAL39540450788% 

Python Unit Test Overview

Tests Skipped Failures Errors Time
7856 34 💤 0 ❌ 0 🔥 2m 7s ⏱️

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

This PR adds opt-in AG-UI thread snapshot persistence to allow FastAPI AG-UI endpoints (agent + workflow) to rehydrate a thread after refresh/reconnect while keeping server-side history authoritative. It introduces a snapshot store protocol, an in-memory implementation for dev/test, wires persistence into runners and endpoint setup (with a mandatory scope resolver), and adds extensive unit tests + documentation.

Changes:

  • Introduce AG-UI Thread Snapshot primitives (AGUIThreadSnapshot, AGUIThreadSnapshotStore) plus a bounded in-memory store implementation.
  • Add snapshot hydration + persistence logic to both agent and workflow runners and integrate snapshot configuration into the FastAPI endpoint helper (including deferred default-state handling when persistence is active).
  • Add tests and documentation for snapshot storage, exports, hydration behavior, and interrupt persistence.

Reviewed changes

Copilot reviewed 14 out of 14 changed files in this pull request and generated 4 comments.

Show a summary per file
File Description
python/packages/core/agent_framework/ag_ui/init.pyi Exposes snapshot primitives via the core agent_framework.ag_ui facade typing stub.
python/packages/core/agent_framework/ag_ui/init.py Exposes snapshot primitives via the core agent_framework.ag_ui facade runtime module.
python/packages/ag-ui/tests/ag_ui/test_snapshots.py New unit tests for snapshot datamodel + in-memory store behavior.
python/packages/ag-ui/tests/ag_ui/test_public_exports.py Verifies runtime exports for snapshot primitives in both agent_framework_ag_ui and agent_framework.ag_ui.
python/packages/ag-ui/tests/ag_ui/test_endpoint.py Adds end-to-end endpoint tests for hydration, dedupe/authoritative history, interrupts, and error behavior with persistence enabled.
python/packages/ag-ui/README.md Documents snapshot persistence, hydration semantics, and security model (scope resolver/authorization boundary).
python/packages/ag-ui/AGENTS.md Updates package overview to include snapshot store primitives.
python/packages/ag-ui/agent_framework_ag_ui/_workflow.py Adds workflow-side snapshot hydration + persistence and a workflow snapshot builder.
python/packages/ag-ui/agent_framework_ag_ui/_snapshots.py New module defining snapshot model/protocol, scope/thread types, constants, and in-memory store.
python/packages/ag-ui/agent_framework_ag_ui/_run_common.py Adds shared message reconstruction helpers used to merge stored history with new turns.
python/packages/ag-ui/agent_framework_ag_ui/_endpoint.py Adds endpoint parameters and validation for snapshot persistence (store + mandatory scope resolver) and defers defaults when persistence is active.
python/packages/ag-ui/agent_framework_ag_ui/_agent.py Adds snapshot store plumbing to agent config/wrapper.
python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py Adds agent-side hydration + persistence, history reconstruction, interrupt persistence/clearing, and state snapshot tracking.
python/packages/ag-ui/agent_framework_ag_ui/init.py Exports snapshot primitives from the agent_framework_ag_ui package root.

Comment thread python/packages/ag-ui/agent_framework_ag_ui/_workflow.py
Comment thread python/packages/ag-ui/agent_framework_ag_ui/_workflow.py Outdated
Comment thread python/packages/ag-ui/agent_framework_ag_ui/__init__.py
Comment thread python/packages/ag-ui/agent_framework_ag_ui/__init__.py

@github-actions github-actions Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Automated Code Review

Reviewers: 5 | Confidence: 89%

✓ Correctness

This is a well-designed opt-in snapshot persistence feature. The implementation correctly handles hydration, error recovery (not overwriting on RUN_ERROR), message reconstruction with deduplication against tampered history, state merging priority, and the confirm_changes interrupt-clearing flow. I found no correctness bugs after tracing through the key code paths and verifying against the test assertions.

✓ Security Reliability

The PR is well-designed from a security perspective with proper scope-based access control, mandatory scope resolvers, server-side authoritative history, and protection against tampered transcripts. The main reliability concern is that in the workflow path, snapshot_store.save() executes after all events (including RunFinishedEvent) have been yielded to the client. If the store's save fails, the exception propagates to the endpoint's error handler which emits an additional RUN_ERROR after the client already received RUN_FINISHED, violating the single-terminal-event protocol invariant.

✓ Test Coverage

The PR adds comprehensive integration tests for the main snapshot persistence paths (hydration, history reconstruction, interrupt persistence, error recovery) for both agent and workflow endpoints. However, there are a few notable test coverage gaps: no test exercises the async snapshot_scope_resolver branch, the InMemoryAGUIThreadSnapshotStore input validation error paths lack tests, and there's no test for workflow resume/interrupt with snapshot persistence configured.

✗ Failure Modes

The snapshot store's save() calls in both agent and workflow paths are not wrapped in exception handling. Since AGUIThreadSnapshotStore is a protocol designed for custom (e.g., database-backed) implementations, a transient save failure after all run events have already been streamed will propagate as an unhandled exception, causing the endpoint to emit RUN_ERROR instead of RUN_FINISHED. This converts a fully-completed, already-streamed agent turn into an apparent failure for the user, and the snapshot for that successful turn is lost. The same issue exists at both save sites (end-of-run and confirmation flow) in the agent path, and at the post-loop save in the workflow path.

✓ Design Approach

The snapshot persistence work is mostly aligned with the stated goals, but there is one blocking design gap in the workflow path: snapshot scope is treated as the authorization boundary for persisted data, yet workflow_factory instances are still cached only by thread_id. That means two different scopes reusing the same AG-UI thread id will share the same in-memory workflow object even though their snapshots are isolated by scope.

Flagged Issues

  • Unhandled exceptions from snapshot_store.save() convert successfully-completed runs into RUN_ERROR for the client. Since the AGUIThreadSnapshotStore protocol is designed for app-owned implementations (Redis, DynamoDB, etc.), transient save failures are expected in production. The save at _agent_run.py:1288 and _workflow.py:338 happen AFTER all run events have been yielded — the agent/workflow completed successfully — but BEFORE RUN_FINISHED. A failing save propagates to the endpoint's except Exception handler, which emits RUN_ERROR, misleading the client into thinking the run failed.

Suggestions

  • Add a test with an async snapshot_scope_resolver (e.g., async def resolve(request): return "tenant-a") to cover the isawaitable branch in _endpoint.py:142-143.

Automated review by moonbox3's agents

Comment thread python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py Outdated
Comment thread python/packages/ag-ui/agent_framework_ag_ui/_workflow.py Outdated
Comment thread python/packages/ag-ui/agent_framework_ag_ui/_workflow.py Outdated
@github-actions

Copy link
Copy Markdown
Contributor

Flagged issue

Unhandled exceptions from snapshot_store.save() convert successfully-completed runs into RUN_ERROR for the client. Since the AGUIThreadSnapshotStore protocol is designed for app-owned implementations (Redis, DynamoDB, etc.), transient save failures are expected in production. The save at _agent_run.py:1288 and _workflow.py:338 happen AFTER all run events have been yielded — the agent/workflow completed successfully — but BEFORE RUN_FINISHED. A failing save propagates to the endpoint's except Exception handler, which emits RUN_ERROR, misleading the client into thinking the run failed.


Source: automated DevFlow PR review

moonbox3 added 2 commits June 11, 2026 18:08
- Prepend stored thread history when persisting snapshots for resume runs on
  both the agent and workflow paths, so a resumed interrupt no longer
  overwrites the stored thread with just the resume turn's output.
- Filter the incoming message suffix during thread reconstruction: only user
  turns and tool results answering backend-issued tool calls (stored tool
  calls or pending interrupts) may extend authoritative history. Client-forged
  assistant and tool messages are dropped and logged instead of being
  persisted and replayed.
- Close the workflow snapshot builder's tool-call group when a tool result or
  text message lands, so synthesized transcripts keep tool results adjacent to
  their tool_calls message and stay valid as provider replay history.
- Export DEFAULT_MAX_THREAD_SNAPSHOTS from agent_framework_ag_ui and expose
  SnapshotScopeResolver through the core ag_ui facade and stub.
- Add regression tests for agent and workflow resume history preservation,
  forged suffix rejection, builder tool-call grouping, and the export surface.
- Wrap snapshot_store.save() on both the agent and workflow paths so a
  transient store failure (timeout, connection refused) is logged instead of
  propagating. Previously a failing save converted an already-streamed
  successful run into RUN_ERROR, and on the workflow path emitted RUN_ERROR
  after RUN_FINISHED, violating the single-terminal-event invariant. The
  previous snapshot stays available for hydration.
- Key the workflow_factory instance cache by (snapshot_scope, thread_id). The
  Snapshot Scope is the authorization boundary, so the same thread id under
  different scopes no longer shares an in-memory workflow instance.
  clear_thread_workflow accepts an optional snapshot_scope and clears all
  scopes for the thread when omitted.
- Add tests: save-failure tolerance for agent and workflow endpoints,
  scope-isolated workflow cache, async snapshot_scope_resolver support, and
  in-memory store key validation errors.
@moonbox3

Copy link
Copy Markdown
Contributor Author

Both review-level test suggestions from the automated review are addressed in 4f0e8c1: an async snapshot_scope_resolver now has coverage (test_endpoint_supports_async_snapshot_scope_resolver) and the in-memory store key validation error paths are tested (test_in_memory_snapshot_store_rejects_invalid_keys). The flagged save-failure and workflow cache scoping issues are fixed in the same commit; see the inline replies for details.

The existing ignore pattern only matched https://dotnet.microsoft.com/download,
but Microsoft sites insert a locale segment between host and path
(e.g. /en-us/download/dotnet/10.0), so localized links slip past the pattern
and get checked. dotnet.microsoft.com bot-blocks CI link checkers with
intermittent 403s across the whole site, which fails markdown-link-check on
unrelated pull requests since linkspector scans the entire repository.

Ignore the domain wholesale, matching how platform.openai.com is already
handled for the same reason. A 403 from bot blocking is indistinguishable
from a removed page, so the checker cannot produce a meaningful signal for
this domain either way.
Comment thread python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py Outdated
Comment thread python/packages/ag-ui/agent_framework_ag_ui/_snapshots.py Outdated
Copilot and others added 3 commits June 12, 2026 06:04
- Replace list(cast(...)) with a typed annotation for raw_messages
  (_agent_run.py:866) per review suggestion
- Replace OrderedDict with a plain dict in InMemoryAGUIThreadSnapshotStore
  (_snapshots.py:136); regular dicts are insertion-order-safe since
  Python 3.7, so OrderedDict is unnecessary. Update _evict_oldest to use
  next(iter(...)) for FIFO removal instead of popitem(last=False).

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
@moonbox3 moonbox3 requested a review from eavanvalkenburg June 12, 2026 07:11
@moonbox3 moonbox3 added this pull request to the merge queue Jun 12, 2026
Merged via the queue into microsoft:main with commit 76b2b1b Jun 12, 2026
38 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

ag-ui documentation Improvements or additions to documentation python

Projects

Status: Done

Development

Successfully merging this pull request may close these issues.

Python: sync AG-UI conversation history from backend

4 participants