mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-08-28 10:49:56 +02:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
a72cb56fed | ||
|
|
99f1e4e962 | ||
|
|
7912a1ce50 |
@@ -435,26 +435,15 @@ class BasePostgresSaver(BaseCheckpointSaver[str]):
|
||||
(c) the next ancestor cid isn't in `parent_of` yet (waiting for
|
||||
a later page; the cursor stays put).
|
||||
|
||||
A walk that hasn't started yet is a fourth state, distinct from all
|
||||
three: pages arrive newest-first from the head of the thread, so a
|
||||
target deeper than one page isn't in `parent_of` on the early pages
|
||||
and the walk can only begin once its own row lands.
|
||||
|
||||
Mutates `chain_by_ch`, `seed_ver_by_ch`, `seed_inline_by_ch`,
|
||||
`walk_cursor_by_ch`, and `seeded` in place.
|
||||
"""
|
||||
for i, ch in enumerate(channels):
|
||||
if ch in seeded:
|
||||
continue
|
||||
# First-time entry: cursor starts at the target's parent, but
|
||||
# only once the target's own row has loaded. Reading it earlier
|
||||
# would record `None` for "target not seen yet" the same way it
|
||||
# records `None` for "target is a root", and this guard fires
|
||||
# only once, so the walk would stay stranded for good.
|
||||
# First-time entry: cursor starts at the target's parent.
|
||||
if ch not in walk_cursor_by_ch:
|
||||
if target_id not in parent_of:
|
||||
continue
|
||||
walk_cursor_by_ch[ch] = parent_of[target_id]
|
||||
walk_cursor_by_ch[ch] = parent_of.get(target_id)
|
||||
cur_cid = walk_cursor_by_ch[ch]
|
||||
ch_chain = chain_by_ch[ch]
|
||||
hb_i = hb_by_i_by_cid[i]
|
||||
|
||||
@@ -1,172 +0,0 @@
|
||||
"""Stage-1 pagination for `DeltaChannel` histories on Postgres.
|
||||
|
||||
`get_delta_channel_history` pages the `checkpoints` table newest-first in
|
||||
chunks of `_DELTA_PAGE_SIZE`, starting from the head of the thread rather than
|
||||
from the target checkpoint. The target's own row therefore only lands once
|
||||
paging has reached back to it, which for a long thread can be several pages in.
|
||||
|
||||
Until then `parent_of` has no entry for it, and reading the walk cursor out of
|
||||
that map records `None`, the same value that means "the target is a root". The
|
||||
cursor is derived once, so a target older than the first page kept that `None`
|
||||
forever: no seed, no writes, and the channel hydrated empty with no error.
|
||||
See #8448.
|
||||
|
||||
These tests shrink the page size instead of writing 1024+ real checkpoints per
|
||||
case. The behaviour under test is "the target is not on the first page", and
|
||||
the page the target lands on is the only thing that decides it.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
from uuid import uuid4
|
||||
|
||||
import pytest
|
||||
from langgraph.checkpoint.base import (
|
||||
Checkpoint,
|
||||
DeltaChannelHistory,
|
||||
empty_checkpoint,
|
||||
)
|
||||
from langgraph.checkpoint.base.id import uuid6
|
||||
from langgraph.checkpoint.serde.types import _DeltaSnapshot
|
||||
|
||||
from langgraph.checkpoint.postgres import PostgresSaver
|
||||
from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
|
||||
from langgraph.checkpoint.postgres.base import _DELTA_PAGE_SIZE
|
||||
from tests.conftest import DEFAULT_URI
|
||||
|
||||
CHANNEL = "items"
|
||||
STEPS = 8
|
||||
SEED_STEP = 1
|
||||
SEED_VALUE = [10, 20]
|
||||
# Deep enough that the smaller page sizes below have to page past it.
|
||||
TARGET_STEP = 4
|
||||
|
||||
# The real page size holds this whole thread on one page, so it is the control.
|
||||
# The rest each leave the target off the first page: there are three
|
||||
# checkpoints newer than `TARGET_STEP`, so any page size at or below 3 does it.
|
||||
PAGE_SIZES = [_DELTA_PAGE_SIZE, 3, 2, 1]
|
||||
|
||||
|
||||
def _step_args(
|
||||
thread_id: str, step: int, parent: dict | None
|
||||
) -> tuple[dict, Checkpoint, dict[str, Any]]:
|
||||
"""Return the `(config, checkpoint, new_versions)` triple for one step.
|
||||
|
||||
Step `SEED_STEP` stores a snapshot for `CHANNEL`; the rest only bump the
|
||||
channel version, which is what a delta channel does between snapshots.
|
||||
"""
|
||||
config: dict = {"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}}
|
||||
if parent is not None:
|
||||
config["configurable"]["checkpoint_id"] = parent["configurable"][
|
||||
"checkpoint_id"
|
||||
]
|
||||
checkpoint: Checkpoint = empty_checkpoint()
|
||||
checkpoint["id"] = str(uuid6(clock_seq=step))
|
||||
checkpoint["channel_versions"][CHANNEL] = f"v{step}"
|
||||
if step == SEED_STEP:
|
||||
checkpoint["channel_values"][CHANNEL] = _DeltaSnapshot(list(SEED_VALUE))
|
||||
return config, checkpoint, {CHANNEL: f"v{step}"}
|
||||
return config, checkpoint, {}
|
||||
|
||||
|
||||
async def _abuild_chain(saver: AsyncPostgresSaver) -> list[dict]:
|
||||
"""Write `STEPS` linked checkpoints, one write each; return their configs."""
|
||||
thread_id = str(uuid4())
|
||||
parent: dict | None = None
|
||||
configs: list[dict] = []
|
||||
for step in range(STEPS):
|
||||
config, checkpoint, new_versions = _step_args(thread_id, step, parent)
|
||||
parent = await saver.aput(
|
||||
config,
|
||||
checkpoint,
|
||||
{"source": "loop", "step": step, "parents": {}},
|
||||
new_versions,
|
||||
)
|
||||
await saver.aput_writes(parent, [(CHANNEL, f"w{step}")], str(uuid4()))
|
||||
configs.append(parent)
|
||||
return configs
|
||||
|
||||
|
||||
def _build_chain(saver: PostgresSaver) -> list[dict]:
|
||||
"""Sync twin of `_abuild_chain`."""
|
||||
thread_id = str(uuid4())
|
||||
parent: dict | None = None
|
||||
configs: list[dict] = []
|
||||
for step in range(STEPS):
|
||||
config, checkpoint, new_versions = _step_args(thread_id, step, parent)
|
||||
parent = saver.put(
|
||||
config,
|
||||
checkpoint,
|
||||
{"source": "loop", "step": step, "parents": {}},
|
||||
new_versions,
|
||||
)
|
||||
saver.put_writes(parent, [(CHANNEL, f"w{step}")], str(uuid4()))
|
||||
configs.append(parent)
|
||||
return configs
|
||||
|
||||
|
||||
def _assert_history(entry: DeltaChannelHistory, page_size: int) -> None:
|
||||
"""Check the walk from `TARGET_STEP` back to the snapshot at `SEED_STEP`.
|
||||
|
||||
The chain is steps 3, 2 and 1: the target's own writes are pending for its
|
||||
next super-step and excluded, and the walk stops at the snapshot. Writes
|
||||
come back oldest first.
|
||||
"""
|
||||
seed = entry.get("seed")
|
||||
assert isinstance(seed, _DeltaSnapshot), (
|
||||
f"page_size={page_size}: expected a snapshot seed, "
|
||||
f"got {entry.get('seed', '<missing>')!r}"
|
||||
)
|
||||
assert seed.value == SEED_VALUE
|
||||
assert [w[2] for w in entry["writes"]] == ["w1", "w2", "w3"], (
|
||||
f"page_size={page_size}: got {[w[2] for w in entry['writes']]}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("page_size", PAGE_SIZES)
|
||||
async def test_async_target_older_than_the_first_page(
|
||||
page_size: int, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
monkeypatch.setattr("langgraph.checkpoint.postgres.aio._DELTA_PAGE_SIZE", page_size)
|
||||
async with AsyncPostgresSaver.from_conn_string(DEFAULT_URI) as saver:
|
||||
await saver.setup()
|
||||
configs = await _abuild_chain(saver)
|
||||
result = await saver.aget_delta_channel_history(
|
||||
config=configs[TARGET_STEP], channels=[CHANNEL]
|
||||
)
|
||||
_assert_history(result[CHANNEL], page_size)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("page_size", PAGE_SIZES)
|
||||
def test_sync_target_older_than_the_first_page(
|
||||
page_size: int, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
monkeypatch.setattr("langgraph.checkpoint.postgres._DELTA_PAGE_SIZE", page_size)
|
||||
with PostgresSaver.from_conn_string(DEFAULT_URI) as saver:
|
||||
saver.setup()
|
||||
configs = _build_chain(saver)
|
||||
result = saver.get_delta_channel_history(
|
||||
config=configs[TARGET_STEP], channels=[CHANNEL]
|
||||
)
|
||||
_assert_history(result[CHANNEL], page_size)
|
||||
|
||||
|
||||
async def test_root_target_has_no_history_and_still_terminates(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""A root target is the case where a `None` cursor is the right answer.
|
||||
|
||||
Its walk has nowhere to go, so it never seeds and collects none of the
|
||||
thread's writes, and paging has to run out on a short page rather than
|
||||
spin. Page size 1 gives one page per checkpoint, so the loop runs the
|
||||
length of the thread before stopping.
|
||||
"""
|
||||
monkeypatch.setattr("langgraph.checkpoint.postgres.aio._DELTA_PAGE_SIZE", 1)
|
||||
async with AsyncPostgresSaver.from_conn_string(DEFAULT_URI) as saver:
|
||||
await saver.setup()
|
||||
configs = await _abuild_chain(saver)
|
||||
result = await saver.aget_delta_channel_history(
|
||||
config=configs[0], channels=[CHANNEL]
|
||||
)
|
||||
assert result[CHANNEL] == {"writes": []}
|
||||
@@ -1146,9 +1146,26 @@ class Pregel(
|
||||
self,
|
||||
config: RunnableConfig,
|
||||
saved: CheckpointTuple | None,
|
||||
*,
|
||||
saver: BaseCheckpointSaver | None,
|
||||
recurse: BaseCheckpointSaver | None = None,
|
||||
apply_pending_writes: bool = False,
|
||||
) -> StateSnapshot:
|
||||
"""Assemble a `StateSnapshot` from a saved checkpoint.
|
||||
|
||||
Args:
|
||||
config: Config identifying the checkpoint being read.
|
||||
saved: The checkpoint tuple to render, or `None` for an empty snapshot.
|
||||
saver: Checkpointer to read with, as resolved by the caller from
|
||||
`CONFIG_KEY_CHECKPOINTER` before falling back to `self.checkpointer`.
|
||||
Required rather than defaulted because `self.checkpointer` is `None`
|
||||
for a subgraph, which borrows the parent's saver through the config,
|
||||
and a `DeltaChannel` silently hydrates empty without one.
|
||||
recurse: When set, resolve subgraph task states with this checkpointer
|
||||
instead of returning a config that merely signals they exist.
|
||||
apply_pending_writes: Apply this checkpoint's pending writes to the
|
||||
returned values.
|
||||
"""
|
||||
if not saved:
|
||||
return StateSnapshot(
|
||||
values={},
|
||||
@@ -1169,9 +1186,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
|
||||
@@ -1269,9 +1284,12 @@ class Pregel(
|
||||
self,
|
||||
config: RunnableConfig,
|
||||
saved: CheckpointTuple | None,
|
||||
*,
|
||||
saver: BaseCheckpointSaver | None,
|
||||
recurse: BaseCheckpointSaver | None = None,
|
||||
apply_pending_writes: bool = False,
|
||||
) -> StateSnapshot:
|
||||
"""Async version of `_prepare_state_snapshot`. See docstring there."""
|
||||
if not saved:
|
||||
return StateSnapshot(
|
||||
values={},
|
||||
@@ -1292,9 +1310,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
|
||||
@@ -1429,6 +1445,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 +1490,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 +1545,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 +1602,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 +1684,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 +2147,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,328 @@
|
||||
"""Tests for reading and updating a subgraph's `DeltaChannel` state.
|
||||
|
||||
Regression suite for #8470.
|
||||
|
||||
A subgraph is compiled without a checkpointer of its own. The parent lends it
|
||||
one through `CONFIG_KEY_CHECKPOINTER` at read time. `DeltaChannel` stores no
|
||||
value in `channel_values`, so hydrating it requires that saver to walk ancestors
|
||||
and replay their writes. Reading with `self.checkpointer` instead of the
|
||||
caller-resolved saver leaves a subgraph with none, and the channel falls through
|
||||
to `from_checkpoint(MISSING)`: an empty value, indistinguishable from a channel
|
||||
that was never written and raising nothing.
|
||||
|
||||
Coverage:
|
||||
|
||||
* `get_state` / `aget_state` and `get_state_history` / `aget_state_history` on a
|
||||
subgraph namespace, at one and two levels of nesting
|
||||
* `get_state(subgraphs=True)` task states, the surface a human-in-the-loop client
|
||||
reads while a subgraph sits interrupted
|
||||
* `update_state` / `aupdate_state`, which persist the hydrated value as a
|
||||
snapshot blob and so turn an empty hydration into permanent data loss
|
||||
* root-graph controls, which own their checkpointer and were never affected
|
||||
"""
|
||||
|
||||
from typing import Annotated, Any
|
||||
|
||||
import pytest
|
||||
from langgraph.checkpoint.memory import InMemorySaver
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.graph import END, START, StateGraph
|
||||
from langgraph.types import interrupt
|
||||
|
||||
pytestmark = pytest.mark.anyio
|
||||
|
||||
|
||||
def _extend(state: list | None, writes: list[Any]) -> list:
|
||||
"""Batching-invariant list reducer: flattens each write batch onto the state."""
|
||||
out = list(state or [])
|
||||
for write in writes:
|
||||
out.extend(write if isinstance(write, list) else [write])
|
||||
return out
|
||||
|
||||
|
||||
def _child_builder(*, snapshot_frequency: int) -> StateGraph:
|
||||
"""Two-node graph writing `["a1"]` then `["b1", "b2"]` to a `DeltaChannel`."""
|
||||
|
||||
class State(TypedDict, total=False):
|
||||
msgs: Annotated[
|
||||
list, DeltaChannel(_extend, snapshot_frequency=snapshot_frequency)
|
||||
]
|
||||
|
||||
builder = StateGraph(State)
|
||||
builder.add_node("a", lambda state: {"msgs": ["a1"]})
|
||||
builder.add_node("b", lambda state: {"msgs": ["b1", "b2"]})
|
||||
builder.add_edge(START, "a")
|
||||
builder.add_edge("a", "b")
|
||||
builder.add_edge("b", END)
|
||||
return builder
|
||||
|
||||
|
||||
def _wrap(inner: StateGraph, *, node_name: str) -> StateGraph:
|
||||
"""Nest `inner` as a single node, sharing its state schema."""
|
||||
builder = StateGraph(inner.state_schema)
|
||||
builder.add_node(node_name, inner.compile())
|
||||
builder.add_edge(START, node_name)
|
||||
builder.add_edge(node_name, END)
|
||||
return builder
|
||||
|
||||
|
||||
def _build_nested_graph(
|
||||
checkpointer: InMemorySaver,
|
||||
*,
|
||||
snapshot_frequency: int = 1000,
|
||||
depth: int = 1,
|
||||
) -> Any:
|
||||
"""Compile a graph whose `child` node is a subgraph nested `depth` levels deep."""
|
||||
graph = _child_builder(snapshot_frequency=snapshot_frequency)
|
||||
for _ in range(depth):
|
||||
graph = _wrap(graph, node_name="child")
|
||||
return graph.compile(checkpointer=checkpointer)
|
||||
|
||||
|
||||
def _namespaced(config: dict, namespace: str) -> dict:
|
||||
return {"configurable": {**config["configurable"], "checkpoint_ns": namespace}}
|
||||
|
||||
|
||||
def _subgraph_namespace(app: Any, config: dict, *, depth: int = 1) -> str:
|
||||
"""Resolve the `child` namespace `depth` levels down, as a client would.
|
||||
|
||||
Each snapshot's tasks carry `{name, state: {"configurable": {"checkpoint_ns"}}}`,
|
||||
and passing that namespace back is the documented way to read a subgraph's own
|
||||
supersteps.
|
||||
"""
|
||||
namespace = ""
|
||||
for level in range(depth):
|
||||
scoped = _namespaced(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
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Reading a subgraph's state
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_subgraph_get_state_replays_delta_channel() -> None:
|
||||
app = _build_nested_graph(InMemorySaver())
|
||||
config = {"configurable": {"thread_id": "1"}}
|
||||
app.invoke({}, config)
|
||||
|
||||
namespace = _subgraph_namespace(app, config)
|
||||
|
||||
assert app.get_state(config).values == {"msgs": ["a1", "b1", "b2"]}
|
||||
assert app.get_state(_namespaced(config, namespace)).values == {
|
||||
"msgs": ["a1", "b1", "b2"]
|
||||
}
|
||||
|
||||
|
||||
async def test_subgraph_aget_state_replays_delta_channel() -> None:
|
||||
app = _build_nested_graph(InMemorySaver())
|
||||
config = {"configurable": {"thread_id": "1"}}
|
||||
await app.ainvoke({}, config)
|
||||
|
||||
namespace = _subgraph_namespace(app, config)
|
||||
|
||||
assert (await app.aget_state(_namespaced(config, namespace))).values == {
|
||||
"msgs": ["a1", "b1", "b2"]
|
||||
}
|
||||
|
||||
|
||||
def test_subgraph_get_state_history_replays_delta_channel() -> None:
|
||||
app = _build_nested_graph(InMemorySaver())
|
||||
config = {"configurable": {"thread_id": "1"}}
|
||||
app.invoke({}, config)
|
||||
|
||||
namespace = _subgraph_namespace(app, config)
|
||||
values = [
|
||||
snapshot.values
|
||||
for snapshot in app.get_state_history(_namespaced(config, namespace))
|
||||
]
|
||||
|
||||
# newest first: after `b`, after `a`, then the two supersteps preceding any write
|
||||
assert values == [
|
||||
{"msgs": ["a1", "b1", "b2"]},
|
||||
{"msgs": ["a1"]},
|
||||
{"msgs": []},
|
||||
{"msgs": []},
|
||||
]
|
||||
|
||||
|
||||
async def test_subgraph_aget_state_history_replays_delta_channel() -> None:
|
||||
app = _build_nested_graph(InMemorySaver())
|
||||
config = {"configurable": {"thread_id": "1"}}
|
||||
await app.ainvoke({}, config)
|
||||
|
||||
namespace = _subgraph_namespace(app, config)
|
||||
values = [
|
||||
snapshot.values
|
||||
async for snapshot in app.aget_state_history(_namespaced(config, namespace))
|
||||
]
|
||||
|
||||
assert values == [
|
||||
{"msgs": ["a1", "b1", "b2"]},
|
||||
{"msgs": ["a1"]},
|
||||
{"msgs": []},
|
||||
{"msgs": []},
|
||||
]
|
||||
|
||||
|
||||
def test_doubly_nested_subgraph_get_state_replays_delta_channel() -> None:
|
||||
app = _build_nested_graph(InMemorySaver(), depth=2)
|
||||
config = {"configurable": {"thread_id": "1"}}
|
||||
app.invoke({}, config)
|
||||
|
||||
namespace = _subgraph_namespace(app, config, depth=2)
|
||||
|
||||
assert app.get_state(_namespaced(config, namespace)).values == {
|
||||
"msgs": ["a1", "b1", "b2"]
|
||||
}
|
||||
|
||||
|
||||
def test_interrupted_subgraph_task_state_replays_delta_channel() -> None:
|
||||
"""`get_state(subgraphs=True)` resolves task states through a separate
|
||||
recursion; it reads the paused subgraph a human-in-the-loop client is shown."""
|
||||
|
||||
class State(TypedDict, total=False):
|
||||
msgs: Annotated[list, DeltaChannel(_extend)]
|
||||
|
||||
def pause(state: State) -> dict:
|
||||
interrupt("pause")
|
||||
return {"msgs": ["b1"]}
|
||||
|
||||
child = StateGraph(State)
|
||||
child.add_node("a", lambda state: {"msgs": ["a1", "a2"]})
|
||||
child.add_node("b", pause)
|
||||
child.add_edge(START, "a")
|
||||
child.add_edge("a", "b")
|
||||
child.add_edge("b", END)
|
||||
|
||||
app = _wrap(child, node_name="child").compile(checkpointer=InMemorySaver())
|
||||
config = {"configurable": {"thread_id": "1"}}
|
||||
app.invoke({}, config)
|
||||
|
||||
task_values = [
|
||||
task.state.values
|
||||
for task in app.get_state(config, subgraphs=True).tasks
|
||||
if task.name == "child"
|
||||
]
|
||||
|
||||
assert task_values == [{"msgs": ["a1", "a2"]}]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Updating a subgraph's state
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_subgraph_update_state_preserves_delta_channel_history() -> None:
|
||||
"""`update_state` persists the hydrated channel, so an empty hydration is
|
||||
written to disk rather than merely returned.
|
||||
|
||||
`snapshot_frequency=2` is what makes that observable: this update reaches the
|
||||
cadence and forces a `_DeltaSnapshot` blob built from the hydrated value. At
|
||||
the default frequency nothing is written and the loss stays latent, so the
|
||||
assertion below would hold either way.
|
||||
"""
|
||||
app = _build_nested_graph(InMemorySaver(), snapshot_frequency=2)
|
||||
config = {"configurable": {"thread_id": "1"}}
|
||||
app.invoke({}, config)
|
||||
|
||||
namespace = _subgraph_namespace(app, config)
|
||||
child_config = _namespaced(config, namespace)
|
||||
app.update_state(child_config, {"msgs": ["manual"]})
|
||||
|
||||
assert app.get_state(child_config).values == {"msgs": ["a1", "b1", "b2", "manual"]}
|
||||
|
||||
|
||||
async def test_subgraph_aupdate_state_preserves_delta_channel_history() -> None:
|
||||
app = _build_nested_graph(InMemorySaver(), snapshot_frequency=2)
|
||||
config = {"configurable": {"thread_id": "1"}}
|
||||
await app.ainvoke({}, config)
|
||||
|
||||
namespace = _subgraph_namespace(app, config)
|
||||
child_config = _namespaced(config, namespace)
|
||||
await app.aupdate_state(child_config, {"msgs": ["manual"]})
|
||||
|
||||
assert (await app.aget_state(child_config)).values == {
|
||||
"msgs": ["a1", "b1", "b2", "manual"]
|
||||
}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Controls: a graph owning its checkpointer resolves the same saver as before
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_root_graph_get_state_replays_delta_channel() -> None:
|
||||
app = _child_builder(snapshot_frequency=1000).compile(checkpointer=InMemorySaver())
|
||||
config = {"configurable": {"thread_id": "1"}}
|
||||
app.invoke({}, config)
|
||||
|
||||
assert app.get_state(config).values == {"msgs": ["a1", "b1", "b2"]}
|
||||
assert [snapshot.values for snapshot in app.get_state_history(config)] == [
|
||||
{"msgs": ["a1", "b1", "b2"]},
|
||||
{"msgs": ["a1"]},
|
||||
{"msgs": []},
|
||||
{"msgs": []},
|
||||
]
|
||||
|
||||
|
||||
def test_root_graph_update_state_preserves_delta_channel_history() -> None:
|
||||
app = _child_builder(snapshot_frequency=2).compile(checkpointer=InMemorySaver())
|
||||
config = {"configurable": {"thread_id": "1"}}
|
||||
app.invoke({}, config)
|
||||
|
||||
app.update_state(config, {"msgs": ["manual"]})
|
||||
|
||||
assert app.get_state(config).values == {"msgs": ["a1", "b1", "b2", "manual"]}
|
||||
|
||||
|
||||
def test_stateless_subgraph_persists_nothing() -> None:
|
||||
"""A `checkpointer=False` subgraph opts out of persistence entirely.
|
||||
|
||||
Its state is unavailable by design, and passing the caller's saver must not
|
||||
start surfacing state for a subgraph that asked not to be checkpointed.
|
||||
"""
|
||||
child = _child_builder(snapshot_frequency=1000).compile(checkpointer=False)
|
||||
builder = StateGraph(child.builder.state_schema)
|
||||
builder.add_node("child", child)
|
||||
builder.add_edge(START, "child")
|
||||
builder.add_edge("child", END)
|
||||
app = builder.compile(checkpointer=InMemorySaver())
|
||||
config = {"configurable": {"thread_id": "1"}}
|
||||
app.invoke({}, config)
|
||||
|
||||
subgraph_namespaces = [
|
||||
task.state["configurable"]["checkpoint_ns"]
|
||||
for snapshot in app.get_state_history(config)
|
||||
for task in snapshot.tasks
|
||||
if task.name == "child" and isinstance(task.state, dict)
|
||||
]
|
||||
|
||||
assert subgraph_namespaces == []
|
||||
assert app.get_state(config).values == {"msgs": ["a1", "b1", "b2"]}
|
||||
|
||||
|
||||
def test_completed_subgraph_exposes_no_task_state() -> None:
|
||||
"""`subgraphs=True` surfaces task state only while a task is pending.
|
||||
|
||||
Once the subgraph has finished there is no task to attach state to, for every
|
||||
channel type alike. This is subgraph behaviour rather than delta replay, and
|
||||
the fix leaves it untouched.
|
||||
"""
|
||||
app = _build_nested_graph(InMemorySaver())
|
||||
config = {"configurable": {"thread_id": "1"}}
|
||||
app.invoke({}, config)
|
||||
|
||||
assert app.get_state(config, subgraphs=True).tasks == ()
|
||||
Reference in New Issue
Block a user