Python: [Bug]: with predict_state_config set, EVERY tool result emits a full StateSnapshotEvent — including for tools the config never names #7490

Description

1. Symptom

An AG-UI app that configures predict_state_config for a single tool emits one StateSnapshotEvent — carrying the whole shared state — after every tool result in the run, no matter which tool ran. In a deployed app with a 16-key-deep shared state, that was 1,632 snapshots against 1,600 TOOL_CALL_START events over two days: a 1:1 ratio with tool calls, not with predictive-state activity (the predictive tool was called a handful of times).

Because each snapshot serializes the entire flow.current_state, the per-tool-call cost grows with the conversation: in one 95-minute session the snapshots grew 7 KB → 15.9 KB each, and StateSnapshotEvent alone accounted for 44% of all bytes the server streamed to the browser (STATE_SNAPSHOT + STATE_DELTA together: 71%; assistant text: 3%).

2. Cause

_emit_tool_result_common decides whether to emit the snapshot by testing whether a predictive handler exists, not whether it updated anything — as its own comment says it intends (python/packages/ag-ui/agent_framework_ag_ui/_run_common.py:735-748):

if predictive_handler:
    predictive_handler.apply_pending_updates()

if state_update:
    flow.current_state.update(state_update)
    ...

# Emit a single coalesced snapshot when either mechanism updated state.
if (predictive_handler or state_update) and flow.current_state:
    events.append(StateSnapshotEvent(snapshot=flow.current_state))

The handler is constructed once per run whenever the app passes any config at all (_agent_run.py:1947-1952):

predictive_handler: PredictiveStateHandler | None = None
if predict_state_config:
    predictive_handler = PredictiveStateHandler(...)

and it only ever accumulates work for its configured tool — both delta paths skip everything else (_orchestration/_predictive_state.py:142-143 and :177-178):

for state_key, config in self.predict_state_config.items():
    if config["tool"] != tool_name:
        continue

So for any other tool, pending_state_updates is empty, apply_pending_updates() is a no-op (:228-232), and the emit still fires. The information needed to make the condition match the comment is already in hand: pending_state_updates is non-empty exactly when the predictive mechanism changed state.

3. Minimal repro

docs/issues/repros/repro_agui_predictive_snapshot_per_tool_result.py in this repository — public API only (AgentFrameworkAgent.run), deterministic, no LLM and no network. A fake chat client calls look_up (never the predictive tool update_mode), and the run is executed twice, differing only in whether predict_state_config is passed:

predict_state_config set    : 2 tool results -> 3 StateSnapshotEvent(s), 177 bytes of state re-sent
predict_state_config absent : 2 tool results -> 1 StateSnapshotEvent(s), 59 bytes of state re-sent
predict_state_config set    : 4 tool results -> 5 StateSnapshotEvent(s), 295 bytes of state re-sent
predict_state_config absent : 4 tool results -> 1 StateSnapshotEvent(s), 59 bytes of state re-sent

One snapshot in each run is the run-boundary one both configurations emit; the remainder are one per tool result, and they scale exactly with the number of tool calls. update_mode is never invoked in either run.

4. Impact

Every extra snapshot is a full re-serialization of shared state on the server, a full re-transmission over SSE, and — on an AG-UI client — a state-changed notification that re-renders whatever subscribes to it. All three costs are borne per tool call and scale with the size of the shared state, which typically grows with the conversation. The measurements in §1 come from a production deployment; the effect a user reports is a UI that gets progressively slower the longer a conversation runs.

Two properties make this hard for an app to mitigate on its own: the emit happens inside the event pipeline, below anything the app can intercept, and opting out means giving up predictive state entirely for the one tool that legitimately uses it.

Code Sample

Error Messages / Stack Traces

Package Versions

agent-framework-ag-ui: 1.0.1, agent-framework-core: 1.13.0

Python Version

No response

Additional Context

5. Suggested fix

Make the condition express the comment: emit when the predictive handler actually applied pending updates, or when the tool result carried a deterministic state_update. PredictiveStateHandler.apply_pending_updates already knows — having it report whether it applied anything (or checking pending_state_updates before the call) is enough for the guard, and leaves the deterministic state_update path untouched.

That keeps every behaviour that depends on a snapshot following a real state change, and removes only the snapshots that repeat a state nothing has modified. A narrower variant — skipping the emit when the just-finished tool is not named by any entry in predict_state_config — would fix the reported case too, but the pending-updates check also covers a predictive tool whose arguments produced no change.

A related question for the same area: when the emit is warranted, the event carries the entire state rather than the keys that changed, so a large shared state is re-sent in full for a one-key update. StateDeltaEvent already exists for that shape and the predictive path uses it while streaming arguments; using it at the tool-result boundary too would bound the payload by the size of the change instead of the size of the conversation.