Python: Avoid unchanged AG-UI predictive state snapshots

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"})
 
+    had_pending_predictive_updates = 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 (had_pending_predictive_updates or state_update) and flow.current_state:
         events.append(StateSnapshotEvent(snapshot=flow.current_state))
 
     flow.tool_call_id = None

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,
@@ -380,6 +380,38 @@ 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,
+        )
+        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}})