Python: Avoid unchanged AG-UI predictive state snapshots (#7766)

* Python: Avoid unchanged AG-UI predictive state snapshots

Only emit the coalesced snapshot when predictive updates were actually pending or a deterministic state update was returned.

Assisted-by: Codex <codex@openai.com>

* Python: Exercise the predictive update path in snapshot tests

Use the handler streaming API to create pending state and narrow snapshot events by their concrete type.

Assisted-by: Codex <codex@openai.com>
This commit is contained in:
ALEX LILUZ
2026-08-20 05:51:13 +00:00
committed by GitHub
parent 435201b71b
commit 26b28b4386
2 changed files with 35 additions and 2 deletions
@@ -812,6 +812,7 @@ def _emit_tool_result_common(
# in stream order instead of grouping B with A (moonbox3's replay concern).
flow.snapshot_segments.append({"kind": "tool_results"})
had_pending_predictive_updates = bool(predictive_handler and predictive_handler.pending_state_updates)
if predictive_handler:
predictive_handler.apply_pending_updates()
@@ -824,7 +825,7 @@ def _emit_tool_result_common(
)
# Emit a single coalesced snapshot when either mechanism updated state.
if (predictive_handler or state_update) and flow.current_state:
if (had_pending_predictive_updates or state_update) and flow.current_state:
events.append(StateSnapshotEvent(snapshot=flow.current_state))
flow.tool_call_id = None
@@ -5,7 +5,7 @@
import logging
import pytest
from ag_ui.core import EventType
from ag_ui.core import EventType, StateSnapshotEvent
from ag_ui.core.events import (
ReasoningMessageContentEvent,
ReasoningMessageStartEvent,
@@ -380,6 +380,38 @@ class TestEmitToolResultWithState:
events = _emit_tool_result(content, flow)
assert all(e.type != EventType.STATE_SNAPSHOT for e in events)
def test_predictive_handler_without_pending_updates_emits_no_snapshot(self):
"""A configured predictive handler must not emit unchanged state for unrelated tools."""
flow = FlowState(current_state={"existing": "value"})
handler = PredictiveStateHandler(
predict_state_config={"draft": {"tool": "write_draft", "tool_argument": "body"}},
current_state=flow.current_state,
)
content = Content.from_function_result(call_id="c1", result="plain")
events = _emit_tool_result(content, flow, predictive_handler=handler)
assert all(e.type != EventType.STATE_SNAPSHOT for e in events)
assert flow.current_state == {"existing": "value"}
def test_predictive_handler_with_pending_updates_emits_snapshot(self):
"""A pending predictive update is applied and emitted as one snapshot."""
flow = FlowState(current_state={"existing": "value"})
handler = PredictiveStateHandler(
predict_state_config={"draft": {"tool": "write_draft", "tool_argument": "body"}},
current_state=flow.current_state,
)
deltas = handler.emit_streaming_deltas("write_draft", '{"body":"updated"}')
content = Content.from_function_result(call_id="c1", result="plain")
events = _emit_tool_result(content, flow, predictive_handler=handler)
assert len(deltas) == 1
snapshots = [event for event in events if isinstance(event, StateSnapshotEvent)]
assert len(snapshots) == 1
assert snapshots[0].snapshot == {"existing": "value", "draft": "updated"}
assert flow.current_state == {"existing": "value", "draft": "updated"}
def test_tool_result_content_text_unchanged(self):
"""The text sent to the LLM must not leak the state marker."""
tool_return = state_update(text="Weather: 14°C", state={"weather": {"temp": 14}})