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:
@@ -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}})
|
||||
|
||||
Reference in New Issue
Block a user