ag-ui-protocol / ag-ui-protocol/ag-ui

ag-ui-langgraph: STATE_SNAPSHOT ignores LangGraph channel reducers during a fan-out

Abierto
#2,628 0 comentarios 0 reacciones 0 asignados Ver en GitHub
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

Abrir la guía de contribución

Evaluación

Este issue todavía no se ha evaluado.

Recibe los nuevos issues en tu correo

Un resumen breve de issues de GitHub para principiantes.