ag-ui-protocol / ag-ui-protocol/ag-ui
ag-ui-langgraph: STATE_SNAPSHOT ignores LangGraph channel reducers during a fan-out
- Lenguaje dominante
- Python
- Estrellas
- 15.9k
- Forks
- 1.4k
- Merge medio
- 1 d 17 h
- PR fusionados (30 d)
- 163
Descripción
LangGraph channel reducers (`Annotated[list[T], operator.add]`, binary operators, custom reducers) are ignored when `ag_ui_langgraph` accumulates streamed state.
## What happens
In `_handle_stream_events`, each `on_chain_end` merges node output with a plain dict update:
```python
if event_type == "on_chain_end" and isinstance(
event.get("data", {}).get("output"), dict
):
output = event["data"]["output"]
current_graph_state.update(output)
```
`dict.update` overwrites the key. For a fan-out of N parallel tasks that each return `{plan_results: [one_item]}` onto a channel declared as `Annotated[list[PlanOutcome], operator.add]`, the in-memory accumulator keeps **only the last write**. Every subsequent `STATE_SNAPSHOT` is built from that lossy dict, so the client never sees the other N-1 items until the join node runs and a later `aget_state` read (REST `/state`, end-of-run snapshot) applies the real reducers.
LangGraph itself is correct: `Pregel.aget_state` calls `_prepare_state_snapshot(..., apply_pending_writes=True)` when the config has no pinned `checkpoint_id`, and that path replays `checkpoint_writes` through the channel reducers. The REST state read therefore already returns the full fan-out; only the streamed snapshots are wrong.
## Expected
Streamed `STATE_SNAPSHOT` events should apply the same merge as the checkpointer — at minimum, `operator.add` (and other declared reducers) on list channels — so a fan-out is visible incrementally.
Two reasonable fixes:
1. Stop accumulating with `dict.update`. After each node exit, snapshot from `await self.graph.aget_state(config)` (reducer-aware, already used at subgraph boundaries and run end).
2. If the in-memory accumulator must stay, apply each node's writes through the graph's channel reducers instead of `dict.update`.
## Workaround
We wrap `_handle_stream_events` in a thin post-processor that rebuilds each `StateSnapshotEvent` from `graph.aget_state` before yielding. That extra checkpoint read is per node-exit snapshot, not per token.
## Environment
- `ag-ui-langgraph==0.0.43`
- LangGraph with `AsyncPostgresSaver`
- Fan-out via `Send` onto an `operator.add` list channel
Guía de contribución
Evaluación
Este issue todavía no se ha evaluado.