State Management with AG-UI

[tschokokuki]

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!

[moonbox3]

Hi @tschokokuki, thanks for bringing this up. It is a current gap. As a potential workaround, one could intercept the agent stream and inject StateSnapshotEvent when specific FunctionResultContent is available. Something like:

from ag_ui.core import StateSnapshotEvent

async def intercept_stream():
    async for update in agent.run_stream(messages):
        yield update
        for content in update.contents:
            if isinstance(content, FunctionResultContent) and content.name == "get_weather":
                # Emit state based on actual tool result
                yield StateSnapshotEvent(snapshot={"weather": content.result})

Otherwise we will need to add handling for tools to directly emit state events. Would you be interested in contributing to this? The relevant code is in agent_framework_ag_ui/_events.py, specifically _handle_function_result_content() which currently doesn't extract state from tool results.

[tschokokuki]

Hi @moonbox3,

Thank you very much for pointing me in the right direction.

I got the workaround running as you suggested.
I also began with extending the _handle_function_result_content().

From my side, I added one more requirement: I want to pass a custom message for the LLM rather than relying on a default one.

I wonder, what would be the best return format for the Tools in context of the agent_framework project.

Option 1: StateDeltaEvent

Pro:

• any Event could be emitted in _handle_function_result_content() without conversions

Cons:

• The StateDeltaEvent seems to create issues when serializing the agent`s thread state. => result: null
• result_message adds a new convention which must be picked up in _handle_function_result_content()

@ai_function()
def get_weather() -> StateDeltaEvent:
    weather = "sunny"
    return StateDeltaEvent(
        delta=[{"op": "replace", "path": "/weather", "value": weather}],
        result_message=f"Weather updated to: {weather}")

Option 2: Dictionary with a marker

Pro:

• The entire Dict is compatible with storing it as a result in the agent' thread state => result = "{ "agui_event": True, "event_type": ..... }"

Cons:

• Actually storing only the "result_message" as a result in thread state would be sufficient
  => This creates unwanted artefacts when converting the result from agent_framework_messages to ag_ui_messages
  => Thus we could extend the converting function to filter out events, which adds more complexity.

@ai_function()
def get_weather() -> Dict:
    weather = "sunny"
    return {
        "__agui_event__": True,
        "event_type": "StateDeltaEvent",
        "event_data": {
            "delta": [{"op": "replace", "path": "/weather", "value": weather}]
        },
        "result_message": f"Weather updated to: {weather}"
    }

In my opinion 1 would be preferable.
But I haven’t looked more deeply yet into when agent_framework loses the results, supposedly due to a serialization failure.

Before dedicating myself to the final solution, what would be the prefered option in your opinion?
Basically its about calculating the efforts for mitigating the contra arguments, which I cannot yet estimate at the moment.

Thanks a lot!

[tschokokuki]

@moonbox3 . Could you please give me a hint, which Option to go for?

[moonbox3]

@tschokokuki Those are both valid options; however, I believe a third approach may be work better.

A few comments about the options you listed:

Option 1 (returning StateDeltaEvent directly) has a fundamental issue: FunctionResultContent.result expects JSON-serializable data, not event objects. This is why you're seeing serialization failures and result: null in thread state. The result_message field is a workaround that introduces coupling between the event schema and the result schema.

Option 2 (Dict with marker) avoids the serialization issue but introduces sentinel detection ( __agui_event__ ), pollutes the result schema, and creates artifacts in thread state that downstream consumers must handle.

I would recommend a structured tool result with state. It's a cleaner pattern that uses a structured return type that separates concerns:

from typing import TypedDict, Any
from dataclasses import dataclass

@dataclass
class ToolResultWithState:
    """Structured return type for tools that emit state."""
    result: Any  # The actual result for LLM context
    state_delta: list[dict[str, Any]] | None = None  # JSON Patch ops
    state_snapshot: dict[str, Any] | None = None  # Full snapshot override

@ai_function()
def get_weather(location: str) -> ToolResultWithState:
    weather = fetch_weather_api(location)  # Actual runtime data
    return ToolResultWithState(
        result=f"Weather in {location}: {weather}",
        state_delta=[{"op": "replace", "path": "/weather", "value": weather}],
    )

Then the implementation in _handle_function_result_content():

def _handle_function_result_content(self, content: FunctionResultContent) -> list[BaseEvent]:
    events: list[BaseEvent] = []
    
    # Extract state events if result is ToolResultWithState
    result_for_llm = content.result
    if isinstance(content.result, dict):
        if "state_delta" in content.result or "state_snapshot" in content.result:
            if content.result.get("state_delta"):
                events.append(StateDeltaEvent(delta=content.result["state_delta"]))
            if content.result.get("state_snapshot"):
                events.append(StateSnapshotEvent(snapshot=content.result["state_snapshot"]))
            result_for_llm = content.result.get("result", "")
    
    # Continue with existing logic using result_for_llm
    ...

Why I think this works better:

1. Clean separation: result goes to the LLM for context, state_* fields go to the UI via AG-UI events
2. Fully serializable: thread state stores the complete dict, which is valid JSON—no serialization failures
3. No magic markers: no sentinel strings like __agui_event__ to detect
4. Type-safe: proper dataclass/TypedDict with IDE support and validation
5. Backward compatible: existing tools without state fields continue working unchanged

The key insight is that _handle_function_result_content() already receives the full result, it just needs to check for known keys (state_delta, state_snapshot) and emit the appropriate AG-UI events before constructing the ToolCallResultEvent. The LLM only sees the result field, while the UI receives the state events.

Next Steps:

This will require changes to the ag-ui package to support natively. I'll file an issue to track this work. In the meantime, you can unblock yourself by either:

• Overriding _handle_function_result_content() in a custom subclass of AgentFrameworkEventBridge
• Using the stream interception workaround mentioned earlier

Once the library support lands, you can migrate to the cleaner ToolResultWithState API.

What are your thoughts?

[tschokokuki]

@moonbox3 : Thank you for your detailed response! I agree with your proposed approach and can see its advantages.

One question for clarification:
With "required changes to ag-ui package" you refer to creating a pull request for updating _handle_function_result_content() (among other things) in the agent_framework_ag_ui/_events.py file?