mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-09-23 18:15:08 +02:00
fix(langgraph): hydrate subgraph delta channels with the resolved saver
A subgraph has no saver of its own: it uses the parent's, passed through `CONFIG_KEY_CHECKPOINTER`, or holds `checkpointer=True`. Every state reader resolved that saver, then called `_prepare_state_snapshot`, which re-derived it from `self.checkpointer` and got `None`. A `DeltaChannel` stores nothing in `channel_values`, so without a saver it hydrated empty while plain channels in the same snapshot read correctly. `bulk_update_state` had the same expression, and once an update reached the channel's snapshot cadence it persisted the empty value as a `_DeltaSnapshot`, losing history on disk. Pass the caller's saver in as a required keyword argument. Fixes #8470 Co-authored-by: gururafiki <22777967+gururafiki@users.noreply.github.com> Co-authored-by: Yuan Gao <119447586+DavidGao520@users.noreply.github.com>
This commit is contained in:
co-authored by
gururafiki
Yuan Gao
parent
bdb85b5aa8
commit
23331a8389
@@ -1146,6 +1146,8 @@ class Pregel(
|
||||
self,
|
||||
config: RunnableConfig,
|
||||
saved: CheckpointTuple | None,
|
||||
*,
|
||||
saver: BaseCheckpointSaver,
|
||||
recurse: BaseCheckpointSaver | None = None,
|
||||
apply_pending_writes: bool = False,
|
||||
) -> StateSnapshot:
|
||||
@@ -1169,9 +1171,7 @@ class Pregel(
|
||||
channels, managed = channels_from_checkpoint(
|
||||
self.channels,
|
||||
saved.checkpoint,
|
||||
saver=self.checkpointer
|
||||
if isinstance(self.checkpointer, BaseCheckpointSaver)
|
||||
else None,
|
||||
saver=saver,
|
||||
config=saved.config,
|
||||
)
|
||||
# tasks for this checkpoint
|
||||
@@ -1186,11 +1186,7 @@ class Pregel(
|
||||
stop,
|
||||
for_execution=True,
|
||||
store=self.store,
|
||||
checkpointer=(
|
||||
self.checkpointer
|
||||
if isinstance(self.checkpointer, BaseCheckpointSaver)
|
||||
else None
|
||||
),
|
||||
checkpointer=saver,
|
||||
manager=None,
|
||||
)
|
||||
# get the subgraphs
|
||||
@@ -1269,6 +1265,8 @@ class Pregel(
|
||||
self,
|
||||
config: RunnableConfig,
|
||||
saved: CheckpointTuple | None,
|
||||
*,
|
||||
saver: BaseCheckpointSaver,
|
||||
recurse: BaseCheckpointSaver | None = None,
|
||||
apply_pending_writes: bool = False,
|
||||
) -> StateSnapshot:
|
||||
@@ -1292,9 +1290,7 @@ class Pregel(
|
||||
channels, managed = await achannels_from_checkpoint(
|
||||
self.channels,
|
||||
saved.checkpoint,
|
||||
saver=self.checkpointer
|
||||
if isinstance(self.checkpointer, BaseCheckpointSaver)
|
||||
else None,
|
||||
saver=saver,
|
||||
config=saved.config,
|
||||
)
|
||||
# tasks for this checkpoint
|
||||
@@ -1309,11 +1305,7 @@ class Pregel(
|
||||
stop,
|
||||
for_execution=True,
|
||||
store=self.store,
|
||||
checkpointer=(
|
||||
self.checkpointer
|
||||
if isinstance(self.checkpointer, BaseCheckpointSaver)
|
||||
else None
|
||||
),
|
||||
checkpointer=saver,
|
||||
manager=None,
|
||||
)
|
||||
# get the subgraphs
|
||||
@@ -1429,6 +1421,7 @@ class Pregel(
|
||||
return self._prepare_state_snapshot(
|
||||
config,
|
||||
saved,
|
||||
saver=checkpointer,
|
||||
recurse=checkpointer if subgraphs else None,
|
||||
apply_pending_writes=CONFIG_KEY_CHECKPOINT_ID not in config[CONF],
|
||||
)
|
||||
@@ -1473,6 +1466,7 @@ class Pregel(
|
||||
return await self._aprepare_state_snapshot(
|
||||
config,
|
||||
saved,
|
||||
saver=checkpointer,
|
||||
recurse=checkpointer if subgraphs else None,
|
||||
apply_pending_writes=CONFIG_KEY_CHECKPOINT_ID not in config[CONF],
|
||||
)
|
||||
@@ -1527,7 +1521,7 @@ class Pregel(
|
||||
checkpointer.list(config, before=before, limit=limit, filter=filter)
|
||||
):
|
||||
yield self._prepare_state_snapshot(
|
||||
checkpoint_tuple.config, checkpoint_tuple
|
||||
checkpoint_tuple.config, checkpoint_tuple, saver=checkpointer
|
||||
)
|
||||
|
||||
async def aget_state_history(
|
||||
@@ -1584,7 +1578,7 @@ class Pregel(
|
||||
)
|
||||
]:
|
||||
yield await self._aprepare_state_snapshot(
|
||||
checkpoint_tuple.config, checkpoint_tuple
|
||||
checkpoint_tuple.config, checkpoint_tuple, saver=checkpointer
|
||||
)
|
||||
|
||||
def bulk_update_state(
|
||||
@@ -1666,10 +1660,7 @@ class Pregel(
|
||||
channels, managed = channels_from_checkpoint(
|
||||
self.channels,
|
||||
checkpoint,
|
||||
saver=self.checkpointer
|
||||
if saved is not None
|
||||
and isinstance(self.checkpointer, BaseCheckpointSaver)
|
||||
else None,
|
||||
saver=checkpointer if saved is not None else None,
|
||||
config=saved.config if saved is not None else None,
|
||||
)
|
||||
values, as_node = updates[0][:2]
|
||||
@@ -2132,10 +2123,7 @@ class Pregel(
|
||||
channels, managed = await achannels_from_checkpoint(
|
||||
self.channels,
|
||||
checkpoint,
|
||||
saver=self.checkpointer
|
||||
if saved is not None
|
||||
and isinstance(self.checkpointer, BaseCheckpointSaver)
|
||||
else None,
|
||||
saver=checkpointer if saved is not None else None,
|
||||
config=saved.config if saved is not None else None,
|
||||
)
|
||||
values, as_node = updates[0][:2]
|
||||
|
||||
@@ -0,0 +1,253 @@
|
||||
import operator
|
||||
from typing import Annotated, Any, Literal
|
||||
|
||||
import pytest
|
||||
from langgraph.checkpoint.base import BaseCheckpointSaver
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.graph import END, START, StateGraph
|
||||
|
||||
pytestmark = pytest.mark.anyio
|
||||
|
||||
|
||||
def _extend(state: list | None, writes: list[Any]) -> list:
|
||||
out = list(state or [])
|
||||
for write in writes:
|
||||
out.extend(write if isinstance(write, list) else [write])
|
||||
return out
|
||||
|
||||
|
||||
def _state_schema(snapshot_frequency: int = 1000) -> type:
|
||||
class State(TypedDict, total=False):
|
||||
delta: Annotated[
|
||||
list, DeltaChannel(_extend, snapshot_frequency=snapshot_frequency)
|
||||
]
|
||||
plain: Annotated[list, operator.add]
|
||||
|
||||
return State
|
||||
|
||||
|
||||
def _both(*items: str) -> dict:
|
||||
return {"delta": list(items), "plain": list(items)}
|
||||
|
||||
|
||||
def _child_builder(*, snapshot_frequency: int = 1000) -> StateGraph:
|
||||
builder = StateGraph(_state_schema(snapshot_frequency))
|
||||
builder.add_node("a", lambda state: _both("a1"))
|
||||
builder.add_node("b", lambda state: _both("b1", "b2"))
|
||||
builder.add_edge(START, "a")
|
||||
builder.add_edge("a", "b")
|
||||
builder.add_edge("b", END)
|
||||
return builder
|
||||
|
||||
|
||||
def _wrap(
|
||||
inner: StateGraph,
|
||||
*,
|
||||
checkpointer: bool | None = None,
|
||||
interrupt_before: list[str] | None = None,
|
||||
) -> StateGraph:
|
||||
builder = StateGraph(inner.state_schema)
|
||||
builder.add_node(
|
||||
"child",
|
||||
inner.compile(checkpointer=checkpointer, interrupt_before=interrupt_before),
|
||||
)
|
||||
builder.add_edge(START, "child")
|
||||
builder.add_edge("child", END)
|
||||
return builder
|
||||
|
||||
|
||||
def _nested_app(
|
||||
checkpointer: BaseCheckpointSaver,
|
||||
*,
|
||||
depth: int = 1,
|
||||
snapshot_frequency: int = 1000,
|
||||
pause_before_b: bool = False,
|
||||
subgraph_checkpointer: bool | None = None,
|
||||
) -> Any:
|
||||
graph = _child_builder(snapshot_frequency=snapshot_frequency)
|
||||
for _ in range(depth):
|
||||
graph = _wrap(
|
||||
graph,
|
||||
checkpointer=subgraph_checkpointer,
|
||||
interrupt_before=["b"] if pause_before_b else None,
|
||||
)
|
||||
return graph.compile(checkpointer=checkpointer)
|
||||
|
||||
|
||||
def _scoped(config: dict, namespace: str) -> dict:
|
||||
return {"configurable": {**config["configurable"], "checkpoint_ns": namespace}}
|
||||
|
||||
|
||||
def _child_namespace(app: Any, config: dict, *, depth: int = 1) -> str:
|
||||
namespace = ""
|
||||
for level in range(depth):
|
||||
scoped = _scoped(config, namespace) if namespace else config
|
||||
namespace = next(
|
||||
(
|
||||
task.state["configurable"]["checkpoint_ns"]
|
||||
for snapshot in app.get_state_history(scoped)
|
||||
for task in snapshot.tasks
|
||||
if task.name == "child" and isinstance(task.state, dict)
|
||||
),
|
||||
"",
|
||||
)
|
||||
assert namespace, f"no `child` subgraph task at nesting level {level}"
|
||||
return namespace
|
||||
|
||||
|
||||
async def _achild_namespace(app: Any, config: dict) -> str:
|
||||
async for snapshot in app.aget_state_history(config):
|
||||
for task in snapshot.tasks:
|
||||
if task.name == "child" and isinstance(task.state, dict):
|
||||
return task.state["configurable"]["checkpoint_ns"]
|
||||
raise AssertionError("no `child` subgraph task")
|
||||
|
||||
|
||||
HISTORY = [_both("a1", "b1", "b2"), _both("a1"), _both(), _both()]
|
||||
|
||||
|
||||
def test_subgraph_get_state(sync_checkpointer: BaseCheckpointSaver) -> None:
|
||||
app = _nested_app(sync_checkpointer)
|
||||
config = {"configurable": {"thread_id": "1"}}
|
||||
app.invoke({}, config)
|
||||
|
||||
child = _scoped(config, _child_namespace(app, config))
|
||||
|
||||
assert app.get_state(config).values == _both("a1", "b1", "b2")
|
||||
assert app.get_state(child).values == _both("a1", "b1", "b2")
|
||||
|
||||
|
||||
async def test_subgraph_aget_state(async_checkpointer: BaseCheckpointSaver) -> None:
|
||||
app = _nested_app(async_checkpointer)
|
||||
config = {"configurable": {"thread_id": "1"}}
|
||||
await app.ainvoke({}, config)
|
||||
|
||||
child = _scoped(config, await _achild_namespace(app, config))
|
||||
|
||||
assert (await app.aget_state(child)).values == _both("a1", "b1", "b2")
|
||||
|
||||
|
||||
def test_subgraph_get_state_history(sync_checkpointer: BaseCheckpointSaver) -> None:
|
||||
app = _nested_app(sync_checkpointer)
|
||||
config = {"configurable": {"thread_id": "1"}}
|
||||
app.invoke({}, config)
|
||||
|
||||
child = _scoped(config, _child_namespace(app, config))
|
||||
|
||||
assert [s.values for s in app.get_state_history(child)] == HISTORY
|
||||
|
||||
|
||||
async def test_subgraph_aget_state_history(
|
||||
async_checkpointer: BaseCheckpointSaver,
|
||||
) -> None:
|
||||
app = _nested_app(async_checkpointer)
|
||||
config = {"configurable": {"thread_id": "1"}}
|
||||
await app.ainvoke({}, config)
|
||||
|
||||
child = _scoped(config, await _achild_namespace(app, config))
|
||||
|
||||
assert [s.values async for s in app.aget_state_history(child)] == HISTORY
|
||||
|
||||
|
||||
def test_doubly_nested_subgraph_get_state(
|
||||
sync_checkpointer: BaseCheckpointSaver,
|
||||
) -> None:
|
||||
app = _nested_app(sync_checkpointer, depth=2)
|
||||
config = {"configurable": {"thread_id": "1"}}
|
||||
app.invoke({}, config)
|
||||
|
||||
child = _scoped(config, _child_namespace(app, config, depth=2))
|
||||
|
||||
assert app.get_state(child).values == _both("a1", "b1", "b2")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("persistence", ["per-invocation", "per-thread"])
|
||||
def test_interrupted_subgraph_task_state(
|
||||
sync_checkpointer: BaseCheckpointSaver,
|
||||
persistence: Literal["per-invocation", "per-thread"],
|
||||
) -> None:
|
||||
app = _nested_app(
|
||||
sync_checkpointer,
|
||||
pause_before_b=True,
|
||||
subgraph_checkpointer=True if persistence == "per-thread" else None,
|
||||
)
|
||||
config = {"configurable": {"thread_id": "1"}}
|
||||
app.invoke({}, config)
|
||||
|
||||
(task,) = app.get_state(config, subgraphs=True).tasks
|
||||
|
||||
assert task.state.values == _both("a1")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("persistence", ["per-invocation", "per-thread"])
|
||||
async def test_interrupted_subgraph_task_state_async(
|
||||
async_checkpointer: BaseCheckpointSaver,
|
||||
persistence: Literal["per-invocation", "per-thread"],
|
||||
) -> None:
|
||||
app = _nested_app(
|
||||
async_checkpointer,
|
||||
pause_before_b=True,
|
||||
subgraph_checkpointer=True if persistence == "per-thread" else None,
|
||||
)
|
||||
config = {"configurable": {"thread_id": "1"}}
|
||||
await app.ainvoke({}, config)
|
||||
|
||||
(task,) = (await app.aget_state(config, subgraphs=True)).tasks
|
||||
|
||||
assert task.state.values == _both("a1")
|
||||
|
||||
|
||||
def test_subgraph_update_state_keeps_history(
|
||||
sync_checkpointer: BaseCheckpointSaver,
|
||||
) -> None:
|
||||
app = _nested_app(sync_checkpointer, snapshot_frequency=2)
|
||||
config = {"configurable": {"thread_id": "1"}}
|
||||
app.invoke({}, config)
|
||||
|
||||
child = _scoped(config, _child_namespace(app, config))
|
||||
app.update_state(child, _both("manual"))
|
||||
|
||||
assert app.get_state(child).values == _both("a1", "b1", "b2", "manual")
|
||||
|
||||
|
||||
async def test_subgraph_aupdate_state_keeps_history(
|
||||
async_checkpointer: BaseCheckpointSaver,
|
||||
) -> None:
|
||||
app = _nested_app(async_checkpointer, snapshot_frequency=2)
|
||||
config = {"configurable": {"thread_id": "1"}}
|
||||
await app.ainvoke({}, config)
|
||||
|
||||
child = _scoped(config, await _achild_namespace(app, config))
|
||||
await app.aupdate_state(child, _both("manual"))
|
||||
|
||||
assert (await app.aget_state(child)).values == _both("a1", "b1", "b2", "manual")
|
||||
|
||||
|
||||
def test_stateless_subgraph_persists_nothing(
|
||||
sync_checkpointer: BaseCheckpointSaver,
|
||||
) -> None:
|
||||
app = _nested_app(sync_checkpointer, subgraph_checkpointer=False)
|
||||
config = {"configurable": {"thread_id": "1"}}
|
||||
app.invoke({}, config)
|
||||
|
||||
child_tasks = [
|
||||
task
|
||||
for snapshot in app.get_state_history(config)
|
||||
for task in snapshot.tasks
|
||||
if task.name == "child" and isinstance(task.state, dict)
|
||||
]
|
||||
|
||||
assert child_tasks == []
|
||||
assert app.get_state(config).values == _both("a1", "b1", "b2")
|
||||
|
||||
|
||||
def test_completed_subgraph_exposes_no_task_state(
|
||||
sync_checkpointer: BaseCheckpointSaver,
|
||||
) -> None:
|
||||
app = _nested_app(sync_checkpointer)
|
||||
config = {"configurable": {"thread_id": "1"}}
|
||||
app.invoke({}, config)
|
||||
|
||||
assert app.get_state(config, subgraphs=True).tasks == ()
|
||||
Reference in New Issue
Block a user