Repository navigation
Python: State Management with AG-UI #3167
Description
Discussed in #2966
Originally posted by tschokokuki December 19, 2025
Dear agent-framework community,
my goal is to deterministically edit the state of an agent via a tool.
However, all I found in the demos was, how the agent state is edited by the predict_state via the predict_state_config.
As described in this demo: https://learn.microsoft.com/en-us/agent-framework/integrations/ag-ui/state-management?pivots=programming-language-python
However, the predict_state consists of non-deterministically generated tool arguments.
Thus, to my understanding predict state should be used for optimistic UI updates but not define final state.
Motivation:
- Lets say we want a get_weather tool to write the current weather into a global state shared with Backend and Frontend.
The state cannot be predicted by the tool call arguments by the llm, but must be retrieved during the tool call. - Furthermore for Human In The Loop (HITL) it would be required to use the Tool's returned results as a state, as the predict_state is emitted instantly as soon as available.
I found this pydantic_ai demo, where the tool emits a StateDeltaEvent, but the agent_framework tools would not let me return events.
Source: https://docs.copilotkit.ai/pydantic-ai/shared-state/predictive-state-updates
@agent.tool_plain
async def update_steps(steps: list[str]) -> StateSnapshotEvent:
"""Update the steps of the agent."""
return StateSnapshotEvent(
type=EventType.STATE_SNAPSHOT,
snapshot={
"observed_steps": steps
}
)How would this be done using agent_framework?
Thanks for your help!
Activity
LEDazzio01 commentedon Feb 25, 2026
I'd like to pick this up. I've been working on the AG-UI event emission layer in
_run_common.py(PRs #4200, #4233) and am familiar withFlowState,_emit_content, and the custom event patterns. Happy to look into deterministic state updates from tools.LEDazzio01 commentedon Feb 25, 2026
After looking into this, here's what I found about the current state of AG-UI state management:
What exists today:
PredictiveStateHandler(_orchestration/_predictive_state.py) — emitsStateDeltaEventoptimistically during streaming based on LLM-predicted tool call arguments and apredict_state_configmapping. This is the non-deterministic approach the issue describes._workflow_run.pyalready supportsStateSnapshotEventpassthrough. A workflow executor canyield_output(StateSnapshotEvent(snapshot={...}))and it flows through to the AG-UI consumer as-is.What's missing (the ask):
StateSnapshotEventorStateDeltaEvent. The_run_common.py_emit_contenthandler processesfunction_resultcontent but doesn't look for state payloads in tool return values.StateSnapshotEvent.Question for maintainers: What's the preferred approach here? Some options:
ContentType(e.g.,state_snapshot) that tools can return, which_emit_contentmaps toStateSnapshotEventHappy to implement once the design direction is clear.
tschokokuki commentedon Mar 26, 2026
I’m still looking for state changes to be driven by the actual tool output, rather than by optimistic input.
In #2966, Evan Mattson (@moonbox3) suggested overriding
_handle_function_result_content()inagent_framework_ag_ui/_events.py.From what I can tell, the correct place for this approach now would be
_emit_tool_result_common().(Version: 1.0.0.r5)