From c7a883c07ee3e6dce66a4fbafe982a42e7c67af2 Mon Sep 17 00:00:00 2001 From: alexliluz <49665315+alexliluz@users.noreply.github.com> Date: Wed, 19 Aug 2026 21:47:06 +0800 Subject: [PATCH 1/2] 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 --- .../agent_framework_ag_ui/_run_common.py | 3 +- .../ag-ui/tests/ag_ui/test_run_common.py | 31 +++++++++++++++++++ 2 files changed, 33 insertions(+), 1 deletion(-) diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_run_common.py b/python/packages/ag-ui/agent_framework_ag_ui/_run_common.py index 3606428f600..26241c120f6 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_run_common.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_run_common.py @@ -783,6 +783,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"}) + predictive_state_updated = bool(predictive_handler and predictive_handler.pending_state_updates) if predictive_handler: predictive_handler.apply_pending_updates() @@ -795,7 +796,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 (predictive_state_updated or state_update) and flow.current_state: events.append(StateSnapshotEvent(snapshot=flow.current_state)) flow.tool_call_id = None diff --git a/python/packages/ag-ui/tests/ag_ui/test_run_common.py b/python/packages/ag-ui/tests/ag_ui/test_run_common.py index fe1f3bc4dce..ce4a1ed480b 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_run_common.py +++ b/python/packages/ag-ui/tests/ag_ui/test_run_common.py @@ -380,6 +380,37 @@ def test_no_state_snapshot_when_result_has_no_state(self): 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, + ) + handler.pending_state_updates["draft"] = "updated" + content = Content.from_function_result(call_id="c1", result="plain") + + events = _emit_tool_result(content, flow, predictive_handler=handler) + + snapshots = [event for event in events if event.type == EventType.STATE_SNAPSHOT] + assert len(snapshots) == 1 + assert snapshots[0].snapshot == {"existing": "value", "draft": "updated"} # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + 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}}) From a4af8bc88d0d7a38ee5fcd881397ddb8f872895a Mon Sep 17 00:00:00 2001 From: alexliluz <49665315+alexliluz@users.noreply.github.com> Date: Wed, 19 Aug 2026 21:52:33 +0800 Subject: [PATCH 2/2] 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 --- .../packages/ag-ui/agent_framework_ag_ui/_run_common.py | 4 ++-- python/packages/ag-ui/tests/ag_ui/test_run_common.py | 9 +++++---- 2 files changed, 7 insertions(+), 6 deletions(-) diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_run_common.py b/python/packages/ag-ui/agent_framework_ag_ui/_run_common.py index 26241c120f6..a83917e6dca 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_run_common.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_run_common.py @@ -783,7 +783,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"}) - predictive_state_updated = bool(predictive_handler and predictive_handler.pending_state_updates) + had_pending_predictive_updates = bool(predictive_handler and predictive_handler.pending_state_updates) if predictive_handler: predictive_handler.apply_pending_updates() @@ -796,7 +796,7 @@ def _emit_tool_result_common( ) # Emit a single coalesced snapshot when either mechanism updated state. - if (predictive_state_updated 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 diff --git a/python/packages/ag-ui/tests/ag_ui/test_run_common.py b/python/packages/ag-ui/tests/ag_ui/test_run_common.py index ce4a1ed480b..a34d0e9b2f3 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_run_common.py +++ b/python/packages/ag-ui/tests/ag_ui/test_run_common.py @@ -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, @@ -401,14 +401,15 @@ def test_predictive_handler_with_pending_updates_emits_snapshot(self): predict_state_config={"draft": {"tool": "write_draft", "tool_argument": "body"}}, current_state=flow.current_state, ) - handler.pending_state_updates["draft"] = "updated" + 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) - snapshots = [event for event in events if event.type == EventType.STATE_SNAPSHOT] + 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"} # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + assert snapshots[0].snapshot == {"existing": "value", "draft": "updated"} assert flow.current_state == {"existing": "value", "draft": "updated"} def test_tool_result_content_text_unchanged(self):