Python: State Management with AG-UI

[moonbox3]

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!

[LEDazzio01]

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 with FlowState, _emit_content, and the custom event patterns. Happy to look into deterministic state updates from tools.

[LEDazzio01]

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) — emits StateDeltaEvent optimistically during streaming based on LLM-predicted tool call arguments and a predict_state_config mapping. This is the non-deterministic approach the issue describes.
• Workflow path — _workflow_run.py already supports StateSnapshotEvent passthrough. A workflow executor can yield_output(StateSnapshotEvent(snapshot={...})) and it flows through to the AG-UI consumer as-is.

What's missing (the ask):

• Agent path — there's no mechanism for a tool's execution result to emit StateSnapshotEvent or StateDeltaEvent. The _run_common.py _emit_content handler processes function_result content but doesn't look for state payloads in tool return values.
• A tool returning state-relevant data has no way to signal "update shared state with this" — unlike the pydantic_ai pattern where a tool can directly return a StateSnapshotEvent.

Question for maintainers: What's the preferred approach here? Some options:

1. A new ContentType (e.g., state_snapshot) that tools can return, which _emit_content maps to StateSnapshotEvent
2. A callback/hook on the agent or tool level that emits state events post-execution
3. Something else entirely

Happy to implement once the design direction is clear.

[tschokokuki]

I’m still looking for state changes to be driven by the actual tool output, rather than by optimistic input.

In #2966, @moonbox3 suggested overriding _handle_function_result_content() in agent_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)