Skip to content

Python: AG-UI deterministic state updates from tool results#5201

Merged
moonbox3 merged 2 commits intomicrosoft:mainfrom
moonbox3:3167-fix
Apr 14, 2026
Merged

Python: AG-UI deterministic state updates from tool results#5201
moonbox3 merged 2 commits intomicrosoft:mainfrom
moonbox3:3167-fix

Conversation

@moonbox3
Copy link
Copy Markdown
Contributor

Motivation and Context

  • Fixes Python: State Management with AG-UI #3167 — AG-UI agents had no way to push a deterministic state update from a tool's actual return value. The only built-in path, predict_state_config, is driven optimistically by LLM-predicted tool arguments, which is wrong when the real state depends on what the tool actually computes (e.g. get_weather needs to fetch data and put that in state).
  • Adds a new state_update(text, state=...) helper in agent_framework_ag_ui. Tools return it as a Content; the snapshot rides in additional_properties under a reserved key, survives parse_result / from_function_result round-tripping, and is extracted by the AG-UI emitter which merges it into FlowState.current_state and emits a StateSnapshotEvent right after the ToolCallResult event.
  • Orthogonal to predict_state_config — both mechanisms can fire on the same tool result, with predictive state applied first.
  • Non-breaking and opt-in: zero behavior change for tools that don't call state_update().

Description

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 moonbox3 self-assigned this Apr 10, 2026
Copilot AI review requested due to automatic review settings April 10, 2026 05:45
@moonbox3 moonbox3 moved this to In Review in Agent Framework Apr 10, 2026
@moonbox3
Copy link
Copy Markdown
Contributor Author

moonbox3 commented Apr 10, 2026

Python Test Coverage

Python Test Coverage Report •
FileStmtsMissCoverMissing
packages/ag-ui/agent_framework_ag_ui
   _run_common.py3071096%381–382, 388–393, 701–702
   _state.py110100% 
packages/ag-ui/agent_framework_ag_ui_examples/agents
   weather_state_agent.py15150%16, 18, 20–21, 23, 27, 36–37, 53–54, 58–59, 68, 70, 82
packages/ag-ui/agent_framework_ag_ui_examples/server
   main.py64640%5, 7–9, 11–17, 19–27, 29–32, 34, 36, 39–40, 43–45, 48–50, 52–54, 56–58, 61–62, 64–65, 67, 69, 80, 88, 95, 102, 109, 116, 125, 132, 139, 147, 154, 156–157, 159–160, 164, 172–173
TOTAL27249320188% 

Python Unit Test Overview

Tests Skipped Failures Errors Time
5466 20 💤 0 ❌ 0 🔥 1m 28s ⏱️

Copy link
Copy Markdown
Contributor

Copilot AI left a comment

Choose a reason for hiding this comment

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

Pull request overview

Adds an opt-in mechanism for AG-UI tools to emit deterministic state snapshots derived from the tool’s actual return value (vs. optimistic predict_state_config based on predicted args), addressing agent-framework issue #3167.

Changes:

  • Introduces a public state_update(text, state=...) helper that encodes a state snapshot into Content.additional_properties.
  • Extracts tool-result state markers during AG-UI emission, merges into FlowState.current_state, and emits STATE_SNAPSHOT after tool results (including MCP tool results).
  • Adds unit + end-to-end + golden tests and an examples endpoint demonstrating deterministic tool-driven state.

Reviewed changes

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

Show a summary per file
File Description
python/packages/core/agent_framework/ag_ui/init.pyi Exposes state_update in the core AG-UI typing surface.
python/packages/core/agent_framework/ag_ui/init.py Adds state_update to lazy exports (but still has export-list mismatches vs stub).
python/packages/ag-ui/agent_framework_ag_ui/_state.py New helper + reserved marker key for deterministic state updates.
python/packages/ag-ui/agent_framework_ag_ui/_run_common.py Extracts/merges tool-result state and emits STATE_SNAPSHOT events.
python/packages/ag-ui/agent_framework_ag_ui/init.py Re-exports state_update from the runtime package.
python/packages/ag-ui/agent_framework_ag_ui_examples/server/main.py Registers a new /deterministic_state example endpoint.
python/packages/ag-ui/agent_framework_ag_ui_examples/agents/weather_state_agent.py New example agent showing tool-driven deterministic state.
python/packages/ag-ui/tests/ag_ui/test_run_common.py Adds focused tests for helper, extraction, and emission behavior.
python/packages/ag-ui/tests/ag_ui/test_public_exports.py Verifies state_update is exported publicly.
python/packages/ag-ui/tests/ag_ui/test_agent_wrapper_comprehensive.py End-to-end regression coverage through real tool invocation.
python/packages/ag-ui/tests/ag_ui/golden/test_scenario_deterministic_state.py Golden stream tests pinning user-visible event ordering/content.

1. Add missing AGUIEventConverter, AGUIHttpService, __version__ to
   _IMPORTS in core ag_ui lazy-export list to match the .pyi stub.

2. Coalesce predictive and deterministic state snapshots into a single
   StateSnapshotEvent when both mechanisms are active on the same tool
   result, reducing redundant snapshot traffic.

3. Update state_update() docstring to clarify that a predictive snapshot
   may be emitted before the deterministic one when predict_state_config
   is active.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
@moonbox3 moonbox3 added this pull request to the merge queue Apr 14, 2026
Merged via the queue into microsoft:main with commit f183f88 Apr 14, 2026
31 checks passed
@github-project-automation github-project-automation bot moved this from In Review to Done in Agent Framework Apr 14, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

Status: Done

Development

Successfully merging this pull request may close these issues.

Python: State Management with AG-UI

4 participants