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!