Skip to content

Python: fix(python): separate provider input from UI snapshot in AG-UI service-session mode - #7770

Open
pratik wayase (PratikWayase) wants to merge 2 commits into
microsoft:mainfrom
PratikWayase:fix/7710-agui-snapshot-block
Open

Python: fix(python): separate provider input from UI snapshot in AG-UI service-session mode#7770
pratik wayase (PratikWayase) wants to merge 2 commits into
microsoft:mainfrom
PratikWayase:fix/7710-agui-snapshot-block

Conversation

@PratikWayase

Copy link
Copy Markdown
Contributor

Motivation & Context

When AgentFrameworkAgent is configured with both use_service_session=True and a snapshot_store, the AG-UI runner unconditionally reconstructs the full stored message history and passes it to the wrapped agent alongside the service_session_id. Foundry Hosted Agents using the Responses protocol reject replayed assistant messages with an HTTP 400 error, because the server already owns conversation history via the service session. This blocks multi-turn applications that require both Foundry-managed conversation continuation and AG-UI snapshot persistence for UI hydration or approval state.

This change enforces a split-authority model consistent with the framework’s “one State Authority” contract and analogous to the fix previously applied to AgentExecutor in #7682.

Fixes #7710

Description & Review Guide

  • What are the major changes?

    • In run_agent_stream, message reconstruction from the thread snapshot is now conditional on config.use_service_session. When service-session mode is active, raw_messages (sent to the provider) contains only incremental input, while snapshot_messages (used for UI hydration and persistence) retains the fully reconstructed transcript.
    • Added a defense-in-depth guard ensuring that in service-session mode, the current turn’s output is always folded into latest_messages_snapshot before saving, even when should_emit_snapshot evaluates to False. This prevents UI history collapse when agents emit content types that don’t trigger the standard snapshot emission path.
    • Added regression test test_service_session_snapshot_split_authority verifying both halves of the split: provider receives incremental-only input, and the snapshot store retains the complete transcript.
  • What is the impact of these changes?

    • Unblocks the integration stack: AG-UI client → AgentFrameworkAgentFoundryAgent → Foundry Hosted Agent (Responses protocol) with snapshot persistence enabled.
    • No behavioral change for prompt-based Foundry agents or any configuration where use_service_session=False.
    • Existing snapshot hydration, approval resume, and predictive state flows continue to work unchanged; the full history is still available in snapshot_messages and persisted to the store.
  • What do you want reviewers to focus on?

    • Correctness of the conditional guards around _reconstruct_messages_from_thread_snapshot and resume_seeded_messages — specifically that no code path leaks reconstructed history into raw_messages when use_service_session=True.
    • Whether the defense-in-depth block at the snapshot save point correctly handles edge cases where should_emit_snapshot is False but the UI still needs the current turn’s output.
    • Adequacy of the regression test coverage for both normal turns and the snapshot hydration assertion.

Related Issue

Fixes #7710

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.
  • This is not a breaking change.

Copilot AI balanced review requested due to automatic review settings August 19, 2026 18:24
@agent-framework-automation agent-framework-automation Bot added the python Usage: [Issues, PRs], Target: Python label Aug 19, 2026
@github-actions github-actions Bot changed the title fix(python): separate provider input from UI snapshot in AG-UI service-session mode Python: fix(python): separate provider input from UI snapshot in AG-UI service-session mode Aug 19, 2026

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

Separates service-managed provider history from AG-UI snapshot history.

Changes:

  • Prevents stored snapshot replay into service-session provider calls.
  • Reconstructs full snapshots for UI hydration.
  • Adds split-authority regression coverage.

Reviewed changes

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

File Description
_agent_run.py Separates provider input and persisted snapshots.
test_snapshots.py Tests incremental provider input and full snapshot history.

💡 Add a code-review agent skill for context-aware, tailored reviews. Learn more in the docs.

Comment thread python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py
Comment thread python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py Outdated
stored_interrupt=stored_snapshot.interrupt,
)
else:
stored_count = len(stored_snapshot.messages)

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.

Could we avoid using len(stored_snapshot.messages) as an unconditional cursor? It assumes every request contains an exact, current stored prefix. Incremental requests then lose the new turn, direct confirm_changes responses can truncate persisted history, and a stale snapshot after a failed write can replay an already-processed turn into the service session. Could we derive a validated suffix by identity while keeping the full input for snapshot reconstruction?

seeded_resume_from_snapshot = True
else:

if not config.use_service_session:

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.

Would it make sense for one snapshot/session helper to produce both views for a run? The new branches at _agent_run.py:2273-2286 use a positional len(stored_snapshot.messages) cursor for provider input, while _reconstruct_messages_from_thread_snapshot() at _agent_run.py:2389-2397 separately decides which messages belong in the UI snapshot. Those paths already disagree for incremental requests and approval responses, so a helper such as _split_service_session_input(...) returning the provider suffix and persisted snapshot would keep the prefix, suffix, and resume rules in one place.

@moonbox3 Evan Mattson (moonbox3) added the ag-ui Usage: [Issues, PRs], Target: AG-UI protocol integration label Aug 20, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

ag-ui Usage: [Issues, PRs], Target: AG-UI protocol integration python Usage: [Issues, PRs], Target: Python

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Python: AG-UI snapshot replay with service sessions blocks Foundry Hosted Agents

3 participants