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 bb947873d..7cf949341 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 @@ -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 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 fe1f3bc4d..a34d0e9b2 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, @@ -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}})