Compare commits

...
Author SHA1 Message Date
Elior Nataf LackritzandGitHub afbe7e48b1 release(checkpoint-postgres): 3.1.3 (#9287)
Version bump only. No library code changes in this PR.

Unreleased since `checkpointpostgres==3.1.2`:

- #8617 - widen `Store` `put` value type to `Mapping[str, Any]`
- #8556 - derive the delta walk cursor once the target loads
- #8544 - order delta channel replay by task path
- #9217 - validate namespace labels before store dispatch
- #9248 - start the delta walk at the target
- #9262 - treat an empty `checkpoint_id` as the latest in delta history

Follows `checkpoint==4.3.0`, which is already live on PyPI. #8544 raised
the floor to `langgraph-checkpoint>=4.3.0`, so this release needs it.
2026-10-12 22:57:04 +00:00
Elior Nataf LackritzandGitHub 1074714c26 release(checkpoint-sqlite): 3.1.2 (#9288)
Version bump only. No library code changes in this PR.

Unreleased since `checkpointsqlite==3.1.1`:

- #8617 - widen `Store` `put` value type to `Mapping[str, Any]`
- #8557 - walk delta ancestors by parent pointer
- #8544 - order delta channel replay by task path
- #9217 - validate namespace labels before store dispatch
- #9262 - treat an empty `checkpoint_id` as the latest in delta history

Follows `checkpoint==4.3.0`, which is already live on PyPI. #8544 raised
the floor to `langgraph-checkpoint>=4.3.0`, so this release needs it.
2026-10-12 22:47:17 +00:00
Elior Nataf LackritzandGitHub 740461ed50 release(sdk-py): 0.4.7 (#9289)
Bumps the Python SDK version from 0.4.6 to 0.4.7. A thread stream's run
now ends only on the root graph's `completed` or `failed` event, not a
subgraph's, so `thread.output` no longer reads the state too early and a
run that fails after a subgraph finishes is reported as failed (#9228).
2026-10-12 22:47:05 +00:00
Igor SoarezandGitHub 08a345637f release(langgraph): 1.2.15 (#9283)
Bumps langgraph from 1.2.14 to 1.2.15. It now requires
langgraph-checkpoint 4.3.0 or later (#8544).

Also includes:

- Order delta channel replay by task path (#8544)
- Replay a resumed exit-mode run's delta writes in live order (#9114)
- Give each `bulk_update_state` update its own task id (#9128)
- Stop `astream_events` dropping control/interrupts on v1/v2 (#9219)
- Never store exit-mode delta writes under the null task id (#9229)
- Mark a resume value that fails an interrupt's `response_schema`, so
`ToolNode` can stop swallowing it once langgraph-prebuilt is released
(#9232)
- Don't save a checkpoint past a failed DeltaChannel write (#9235)
- Keep a fork's DeltaChannel snapshot from starting nodes that never ran
(#9264)
- Snapshot a DeltaChannel that an `update_state` Overwrite resets
(#9263)
- Read a run's DeltaChannel input back only on its own checkpoints
(#9260)
- Store an `update_state` input's writes on the checkpoint it updates
(#9258)
- Keep an exit run from a checkpoint the thread moved past out of its
newer branch (#9269)
- Stop at an `interrupt_before` breakpoint when a drain is requested
before it (#9281)
- Save a cached node's writes like a node that ran (#9259)
- Give every message written to a messages DeltaChannel an id (#9275)
2026-10-12 18:34:01 -04:00
b7d053b256 fix(langgraph): give every message written to a messages DeltaChannel an id (#9275)
A DeltaChannel replays its saved writes on every read, and
`_messages_delta_reducer` doesn't fill in missing ids, so the loop gives
messages their ids before saving them (`ensure_message_ids`). That
missed several cases:

- Strings and `(role, content)` tuples, alone or in a list (`[("user",
"hi")]`), and messages inside an `Overwrite` were saved with no id, so
nothing could target them by id, for example `RemoveMessage`. A regular
messages channel gives them ids.
- A raw `Pregel` graph's input to a DeltaChannel got its ids after it
was applied, and none with exit durability or on a new thread; a dict
given as input read back with an id its nodes never saw. A
`Command(update=...)` sent to an older checkpoint got none either.

Now, in a DeltaChannel whose reducer is `_messages_delta_reducer`, every
message-like the reducer accepts becomes a message with an id before the
writes are saved or applied: in `put_writes`, before input is applied,
for a forked `Command`, and in `update_state`. Anything that can't
become a message is left as it is, so it fails where it always has.
Other DeltaChannels keep the rule they had (message objects, and message
dicts in a list, get ids), since they may hold other data and
deepagents' messages channel, which uses its own copy of the reducer,
relies on it; the dicts in an `Overwrite` there are left as they are.
Lists and dicts the caller passed in are never changed: a write that
changes is replaced. So the id step can recognize the reducer without an
import cycle, `_messages_delta_reducer` now lives in
`langgraph.pregel._messages`; `langgraph.graph.message` re-exports it.

Visible changes: in a messages DeltaChannel, a node's update streams as
messages with ids, and a string is stored as a message (about 170 bytes
more each).

JS twin: langchain-ai/langgraphjs#2992. Thanks to @HarshShroff for
testing every `Overwrite` form and for the tuple fix.

## Tests

`test_delta_channel_id_stability.py` and
`test_delta_channel_update_state.py`: a write of every shape the reducer
accepts gets an id that every read keeps, as input, from a node and
through `update_state` (sync, async and as input); a list written twice
adds its messages twice and stays as it was; a forked `Command` gets an
id; a message dict the reducer rejects in a list still fails the run
before it's saved; a raw `Pregel` input gets an id on a new thread with
every durability, and its nodes see the ids that are saved; another
DeltaChannel keeps the dicts in an `Overwrite`. 36 of the 87 tests in
these files fail on main. The langgraph and prebuilt suites pass, and
lint is clean.

---------

Co-authored-by: Harsh Shroff <29751900+HarshShroff@users.noreply.github.com>
2026-10-12 20:33:45 +00:00
Elior Nataf LackritzandGitHub 395c145c04 fix(langgraph): save a cached node's writes like a node that ran (#9259)
A node served from the node cache had its writes applied but never
saved: the loop only streamed them. Plain channels didn't notice, since
the next checkpoint stores their value, but a `DeltaChannel` is rebuilt
from the saved writes, so its saved state lost the cached node's write
for good, under every durability. With `a -> b -> c` and `b` cached, a
second thread read back `log` as `a, c` next to `plain` as `a, b, c`.

The loop now saves a cache hit with `put_writes`, like a node that ran,
in all six places it matches cached writes. A new `cached` flag keeps
the hit streamed as cached and stops `put_writes` from writing it back
to the cache, which would also restart its TTL. A hit now costs the
saver write the node makes when it runs, and a cached `@task` call's
result is saved the same way.

The sync loop also stops caching a node that failed or was interrupted,
as the async loop already did: a cached error made the next run with the
same input skip the node without raising.

Found in review of langchain-ai/langgraphjs#2969. JS twin:
langchain-ai/langgraphjs#2980

## Tests

- `test_delta_channel_node_cache.py`: `b` runs on the first thread and
is a cache hit on the second, under every durability, with `invoke` and
`ainvoke`. Every checkpoint of the second thread reads the same in the
DeltaChannel and a plain channel, and the cache is written once. All 6
cases fail on main; without the `cached` check in `put_writes`, they
fail on the cache count.
- `test_a_node_that_fails_is_not_cached`, with `invoke` and `ainvoke`: a
cached node that raises raises again on the next thread. The sync case
fails on main.
- The two tests that spy on `PregelLoop.put_writes` pass the new keyword
through.
- `make format`, `make lint_package`, `make lint_tests`, and the
`langgraph`, `prebuilt` and `cli` suites pass (without Redis).
2026-10-12 14:24:38 -04:00
Igor SoarezandGitHub 10f1fabb3b fix(langgraph): stop at an interrupt_before breakpoint when a drain is requested before it (#9281)
When a drain is requested during the step right before an
interrupt_before breakpoint, the run stops with GraphDrained and the
breakpoint is lost: resuming the run skips it unintentionally.

The fix moves the drain check in tick() after the breakpoint check, so
breakpoints win over draining.

JS twin: langchain-ai/langgraphjs#3009
2026-10-12 18:29:56 +03:00
Elior Nataf LackritzandGitHub 799c2f04fd fix(langgraph): keep an exit run from a checkpoint the thread moved past out of its newer branch (#9269)
An exit-durability run that starts from a checkpoint the thread has
moved past, addressed by `checkpoint_id`, stores its DeltaChannel writes
on that checkpoint. The branch that already grew from it reads the
writes stored there, so its checkpoints picked up this run's writes too:

```python
graph.invoke({"log": ["in1"]}, config)  # nodes a and b add "a1", "b1"
turn_1 = graph.get_state(config).config
graph.invoke({"log": ["in2"]}, config)
turn_2 = graph.get_state(config).config
graph.invoke({"log": ["in3"]}, turn_1, durability="exit")
graph.get_state(turn_2).values["log"]
# ["in1", "a1", "b1", "in3", "a3", "b3", "in2", "a2", "b2"], was ["in1", "a1", "b1", "in2", "a2", "b2"]
```

Sync and async durability store each step's writes on the run's own
checkpoints, so they weren't affected. The run now snapshots the
DeltaChannels it wrote instead, as a run's input already does on a
checkpoint a `checkpoint_id` addressed (#9260). That costs one snapshot
of each channel the run wrote, even when the addressed checkpoint is the
latest one. Threads that already hit this keep the wrong values on the
existing branch.

Found by @hntrl in review of langchain-ai/langgraphjs#2974
([repro](https://gist.github.com/hntrl/2185e489abe22147b3f6f77010c18264)).
JS twin: langchain-ai/langgraphjs#2989

## Tests

`test_delta_channel_exit_mode.py`: one and two exit runs from the end of
turn 1, after turn 2, leave every checkpoint's DeltaChannel equal to a
plain reducer channel. Both fail on main on every saver in the fixture.
The `langgraph` suite passes (without Redis), and lint is clean.
2026-10-11 08:51:09 -04:00
Elior Nataf LackritzandGitHub b4991f1ba3 fix(langgraph): store an update_state input's writes on the checkpoint it updates (#9258)
`update_state(..., as_node="__input__")` stored its writes on the
checkpoint it saved. A DeltaChannel reads the writes stored on a
checkpoint's ancestors, not its own, so that checkpoint read back
without the input unless the channel snapshotted there, and once a later
checkpoint snapshotted, the input was gone for good. Only a raw `Pregel`
graph whose input channel is a DeltaChannel hits it, since a
`StateGraph`'s input goes to its start channel.

The input's writes now go on the checkpoint the update builds on, as a
node update's do. An update on a checkpoint the thread has moved past
stores nothing there and snapshots the DeltaChannels it writes on its
own checkpoint instead, the rule a node update already follows (#9165).

Both kinds of update now go through one helper, which also gives a
message an id before saving it, as the loop's `put_writes` does with
`ensure_message_ids`. Before, a message saved through `update_state`
without an id read back with id `None`, on the node path too, while the
same message from a node got one.

JS twin: langchain-ai/langgraphjs#2973, where review found it.

## Tests

`test_delta_channel_update_state.py`, sync and async: an update as input
reads back on its checkpoint and after the next run, with
`snapshot_frequency` 1 and 2, and an update as input to an older
checkpoint stays out of its other branch. The frequency 2 and
older-checkpoint cases fail on main. Five more check that a message
saved through `update_state` gets an id every read keeps, on the node
path (latest and older checkpoint) and the input path, sync and async;
they fail on main. The langgraph suite passes, and `lint_package` and
`lint_tests` are clean.
2026-10-10 11:57:46 +00:00
Elior Nataf LackritzandGitHub 6aa0afba68 fix(checkpoint): default InMemorySaver delta history to the latest checkpoint (#9257)
`InMemorySaver.get_delta_channel_history` looked its target up by
`config["configurable"].get("checkpoint_id", "")`, so a config without a
`checkpoint_id` matched no checkpoint, and every channel came back with
no seed and no writes, which a `DeltaChannel` reads as empty. The base
implementation, `PostgresSaver` and `SqliteSaver` read from the thread's
latest checkpoint instead, as `get_tuple` does.

A missing id now resolves to the namespace's newest checkpoint id, with
the same `max` that `InMemorySaver.get_tuple` uses. Graphs always pass a
checkpoint id here, so only direct callers of the saver API were
affected, and a call with an id does the same work as before.

Thanks @Hotragn, who pinned this down on #8242, where @longquanzheng's
branch had already fixed it.

JS twin: langchain-ai/langgraphjs#2979

## Tests

- `test_get_channel_writes_without_checkpoint_id_reads_the_latest` fails
on main (`{'messages': {'writes': []}}`).
- The `checkpoint`, `checkpoint-conformance`, `checkpoint-sqlite`,
`checkpoint-postgres` (Postgres 16), `prebuilt` and `langgraph` suites
pass, the last two without their Redis cases. The DeltaChannel
conformance suite also passes against `InMemorySaver`; CI skips that run
in `libs/checkpoint`, where the conformance package isn't installed, so
the new test lives in `test_memory.py`.
2026-10-10 00:59:20 +00:00
Elior Nataf LackritzandGitHub 12aeb0fddb fix(langgraph): read a run's DeltaChannel input back only on its own checkpoints (#9260)
A raw `Pregel` graph whose input channel is a DeltaChannel saved a run's
input to it as a `NULL_TASK_ID` write on the checkpoint the run started
from, under `"sync"` and `"async"` durability. Readers apply a
checkpoint's `NULL_TASK_ID` writes as its own state, so:

- on a new thread that checkpoint is never saved, and the first run's
input was lost for good: after `{"log": [0], "plain": [0], ...}`, `log`
read `[2]` next to `plain` `[0, 2]`;
- the last checkpoint of a run read the next run's input;
- a run with input from an older checkpoint (`checkpoint_id`) leaked its
input into the branch that already grew from it, and read that branch's
input when it had one.

The input's writes now go on the checkpoint the run starts from under a
task id of their own, `uuid5(checkpoint_id, INPUT)`, the one
`update_state` uses for an input update. On a new thread, or a
checkpoint a `checkpoint_id` addressed, which may already have children
that would read them, nothing is stored there and the input checkpoint
snapshots the channel instead. A normal run stores the same write as
before, and `durability="exit"`, which already kept the input right, is
unchanged. Threads saved before this keep the input writes already
stored on their checkpoints.

A `StateGraph` routes its input through its start channel, so only raw
`Pregel` graphs hit this.

JS twin: langchain-ai/langgraphjs#2981

## Tests

`test_delta_channel_run_input.py`, under every durability, with `invoke`
and `ainvoke`: two runs on a new thread, and a run with input from an
older checkpoint whose other branch had DeltaChannel input or none.
Every checkpoint reads the same in the DeltaChannel and a plain channel.
The 12 sync and async cases fail on main, and dropping the task id, the
new-thread snapshot or the addressed-checkpoint snapshot each fails its
own cases. The `langgraph` and `prebuilt` suites pass (without Redis),
and `lint_package` and `lint_tests` are clean.
2026-10-09 15:59:10 -04:00
Elior Nataf LackritzandGitHub d05236f805 fix(langgraph): snapshot a DeltaChannel that an update_state Overwrite resets (#9263)
A node that writes an `Overwrite` to a DeltaChannel makes the loop
snapshot the channel on the checkpoint that superstep saves, so no read
replays across the reset. The same write through `update_state`, as a
node or as input, didn't snapshot: the value read back right, but reads
of the update's checkpoint and the ones after it walked back past the
reset to the last snapshot. `update_state` now snapshots a DeltaChannel
an `Overwrite` resets, as the loop does.

Raised in review of langchain-ai/langgraphjs#2973. JS twin:
langchain-ai/langgraphjs#2982

## Tests

`test_delta_channel_update_overwrite.py`, sync and async: an `Overwrite`
through `update_state` as a node, and as input to a raw `Pregel` graph,
snapshots the channel on the checkpoint the update saves, and the value
reads back as the overwrite. All four fail on main. The `langgraph`
suite passes (without Redis), and `lint_package` and `lint_tests` are
clean.
2026-10-09 15:26:16 -04:00
Elior Nataf LackritzandGitHub 9d92f33cca fix(langgraph): keep a fork's DeltaChannel snapshot from starting nodes that never ran (#9264)
A fork from a checkpoint with pending writes to a DeltaChannel snapshots
the channel on the fork's first checkpoint, so the fork doesn't replay
the other branch's writes. If the channel had never been written there,
the snapshot also gives it its first version, and `_mark_bumps_seen`
only marked the bump seen for nodes that already have a `versions_seen`
entry. A node subscribed to the channel that never ran has none, so the
bump started it, and it read the empty value before the node that writes
the channel had run:

```python
graph.invoke("go", config)  # writer writes [1] to d, reader reads [1]
fork = graph.update_state(first_checkpoint, {"a": "go"}, as_node="__input__")
graph.get_state(fork).next  # ("writer", "reader"): resuming runs reader on [] first
```

A replay from that checkpoint saves the same kind of fork from the loop.
A channel bumped from no version was never written, so
`create_checkpoint` now also marks the bump seen for the nodes the
channel triggers, from the graph's `trigger_to_nodes`, at every call
that can bump. Only raw `Pregel` graphs hit this: `StateGraph` nodes
trigger on their `branch:to:*` channels, not on state keys. An
`as_node="__copy__"` fork still starts such a node, with a plain list
channel as well, so that case isn't DeltaChannel-specific and isn't
changed here.

Raised in review of langchain-ai/langgraphjs#2974. JS twin:
langchain-ai/langgraphjs#2983

## Tests

`test_delta_channel_seal_subscribers.py`: an update as input (sync and
async) and a replay from the checkpoint before a DeltaChannel's first
write leave only the writer next, and the reader then reads the written
value once. All three fail on main. The `langgraph` suite passes
(without Redis), and `lint_package` and `lint_tests` are clean.
2026-10-09 14:36:27 -04:00
Elior Nataf LackritzandGitHub 26356227c4 test(checkpoint): run the delta conformance suite against InMemorySaver in CI (#9261)
The conformance self-test in `libs/checkpoint-conformance` ran every
capability against `InMemorySaver` but asserted only the base ones.
`libs/checkpoint/tests/test_conformance_delta.py`, which ran the delta
conformance tests against it, always skips in CI: `libs/checkpoint`
can't depend on the conformance package, which depends on it. So no CI
job checked `InMemorySaver.get_delta_channel_history` against the delta
conformance tests.

The self-test now asserts every capability `InMemorySaver` implements,
and the skipped copy is gone.

## Tests

- With `InMemorySaver.get_delta_channel_history` broken to return no
history, the self-test now fails on 7 of the 10 delta conformance tests;
on main it still passes.
- `make format`, `make lint` and `make test` pass in
`libs/checkpoint-conformance`.
2026-10-09 13:44:55 -04:00
Elior Nataf LackritzandGitHub a5dbacae0d fix(checkpoint-postgres,checkpoint-sqlite): treat an empty checkpoint_id as the latest in delta history (#9262)
`PostgresSaver` and `SqliteSaver`, sync and async, read a DeltaChannel's
history from the latest checkpoint only when the config's
`checkpoint_id` was missing or `None`. An empty one was looked up as an
id, matched nothing, and came back with no seed and no writes, while
their own `get_tuple` treats an empty id as the latest checkpoint, as
the default implementation does. They now check for a missing or empty
id the way `get_tuple` does. The in-memory saver gets the same in #9257.

Graphs always pass the id of a saved checkpoint, so only direct callers
of the saver API hit this.

## Tests

-
`test_{sync,async}_empty_checkpoint_id_reads_from_the_latest_checkpoint`
in `checkpoint-postgres/tests/test_delta_pagination.py` and
`checkpoint-sqlite/tests/test_delta_parent_walk.py`: all four fail on
main.
- The `checkpoint-postgres` (Postgres 16) and `checkpoint-sqlite` suites
pass, and `lint_package` and `lint_tests` are clean in both.
2026-10-09 13:44:51 -04:00
cba111d8d6 fix(langgraph): don't save a checkpoint past a failed DeltaChannel write (#9235)
A DeltaChannel is rebuilt from its writes along the parent chain, so a
checkpoint saved without one of them reads back short for good. If
saving a DeltaChannel write failed (a database blip, say), the next
checkpoints were saved anyway and the channel lost that write, while
plain channels kept it. Three things let that happen:

- the sync loop waited for the delta writes before saving a checkpoint,
but never checked that they succeeded;
- both loops saved a checkpoint even when the previous save had failed,
so it pointed at a parent that was never saved;
- each save took the pending write list only when it started in the
background, and saves can start out of order, so a save could take a
later checkpoint's writes and not wait for its own. Waiting on writes
queued behind it, a save could also hang a sync run with
`durability="async"` (the default) once every worker was such a save: on
a 15-core machine a loop of 100 fast steps hung every time with default
settings, and with `max_concurrency=1` three steps are enough.

Now `_put_checkpoint` hands each save the writes submitted before it,
and a save fails if one of them or the previous save failed. The run
raises the write's error as before, the thread stays at its last
checkpoint that has all its writes, and running it again reruns the
task. For graphs without a DeltaChannel, a failed save now also stops
the later saves of that run, as JS already does, instead of saving
checkpoints whose parent is missing.

Thanks @shivangsharma01 for reproducing this in #8234, and @iroiro147,
whose #8299 had the first two changes before the issue-assignment check
closed it. Thanks also to @HarshShroff for the failed-save test case.

JS twin: langchain-ai/langgraphjs#2969

## Tests

`test_delta_channel_write_failures.py`:

- a saver fails one delta write once; on every durability, with `invoke`
and `ainvoke`, every saved checkpoint matches a plain channel and a
retry gets the write back. The `invoke` cases and `ainvoke` with
`durability="async"` fail on main.
- a saver fails the save of the checkpoint after `a` once; with `invoke`
and `ainvoke`, every saved checkpoint matches a plain channel and a
retry finishes the run. The `durability="async"` cases fail on main,
which saves the later checkpoints on top of the missing one.
- a delta graph run with `max_concurrency=1` finishes. It hangs on main.

Reverting any one of the three changes fails at least one of these.
`make format`, `make lint` and the langgraph suite pass. A 100-step
benchmark (delta and plain channels, every durability, with and without
1 ms of saver latency) makes the same saver calls and runs in the same
time as main, except for the case that hangs on main.

---------

Co-authored-by: iroiro147 <265728356+iroiro147@users.noreply.github.com>
2026-10-09 12:48:24 +00:00
Elior Nataf LackritzandGitHub 93a5a28008 perf(checkpoint-postgres): start the delta walk at the target (#9248)
Reading a DeltaChannel at an older checkpoint (the state at a past
checkpoint, each entry of `get_state_history`, a fork or replay from
one) walked `checkpoints` down from the newest row, so every newer row
was read before the target's chain. A checkpoint's ancestors have
smaller ids, the order `get_tuple` already relies on to find the latest
checkpoint, so the walk now starts at the target: the first page is
`checkpoint_id <= target`, later pages stay `< cursor`.

The cursor is never NULL anymore, so the `(cursor IS NULL OR
checkpoint_id < cursor)` predicate goes too. Once Postgres caches a
generic plan for the prepared statement (`from_conn_string` sets
`prepare_threshold=0`), that predicate is a filter rather than an index
condition, and each later page re-read every newer row.

`get_delta_channel_history` on a 6,000-checkpoint thread with a snapshot
every 200 checkpoints, median of three runs of 60 reads each, same
results before and after:

| target | main | this PR |
|---|---|---|
| newest | 3.1 ms | 2.7 ms |
| 1,000 back | 4.5 ms | 2.9 ms |
| 3,000 back | 8.9 ms | 2.8 ms |
| 5,000 back | 15.7 ms | 2.8 ms |

The target is now always the first row read, so the pagination tests
from #8556 can no longer leave it off the first page; they are renamed
for what they still cover, a walk that continues across pages.

## Test plan
- [x] `test_{sync,async}_walk_reads_nothing_newer_than_the_target`
record the rows the walk reads; both fail on main, which also reads the
three newer checkpoints.
- [x] `checkpoint-postgres` suite on Postgres 16: 281 passed, 3 skipped.
- [x] `langgraph` DeltaChannel tests against this saver (`-k "delta and
not redis"`): 1,154 passed, 1 skipped.
2026-10-09 08:31:15 -04:00
Baihao WangGitHubBaihao Wangopen-swe[bot] <open-swe@users.noreply.github.com>
bfcfea554e fix(checkpoint,checkpoint-postgres,checkpoint-sqlite): validate namespace labels before store dispatch (#9217)
The SQLite and Postgres stores save a namespace as its labels joined
with `.`, so a label containing `.` could resolve to a different
namespace: `("foo.bar",)` and `("foo", "bar")` share the same row. This
rejects those labels at every entry point.

- `BaseStore` and `AsyncBatchedBaseStore`: `get`, `search`, `delete`,
and `list_namespaces` (sync and async) reject labels that are empty, not
strings, or contain `.`. The check runs before an op is queued, so one
caller's bad label can't fail other callers' concurrent requests.
- New `langgraph.store.base.validate_op_namespace(op)` applies the same
label checks to a single op. `SqliteStore`/`AsyncSqliteStore` and
`PostgresStore`/`AsyncPostgresStore` call it from `_group_ops` before
any SQL runs, which covers direct `batch`/`abatch` calls. Other stores
that serialize namespaces as delimited text can call it too.
- Unchanged: hierarchical prefix search, empty-prefix `search(())`, `*`
wildcards in `list_namespaces`, and the existing `put` rules.

## Release note

`get`, `search`, `delete`, and `list_namespaces` now raise
`InvalidNamespaceError` for labels that are empty, not strings, or
contain `.`, where they previously returned `None` or `[]` (or, on the
SQLite and Postgres stores, matched a different namespace). `put`
already rejected these labels, so nothing written through `put` becomes
unreachable.

## Release order

Release `langgraph-checkpoint` first. `langgraph-checkpoint-sqlite` and
`langgraph-checkpoint-postgres` import `validate_op_namespace`, so they
need the `langgraph-checkpoint` release that ships it. `main` already
has `libs/checkpoint` at 4.3.0 and requires
`langgraph-checkpoint>=4.3.0` in both stores (from #8544), so this must
merge before 4.3.0 is released; if 4.3.0 ships first, bump to the next
version and raise both requirements to it. That also guarantees the
pre-queue checks are installed: with an older `langgraph-checkpoint`, an
invalid op would fail inside the shared `abatch` and reject the calls
batched with it.

---------

Co-authored-by: Baihao Wang <byhow@users.noreply.github.com>
Co-authored-by: open-swe[bot] <open-swe@users.noreply.github.com>
2026-10-08 16:10:55 -07:00
46 changed files with 2216 additions and 335 deletions
@@ -14,8 +14,8 @@ async def memory_checkpointer():
@pytest.mark.asyncio
async def test_validate_memory_base():
"""InMemorySaver passes all base capability tests."""
async def test_validate_memory():
"""InMemorySaver passes the tests of every capability it implements."""
report = await validate(memory_checkpointer)
report.print_report()
assert report.passed_all_base(), f"Base tests failed: {report.to_dict()}"
assert report.passed_all(), f"Capability tests failed: {report.to_dict()}"
@@ -451,7 +451,8 @@ class PostgresSaver(BasePostgresSaver):
* Stage 1 (paged): dynamic SELECT over `checkpoints` with three
columns per channel: its version, an `EXISTS` probe for a stored
blob at that version, and its inline value. Pages newest-first by
`checkpoint_id` with a cursor; page size is `_DELTA_PAGE_SIZE`.
`checkpoint_id`, starting at the target; page size is
`_DELTA_PAGE_SIZE`.
Stops paging when every channel has found its seed or a page comes
back short.
@@ -468,7 +469,7 @@ class PostgresSaver(BasePostgresSaver):
thread_id = config["configurable"]["thread_id"]
checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
checkpoint_id = get_checkpoint_id(config)
if checkpoint_id is None:
if not checkpoint_id:
target = self.get_tuple(config)
if target is None:
return {ch: {"writes": []} for ch in channels}
@@ -476,7 +477,7 @@ class PostgresSaver(BasePostgresSaver):
# Stage 1: paged K-JSONB-lookup scan, walking the parent chain in
# Python after each page. Stops as soon as every channel has its seed.
stage1_sql = _build_delta_stage1_sql(channels, paged=True)
stage1_sql = _build_delta_stage1_sql(channels, paged=True, include_cursor=True)
parent_of: dict[str, str | None] = {}
ver_by_i_by_cid: list[dict[str, str | None]] = [{} for _ in channels]
hb_by_i_by_cid: list[dict[str, bool]] = [{} for _ in channels]
@@ -486,7 +487,7 @@ class PostgresSaver(BasePostgresSaver):
seed_inline_by_ch: dict[str, Any] = {}
walk_cursor_by_ch: dict[str, str | None] = {}
seeded: set[str] = set()
cursor: str | None = None
cursor: str | None = checkpoint_id
with self._cursor() as cur:
while True:
@@ -495,7 +496,7 @@ class PostgresSaver(BasePostgresSaver):
# ver_i, blob channel, blob version, inline_i
stage1_params.extend([ch, ch, ch, ch])
stage1_params.extend(
[thread_id, checkpoint_ns, cursor, cursor, _DELTA_PAGE_SIZE]
[thread_id, checkpoint_ns, cursor, _DELTA_PAGE_SIZE]
)
cur.execute(stage1_sql, stage1_params)
page = cur.fetchall()
@@ -527,6 +528,7 @@ class PostgresSaver(BasePostgresSaver):
if len(seeded) == len(channels) or len(page) < _DELTA_PAGE_SIZE:
break
cursor = oldest
stage1_sql = _build_delta_stage1_sql(channels, paged=True)
# Stage 2: per-channel UNION ALL — one writes branch per channel
# with non-empty chain, plus one blob branch per seeded channel.
@@ -417,13 +417,13 @@ class AsyncPostgresSaver(BasePostgresSaver):
thread_id = config["configurable"]["thread_id"]
checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
checkpoint_id = get_checkpoint_id(config)
if checkpoint_id is None:
if not checkpoint_id:
target = await self.aget_tuple(config)
if target is None:
return {ch: {"writes": []} for ch in channels}
checkpoint_id = target.config["configurable"]["checkpoint_id"]
stage1_sql = _build_delta_stage1_sql(channels, paged=True)
stage1_sql = _build_delta_stage1_sql(channels, paged=True, include_cursor=True)
parent_of: dict[str, str | None] = {}
ver_by_i_by_cid: list[dict[str, str | None]] = [{} for _ in channels]
hb_by_i_by_cid: list[dict[str, bool]] = [{} for _ in channels]
@@ -433,7 +433,7 @@ class AsyncPostgresSaver(BasePostgresSaver):
seed_inline_by_ch: dict[str, Any] = {}
walk_cursor_by_ch: dict[str, str | None] = {}
seeded: set[str] = set()
cursor: str | None = None
cursor: str | None = checkpoint_id
async with self._cursor() as cur:
while True:
@@ -442,7 +442,7 @@ class AsyncPostgresSaver(BasePostgresSaver):
# ver_i, blob channel, blob version, inline_i
stage1_params.extend([ch, ch, ch, ch])
stage1_params.extend(
[thread_id, checkpoint_ns, cursor, cursor, _DELTA_PAGE_SIZE]
[thread_id, checkpoint_ns, cursor, _DELTA_PAGE_SIZE]
)
await cur.execute(stage1_sql, stage1_params)
page = await cur.fetchall()
@@ -472,6 +472,7 @@ class AsyncPostgresSaver(BasePostgresSaver):
if len(seeded) == len(channels) or len(page) < _DELTA_PAGE_SIZE:
break
cursor = oldest
stage1_sql = _build_delta_stage1_sql(channels, paged=True)
channels_with_chain = [ch for ch in channels if chain_by_ch[ch]]
channels_with_seed = [ch for ch in channels if seed_ver_by_ch[ch] is not None]
@@ -178,10 +178,13 @@ class _DeltaStage2Row(TypedDict, total=False):
# `_build_delta_stage2_sql` document their shapes.
def _build_delta_stage1_sql(channels: Sequence[str], *, paged: bool) -> str:
def _build_delta_stage1_sql(
channels: Sequence[str], *, paged: bool, include_cursor: bool = False
) -> str:
"""Build stage 1 SQL with K parallel version lookups + seed probes.
For channels=["messages", "files"] (with `paged=True`) the result is::
For channels=["messages", "files"] (with `paged=True, include_cursor=True`)
the result is::
SELECT checkpoint_id, parent_checkpoint_id,
checkpoint -> 'channel_versions' ->> %s AS ver_0,
@@ -197,7 +200,7 @@ def _build_delta_stage1_sql(channels: Sequence[str], *, paged: bool) -> str:
checkpoint -> 'channel_values' -> %s AS inline_1
FROM checkpoints
WHERE thread_id = %s AND checkpoint_ns = %s
AND (%s::text IS NULL OR checkpoint_id < %s)
AND checkpoint_id <= %s
ORDER BY checkpoint_id DESC
LIMIT %s
@@ -238,9 +241,16 @@ def _build_delta_stage1_sql(channels: Sequence[str], *, paged: bool) -> str:
and uses safe identifiers).
Caller must extend params with `[ch_0 x4, ch_1 x4, ..., thread_id, ns,
cursor, cursor, page_size]` when `paged=True` — four per channel: the
version lookup, the blob's channel, the version the blob must match, and the
inline lookup.
cursor, page_size]` when `paged=True` — four per channel: the version
lookup, the blob's channel, the version the blob must match, and the inline
lookup.
Pages run newest-first from the target down. The first page passes the
target as the cursor with `include_cursor=True`, so it opens with the
target's own row, whose parent starts the walk; each later page continues
below the oldest row read. A checkpoint's ancestors have smaller ids (uuid6
is time-ordered, which `get_tuple` also relies on to find the latest
checkpoint), so no row newer than the target is part of its chain.
When `paged=False`, the WHERE has no cursor predicate and there's no
LIMIT/ORDER BY — kept as a non-public helper for tests/diagnostics.
@@ -264,7 +274,7 @@ def _build_delta_stage1_sql(channels: Sequence[str], *, paged: bool) -> str:
)
if paged:
sql += (
" AND (%s::text IS NULL OR checkpoint_id < %s)"
f" AND checkpoint_id {'<=' if include_cursor else '<'} %s"
" ORDER BY checkpoint_id DESC LIMIT %s"
)
return sql
@@ -413,8 +423,8 @@ class BasePostgresSaver(BaseCheckpointSaver[str]):
materialized at this point),
(c) the next ancestor cid isn't in `parent_of` yet (waiting for
a later page; the cursor stays put), or
(d) the target's own row isn't in `parent_of` yet (the walk has
not started; no cursor is set, so a later page retries).
(d) the target's own row isn't in `parent_of` (the target doesn't
exist, so the walk never starts).
Mutates `chain_by_ch`, `seed_ver_by_ch`, `seed_inline_by_ch`,
`walk_cursor_by_ch`, and `seeded` in place.
@@ -422,8 +432,9 @@ class BasePostgresSaver(BaseCheckpointSaver[str]):
for i, ch in enumerate(channels):
if ch in seeded:
continue
# Pages start at the thread head, so the target may not have
# loaded yet; a `None` cursor would read as "target is a root".
# The first page opens with the target's row, so it's missing only
# when the target doesn't exist; a `None` cursor would read as
# "target is a root".
if ch not in walk_cursor_by_ch:
if target_id not in parent_of:
continue
@@ -36,6 +36,7 @@ from langgraph.store.base import (
ensure_embeddings,
get_text_at_path,
tokenize_path,
validate_op_namespace,
)
from psycopg import Capabilities, Connection, Cursor, Pipeline
from psycopg.rows import DictRow, dict_row
@@ -1386,6 +1387,7 @@ def _group_ops(ops: Iterable[Op]) -> tuple[dict[type, list[tuple[int, Op]]], int
grouped_ops: dict[type, list[tuple[int, Op]]] = defaultdict(list)
tot = 0
for idx, op in enumerate(ops):
validate_op_namespace(op)
grouped_ops[type(op)].append((idx, op))
tot += 1
return grouped_ops, tot
+1 -1
View File
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
[project]
name = "langgraph-checkpoint-postgres"
version = "3.1.2"
version = "3.1.3"
description = "Library with a Postgres implementation of LangGraph checkpoint saver."
authors = []
requires-python = ">=3.10"
@@ -13,8 +13,10 @@ import pytest
from langchain_core.embeddings import Embeddings
from langgraph.store.base import (
GetOp,
InvalidNamespaceError,
Item,
ListNamespacesOp,
MatchCondition,
PutOp,
SearchOp,
)
@@ -871,3 +873,65 @@ async def test_omit_expired_search_pagination(store: AsyncPostgresStore) -> None
page2 = await store.asearch(ns, limit=2, offset=2)
assert [i.key for i in page1] == ["a", "b"]
assert [i.key for i in page2] == ["c"]
@pytest.mark.parametrize("namespace", [("foo.bar",), ("foo", ""), ("foo", 1)])
async def test_abatch_rejects_invalid_namespace_labels(
store: AsyncPostgresStore, namespace: tuple
) -> None:
await store.aput(("foo", "bar"), "key", {"original": True})
for op in (
GetOp(namespace, "key"),
GetOp(namespace, "key", refresh_ttl=True),
PutOp(namespace, "key", {"changed": True}),
PutOp(namespace, "key", None),
SearchOp(namespace),
ListNamespacesOp((MatchCondition("prefix", namespace),)),
ListNamespacesOp((MatchCondition("suffix", namespace),)),
):
with pytest.raises(InvalidNamespaceError):
await store.abatch([op])
item = await store.aget(("foo", "bar"), "key")
assert item is not None and item.value == {"original": True}
async def test_invalid_namespace_only_fails_its_own_call(
store: AsyncPostgresStore,
) -> None:
"""Concurrent calls share one `abatch`, which fails every op if it raises.
Labels are checked before an op is queued, so one caller's bad label cannot
fail another caller's request.
"""
await store.aput(("foo", "bar"), "key", {"original": True})
valid, invalid = await asyncio.gather(
store.aget(("foo", "bar"), "key"),
store.aget(("foo.bar",), "key"),
return_exceptions=True,
)
assert isinstance(valid, Item) and valid.value == {"original": True}
assert isinstance(invalid, InvalidNamespaceError)
async def test_sync_methods_reject_invalid_namespace_labels(
store: AsyncPostgresStore,
) -> None:
"""The sync wrappers run off the event loop thread and must validate too."""
await store.aput(("foo", "bar"), "key", {"original": True})
for call in (
lambda: store.get(("foo.bar",), "key"),
lambda: store.search(("foo.bar",)),
lambda: store.delete(("foo.bar",), "key"),
lambda: store.list_namespaces(prefix=("foo.bar",)),
lambda: store.batch([GetOp(("foo.bar",), "key")]),
):
with pytest.raises(InvalidNamespaceError):
await asyncio.to_thread(call)
item = await store.aget(("foo", "bar"), "key")
assert item is not None and item.value == {"original": True}
@@ -1,5 +1,6 @@
from __future__ import annotations
from collections.abc import Mapping, Sequence
from typing import Any
from uuid import uuid4
@@ -14,7 +15,7 @@ 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 langgraph.checkpoint.postgres.base import _DELTA_PAGE_SIZE, BasePostgresSaver
from tests.conftest import DEFAULT_URI
CHANNEL = "items"
@@ -23,8 +24,8 @@ SEED_STEP = 1
SEED_VALUE = [10, 20]
TARGET_STEP = 4
# The real page size is the control; the rest leave the target off the first
# page (three checkpoints are newer than it).
# The real page size is the control; the rest split the walk from the target
# to its seed across pages.
PAGE_SIZES = [_DELTA_PAGE_SIZE, 3, 2, 1]
@@ -92,7 +93,7 @@ def _assert_history(entry: DeltaChannelHistory, page_size: int) -> None:
@pytest.mark.parametrize("page_size", PAGE_SIZES)
async def test_async_target_older_than_the_first_page(
async def test_async_walk_continues_across_pages(
page_size: int, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setattr("langgraph.checkpoint.postgres.aio._DELTA_PAGE_SIZE", page_size)
@@ -106,7 +107,7 @@ async def test_async_target_older_than_the_first_page(
@pytest.mark.parametrize("page_size", PAGE_SIZES)
def test_sync_target_older_than_the_first_page(
def test_sync_walk_continues_across_pages(
page_size: int, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setattr("langgraph.checkpoint.postgres._DELTA_PAGE_SIZE", page_size)
@@ -119,6 +120,84 @@ def test_sync_target_older_than_the_first_page(
_assert_history(result[CHANNEL], page_size)
def _record_rows_read(monkeypatch: pytest.MonkeyPatch) -> list[str]:
read: list[str] = []
ingest = BasePostgresSaver._ingest_stage1_page
def record(rows: Sequence[Mapping[str, Any]], *args: Any) -> str | None:
read.extend(row["checkpoint_id"] for row in rows)
return ingest(rows, *args)
monkeypatch.setattr(BasePostgresSaver, "_ingest_stage1_page", staticmethod(record))
return read
def _ids_from_target_down(configs: list[dict]) -> list[str]:
return [c["configurable"]["checkpoint_id"] for c in configs[TARGET_STEP::-1]]
async def test_async_walk_reads_nothing_newer_than_the_target(
monkeypatch: pytest.MonkeyPatch,
) -> None:
read = _record_rows_read(monkeypatch)
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], _DELTA_PAGE_SIZE)
assert read == _ids_from_target_down(configs)
def test_sync_walk_reads_nothing_newer_than_the_target(
monkeypatch: pytest.MonkeyPatch,
) -> None:
read = _record_rows_read(monkeypatch)
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], _DELTA_PAGE_SIZE)
assert read == _ids_from_target_down(configs)
def _empty_checkpoint_id(config: dict) -> dict:
return {"configurable": {**config["configurable"], "checkpoint_id": ""}}
async def test_async_empty_checkpoint_id_reads_from_the_latest_checkpoint() -> None:
async with AsyncPostgresSaver.from_conn_string(DEFAULT_URI) as saver:
await saver.setup()
configs = await _abuild_chain(saver)
latest = await saver.aget_delta_channel_history(
config=configs[-1], channels=[CHANNEL]
)
result = await saver.aget_delta_channel_history(
config=_empty_checkpoint_id(configs[-1]), channels=[CHANNEL]
)
assert latest[CHANNEL]["writes"]
assert result == latest
def test_sync_empty_checkpoint_id_reads_from_the_latest_checkpoint() -> None:
with PostgresSaver.from_conn_string(DEFAULT_URI) as saver:
saver.setup()
configs = _build_chain(saver)
latest = saver.get_delta_channel_history(config=configs[-1], channels=[CHANNEL])
result = saver.get_delta_channel_history(
config=_empty_checkpoint_id(configs[-1]), channels=[CHANNEL]
)
assert latest[CHANNEL]["writes"]
assert result == latest
async def test_root_target_has_no_history_and_still_terminates(
monkeypatch: pytest.MonkeyPatch,
) -> None:
@@ -11,6 +11,7 @@ import pytest
from langchain_core.embeddings import Embeddings
from langgraph.store.base import (
GetOp,
InvalidNamespaceError,
Item,
ListNamespacesOp,
MatchCondition,
@@ -1164,3 +1165,51 @@ def test_namespace_labels_with_trailing_newline(store) -> None:
assert set(store.list_namespaces(prefix=["users", "alice"], limit=100)) == {
("users", "alice"),
}
@pytest.mark.parametrize("namespace", [("foo.bar",), ("foo", ""), ("foo", 1)])
@pytest.mark.parametrize(
"kind", ["get", "put", "delete", "search", "list_prefix", "list_suffix"]
)
def test_batch_rejects_invalid_namespace_labels(
store, namespace: tuple, kind: str
) -> None:
"""Ops passed straight to `batch` must not reach another namespace.
Namespaces are stored dot-joined, so `("foo.bar",)` flattens to the same
text as `("foo", "bar")`. `BaseStore` methods validate labels themselves,
but `batch` takes ops as given.
"""
op = {
"get": GetOp(namespace, "key"),
"put": PutOp(namespace, "key", {"changed": True}),
"delete": PutOp(namespace, "key", None),
"search": SearchOp(namespace),
"list_prefix": ListNamespacesOp((MatchCondition("prefix", namespace),)),
"list_suffix": ListNamespacesOp((MatchCondition("suffix", namespace),)),
}[kind]
store.put(("foo", "bar"), "key", {"original": True})
with pytest.raises(InvalidNamespaceError):
store.batch([PutOp(("valid",), "key", {}), op])
item = store.get(("foo", "bar"), "key")
assert item is not None and item.value == {"original": True}
# The whole batch is rejected before any SQL runs.
assert store.get(("valid",), "key") is None
def test_batch_allows_empty_search_prefix_and_listing_wildcards(
store,
) -> None:
store.put(("foo", "bar"), "key", {"v": 1})
found, listed = store.batch(
[
SearchOp(()),
ListNamespacesOp((MatchCondition("prefix", ("foo", "*")),)),
]
)
assert [item.namespace for item in found] == [("foo", "bar")]
assert listed == [("foo", "bar")]
+1 -1
View File
@@ -353,7 +353,7 @@ test = [
[[package]]
name = "langgraph-checkpoint-postgres"
version = "3.1.2"
version = "3.1.3"
source = { editable = "." }
dependencies = [
{ name = "langgraph-checkpoint" },
@@ -539,7 +539,7 @@ class SqliteSaver(BaseCheckpointSaver[str]):
thread_id = str(config["configurable"]["thread_id"])
checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
checkpoint_id = get_checkpoint_id(config)
if checkpoint_id is None:
if not checkpoint_id:
target = self.get_tuple(config)
if target is None:
return {ch: {"writes": []} for ch in channels}
@@ -653,7 +653,7 @@ class AsyncSqliteSaver(BaseCheckpointSaver[str]):
thread_id = str(config["configurable"]["thread_id"])
checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
checkpoint_id = get_checkpoint_id(config)
if checkpoint_id is None:
if not checkpoint_id:
target = await self.aget_tuple(config)
if target is None:
return {ch: {"writes": []} for ch in channels}
@@ -28,6 +28,7 @@ from langgraph.store.base import (
ensure_embeddings,
get_text_at_path,
tokenize_path,
validate_op_namespace,
)
_AIO_ERROR_MSG = (
@@ -257,6 +258,7 @@ def _group_ops(ops: Iterable[Op]) -> tuple[dict[type, list[tuple[int, Op]]], int
grouped_ops: dict[type, list[tuple[int, Op]]] = defaultdict(list)
tot = 0
for idx, op in enumerate(ops):
validate_op_namespace(op)
grouped_ops[type(op)].append((idx, op))
tot += 1
return grouped_ops, tot
+1 -1
View File
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
[project]
name = "langgraph-checkpoint-sqlite"
version = "3.1.1"
version = "3.1.2"
description = "Library with a SQLite implementation of LangGraph checkpoint saver."
authors = []
requires-python = ">=3.10"
@@ -9,8 +9,10 @@ from typing import cast
import pytest
from langgraph.store.base import (
GetOp,
InvalidNamespaceError,
Item,
ListNamespacesOp,
MatchCondition,
PutOp,
SearchOp,
)
@@ -745,3 +747,65 @@ async def test_async_namespace_segment_boundary(store: AsyncSqliteStore) -> None
assert set(await store.alist_namespaces(suffix=["alice"], limit=100)) == {
("uid", "users", "alice"),
}
@pytest.mark.parametrize("namespace", [("foo.bar",), ("foo", ""), ("foo", 1)])
async def test_abatch_rejects_invalid_namespace_labels(
store: AsyncSqliteStore, namespace: tuple
) -> None:
await store.aput(("foo", "bar"), "key", {"original": True})
for op in (
GetOp(namespace, "key"),
GetOp(namespace, "key", refresh_ttl=True),
PutOp(namespace, "key", {"changed": True}),
PutOp(namespace, "key", None),
SearchOp(namespace),
ListNamespacesOp((MatchCondition("prefix", namespace),)),
ListNamespacesOp((MatchCondition("suffix", namespace),)),
):
with pytest.raises(InvalidNamespaceError):
await store.abatch([op])
item = await store.aget(("foo", "bar"), "key")
assert item is not None and item.value == {"original": True}
async def test_invalid_namespace_only_fails_its_own_call(
store: AsyncSqliteStore,
) -> None:
"""Concurrent calls share one `abatch`, which fails every op if it raises.
Labels are checked before an op is queued, so one caller's bad label cannot
fail another caller's request.
"""
await store.aput(("foo", "bar"), "key", {"original": True})
valid, invalid = await asyncio.gather(
store.aget(("foo", "bar"), "key"),
store.aget(("foo.bar",), "key"),
return_exceptions=True,
)
assert isinstance(valid, Item) and valid.value == {"original": True}
assert isinstance(invalid, InvalidNamespaceError)
async def test_sync_methods_reject_invalid_namespace_labels(
store: AsyncSqliteStore,
) -> None:
"""The sync wrappers run off the event loop thread and must validate too."""
await store.aput(("foo", "bar"), "key", {"original": True})
for call in (
lambda: store.get(("foo.bar",), "key"),
lambda: store.search(("foo.bar",)),
lambda: store.delete(("foo.bar",), "key"),
lambda: store.list_namespaces(prefix=("foo.bar",)),
lambda: store.batch([GetOp(("foo.bar",), "key")]),
):
with pytest.raises(InvalidNamespaceError):
await asyncio.to_thread(call)
item = await store.aget(("foo", "bar"), "key")
assert item is not None and item.value == {"original": True}
@@ -65,6 +65,39 @@ async def test_async_walk_reaches_parent_whatever_the_id_order(
assert got[CHANNEL] == EXPECTED
EMPTY_CHECKPOINT_ID: dict[str, Any] = {
"configurable": {**CONFIG["configurable"], "checkpoint_id": ""}
}
def test_sync_empty_checkpoint_id_reads_from_the_latest_checkpoint() -> None:
with SqliteSaver.from_conn_string(":memory:") as saver:
root = saver.put(CONFIG, _checkpoint("a-older", {CHANNEL: "seed"}), {}, {})
saver.put_writes(root, [(CHANNEL, "write-root")], "task")
saver.put(root, _checkpoint("z-newer", {}), {}, {})
got = saver.get_delta_channel_history(
config=EMPTY_CHECKPOINT_ID, channels=[CHANNEL]
)
assert got[CHANNEL] == EXPECTED
async def test_async_empty_checkpoint_id_reads_from_the_latest_checkpoint() -> None:
async with AsyncSqliteSaver.from_conn_string(":memory:") as saver:
root = await saver.aput(
CONFIG, _checkpoint("a-older", {CHANNEL: "seed"}), {}, {}
)
await saver.aput_writes(root, [(CHANNEL, "write-root")], "task")
await saver.aput(root, _checkpoint("z-newer", {}), {}, {})
got = await saver.aget_delta_channel_history(
config=EMPTY_CHECKPOINT_ID, channels=[CHANNEL]
)
assert got[CHANNEL] == EXPECTED
def test_walk_reaches_root_of_long_chain_with_descending_ids() -> None:
steps = 40
with SqliteSaver.from_conn_string(":memory:") as saver:
@@ -14,6 +14,7 @@ import pytest
from langchain_core.embeddings import Embeddings
from langgraph.store.base import (
GetOp,
InvalidNamespaceError,
Item,
ListNamespacesOp,
MatchCondition,
@@ -1435,3 +1436,51 @@ def test_list_namespaces_metacharacter_labels(store: SqliteStore) -> None:
assert set(store.list_namespaces(prefix=[label, "child"], limit=100)) == {
(label, "child"),
}
@pytest.mark.parametrize("namespace", [("foo.bar",), ("foo", ""), ("foo", 1)])
@pytest.mark.parametrize(
"kind", ["get", "put", "delete", "search", "list_prefix", "list_suffix"]
)
def test_batch_rejects_invalid_namespace_labels(
store: SqliteStore, namespace: tuple, kind: str
) -> None:
"""Ops passed straight to `batch` must not reach another namespace.
Namespaces are stored dot-joined, so `("foo.bar",)` flattens to the same
text as `("foo", "bar")`. `BaseStore` methods validate labels themselves,
but `batch` takes ops as given.
"""
op = {
"get": GetOp(namespace, "key"),
"put": PutOp(namespace, "key", {"changed": True}),
"delete": PutOp(namespace, "key", None),
"search": SearchOp(namespace),
"list_prefix": ListNamespacesOp((MatchCondition("prefix", namespace),)),
"list_suffix": ListNamespacesOp((MatchCondition("suffix", namespace),)),
}[kind]
store.put(("foo", "bar"), "key", {"original": True})
with pytest.raises(InvalidNamespaceError):
store.batch([PutOp(("valid",), "key", {}), op])
item = store.get(("foo", "bar"), "key")
assert item is not None and item.value == {"original": True}
# The whole batch is rejected before any SQL runs.
assert store.get(("valid",), "key") is None
def test_batch_allows_empty_search_prefix_and_listing_wildcards(
store: SqliteStore,
) -> None:
store.put(("foo", "bar"), "key", {"v": 1})
found, listed = store.batch(
[
SearchOp(()),
ListNamespacesOp((MatchCondition("prefix", ("foo", "*")),)),
]
)
assert [item.namespace for item in found] == [("foo", "bar")]
assert listed == [("foo", "bar")]
+1 -1
View File
@@ -360,7 +360,7 @@ test = [
[[package]]
name = "langgraph-checkpoint-sqlite"
version = "3.1.1"
version = "3.1.2"
source = { editable = "." }
dependencies = [
{ name = "aiosqlite" },
@@ -170,8 +170,8 @@ class InMemorySaver(
return {}
thread_id = config["configurable"]["thread_id"]
checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
checkpoint_id = config["configurable"].get("checkpoint_id", "")
ns_storage = self.storage.get(thread_id, {}).get(checkpoint_ns, {})
checkpoint_id = get_checkpoint_id(config) or max(ns_storage, default="")
chain: list[str] = []
target_entry = ns_storage.get(checkpoint_id)
@@ -771,7 +771,12 @@ class BaseStore(ABC):
Returns:
The retrieved item or `None` if not found.
Raises:
InvalidNamespaceError: If a namespace label is empty, is not a string,
or contains a period (`.`).
"""
_validate_namespace_labels(namespace)
return self.batch(
[GetOp(namespace, str(key), _ensure_refresh(self.ttl_config, refresh_ttl))]
)[0]
@@ -801,6 +806,10 @@ class BaseStore(ABC):
Returns:
List of items matching the search criteria.
Raises:
InvalidNamespaceError: If a `namespace_prefix` label is empty, is not a
string, or contains a period (`.`).
???+ example "Examples"
Basic filtering:
@@ -840,6 +849,7 @@ class BaseStore(ABC):
Natural language search support depends on your store implementation
and requires proper embedding configuration.
"""
_validate_namespace_labels(namespace_prefix)
return self.batch(
[
SearchOp(
@@ -887,6 +897,11 @@ class BaseStore(ABC):
By default, the expiration timer refreshes on both read operations (get/search)
and write operations (put/update), whenever the item is included in the operation.
Raises:
InvalidNamespaceError: If the namespace is empty, its root label is
`"langgraph"`, or a label is empty, is not a string, or contains a
period (`.`).
Note:
Indexing support depends on your store implementation.
If you do not initialize the store with indexing capabilities,
@@ -940,7 +955,12 @@ class BaseStore(ABC):
Args:
namespace: Hierarchical path for the item.
key: Unique identifier within the namespace.
Raises:
InvalidNamespaceError: If a namespace label is empty, is not a string,
or contains a period (`.`).
"""
_validate_namespace_labels(namespace)
self.batch([PutOp(namespace, str(key), None, ttl=None)])
def list_namespaces(
@@ -969,6 +989,10 @@ class BaseStore(ABC):
A list of namespace tuples that match the criteria. Each tuple represents a
full namespace path up to `max_depth`.
Raises:
InvalidNamespaceError: If a `prefix` or `suffix` label is empty, is not a
string, or contains a period (`.`).
???+ example "Examples":
Setting `max_depth=3`. Given the namespaces:
@@ -984,6 +1008,8 @@ class BaseStore(ABC):
# [("a", "b", "c"), ("a", "b", "d"), ("a", "b", "f")]
```
"""
_validate_namespace_labels(prefix or ())
_validate_namespace_labels(suffix or ())
match_conditions = []
if prefix:
match_conditions.append(MatchCondition(match_type="prefix", path=prefix))
@@ -1013,7 +1039,12 @@ class BaseStore(ABC):
Returns:
The retrieved item or `None` if not found.
Raises:
InvalidNamespaceError: If a namespace label is empty, is not a string,
or contains a period (`.`).
"""
_validate_namespace_labels(namespace)
return (
await self.abatch(
[
@@ -1052,6 +1083,10 @@ class BaseStore(ABC):
Returns:
List of items matching the search criteria.
Raises:
InvalidNamespaceError: If a `namespace_prefix` label is empty, is not a
string, or contains a period (`.`).
???+ example "Examples"
Basic filtering:
@@ -1091,6 +1126,7 @@ class BaseStore(ABC):
Natural language search support depends on your store implementation
and requires proper embedding configuration.
"""
_validate_namespace_labels(namespace_prefix)
return (
await self.abatch(
[
@@ -1140,6 +1176,11 @@ class BaseStore(ABC):
By default, the expiration timer refreshes on both read operations (get/search)
and write operations (put/update), whenever the item is included in the operation.
Raises:
InvalidNamespaceError: If the namespace is empty, its root label is
`"langgraph"`, or a label is empty, is not a string, or contains a
period (`.`).
Note:
Indexing support depends on your store implementation.
If you do not initialize the store with indexing capabilities,
@@ -1201,7 +1242,12 @@ class BaseStore(ABC):
Args:
namespace: Hierarchical path for the item.
key: Unique identifier within the namespace.
Raises:
InvalidNamespaceError: If a namespace label is empty, is not a string,
or contains a period (`.`).
"""
_validate_namespace_labels(namespace)
await self.abatch([PutOp(namespace, str(key), None)])
async def alist_namespaces(
@@ -1230,6 +1276,10 @@ class BaseStore(ABC):
A list of namespace tuples that match the criteria. Each tuple represents a
full namespace path up to `max_depth`.
Raises:
InvalidNamespaceError: If a `prefix` or `suffix` label is empty, is not a
string, or contains a period (`.`).
???+ example "Examples"
Setting `max_depth=3` with existing namespaces:
@@ -1245,6 +1295,8 @@ class BaseStore(ABC):
# Returns: [("a", "b", "c"), ("a", "b", "d"), ("a", "b", "f")]
```
"""
_validate_namespace_labels(prefix or ())
_validate_namespace_labels(suffix or ())
match_conditions = []
if prefix:
match_conditions.append(MatchCondition(match_type="prefix", path=prefix))
@@ -1263,6 +1315,14 @@ class BaseStore(ABC):
def _validate_namespace(namespace: tuple[str, ...]) -> None:
if not namespace:
raise InvalidNamespaceError("Namespace cannot be empty.")
_validate_namespace_labels(namespace)
if namespace[0] == "langgraph":
raise InvalidNamespaceError(
f'Root label for namespace cannot be "langgraph". Got: {namespace}'
)
def _validate_namespace_labels(namespace: tuple[str, ...]) -> None:
for label in namespace:
if not isinstance(label, str):
raise InvalidNamespaceError(
@@ -1277,10 +1337,27 @@ def _validate_namespace(namespace: tuple[str, ...]) -> None:
raise InvalidNamespaceError(
f"Namespace labels cannot be empty strings. Got {label} in {namespace}"
)
if namespace[0] == "langgraph":
raise InvalidNamespaceError(
f'Root label for namespace cannot be "langgraph". Got: {namespace}'
)
def validate_op_namespace(op: Op) -> None:
"""Validate the namespace labels an op carries before a store executes it.
`BaseStore` methods check labels before batching, but ops passed directly to
`batch`/`abatch` skip those methods. Stores that serialize namespaces as
delimited text should call this for every op they execute, so a label such
as `"foo.bar"` cannot address the namespace `("foo", "bar")`.
Raises:
InvalidNamespaceError: If a label is empty, is not a string, or contains
a period (`.`).
"""
if isinstance(op, (GetOp, PutOp)):
_validate_namespace_labels(op.namespace)
elif isinstance(op, SearchOp):
_validate_namespace_labels(op.namespace_prefix)
elif isinstance(op, ListNamespacesOp):
for condition in op.match_conditions or ():
_validate_namespace_labels(condition.path)
def _ensure_refresh(
@@ -1319,4 +1396,5 @@ __all__ = [
"ensure_embeddings",
"tokenize_path",
"get_text_at_path",
"validate_op_namespace",
]
@@ -25,6 +25,7 @@ from langgraph.store.base import (
_ensure_refresh,
_ensure_ttl,
_validate_namespace,
_validate_namespace_labels,
)
F = TypeVar("F", bound=Callable)
@@ -86,6 +87,7 @@ class AsyncBatchedBaseStore(BaseStore):
*,
refresh_ttl: bool | None = None,
) -> Item | None:
_validate_namespace_labels(namespace)
self._ensure_task()
fut = self._loop.create_future()
self._aqueue.put_nowait(
@@ -111,6 +113,7 @@ class AsyncBatchedBaseStore(BaseStore):
offset: int = 0,
refresh_ttl: bool | None = None,
) -> list[SearchItem]:
_validate_namespace_labels(namespace_prefix)
self._ensure_task()
fut = self._loop.create_future()
self._aqueue.put_nowait(
@@ -155,6 +158,7 @@ class AsyncBatchedBaseStore(BaseStore):
namespace: tuple[str, ...],
key: str,
) -> None:
_validate_namespace_labels(namespace)
self._ensure_task()
fut = self._loop.create_future()
self._aqueue.put_nowait((fut, PutOp(namespace, key, None)))
@@ -169,6 +173,8 @@ class AsyncBatchedBaseStore(BaseStore):
limit: int = 100,
offset: int = 0,
) -> list[tuple[str, ...]]:
_validate_namespace_labels(prefix or ())
_validate_namespace_labels(suffix or ())
self._ensure_task()
fut = self._loop.create_future()
match_conditions = []
@@ -1,37 +0,0 @@
"""Run delta-channel conformance capabilities against InMemorySaver."""
from __future__ import annotations
import pytest
conformance = pytest.importorskip(
"langgraph.checkpoint.conformance",
reason="langgraph-checkpoint-conformance not installed",
)
@pytest.mark.asyncio
async def test_delta_channel_conformance():
# Imported inside the test: the module-level importorskip above is what
# makes these safe, so they cannot move to the top of the file.
from langgraph.checkpoint.conformance import validate # noqa: PLC0415
from langgraph.checkpoint.conformance.initializer import ( # noqa: PLC0415
checkpointer_test,
)
from langgraph.checkpoint.memory import InMemorySaver # noqa: PLC0415
@checkpointer_test(name="InMemorySaver")
async def mem_saver():
yield InMemorySaver()
report = await validate(
mem_saver,
capabilities={
"delta_channel_history",
},
)
for cap, result in report.results.items():
if result.passed is False:
details = "\n".join(result.failures or [])
pytest.fail(f"Capability {cap} failed:\n{details}")
+13
View File
@@ -420,6 +420,19 @@ class TestInMemorySaverDeltaChannel:
assert "seed" not in result
assert result["writes"] == []
def test_get_channel_writes_without_checkpoint_id_reads_the_latest(self) -> None:
saver = InMemorySaver()
thread: RunnableConfig = {
"configurable": {"thread_id": "t1", "checkpoint_ns": ""}
}
parent = saver.put(thread, empty_checkpoint(), {}, {})
saver.put_writes(parent, [("messages", "hi")], "task1")
saver.put(parent, empty_checkpoint(), {}, {})
result = saver.get_delta_channel_history(config=thread, channels=["messages"])
assert result == {"messages": {"writes": [("task1", "messages", "hi")]}}
class TestBaseFallbackGetChannelWrites:
"""Exercises the `BaseCheckpointSaver.get_delta_channel_history` default
+125
View File
@@ -13,10 +13,14 @@ from langgraph.store.base import (
GetOp,
InvalidNamespaceError,
Item,
ListNamespacesOp,
MatchCondition,
Op,
PutOp,
Result,
SearchOp,
get_text_at_path,
validate_op_namespace,
)
from langgraph.store.base.batch import AsyncBatchedBaseStore
from langgraph.store.memory import InMemoryStore
@@ -528,6 +532,127 @@ async def test_cannot_put_empty_namespace() -> None:
assert (await async_store.aget(("valid", "namespace"), "key")) is None
INVALID_NAMESPACES = [("foo.bar",), ("foo", ""), (123,)]
NAMESPACE_METHODS = ["get", "delete", "search", "prefix", "suffix"]
@pytest.mark.parametrize("namespace", INVALID_NAMESPACES)
@pytest.mark.parametrize("method", NAMESPACE_METHODS)
def test_rejects_invalid_namespace_labels(
mocker: MockerFixture, namespace: tuple, method: str
) -> None:
store = InMemoryStore()
batch = mocker.spy(InMemoryStore, "batch")
call = {
"get": lambda: store.get(namespace, "key"),
"delete": lambda: store.delete(namespace, "key"),
"search": lambda: store.search(namespace),
"prefix": lambda: store.list_namespaces(prefix=namespace),
"suffix": lambda: store.list_namespaces(suffix=namespace),
}[method]
with pytest.raises(InvalidNamespaceError):
call()
batch.assert_not_called()
@pytest.mark.parametrize("batched", [False, True])
@pytest.mark.parametrize("namespace", INVALID_NAMESPACES)
@pytest.mark.parametrize("method", NAMESPACE_METHODS)
async def test_async_rejects_invalid_namespace_labels(
mocker: MockerFixture, batched: bool, namespace: tuple, method: str
) -> None:
# The batched store must reject before queueing: a failure inside the
# shared `abatch` would fail every op queued alongside this one.
store = MockAsyncBatchedStore() if batched else InMemoryStore()
# `MockAsyncBatchedStore` dispatches through `InMemoryStore.batch`.
batch = mocker.spy(InMemoryStore, "batch")
abatch = mocker.spy(InMemoryStore, "abatch")
call = {
"get": lambda: store.aget(namespace, "key"),
"delete": lambda: store.adelete(namespace, "key"),
"search": lambda: store.asearch(namespace),
"prefix": lambda: store.alist_namespaces(prefix=namespace),
"suffix": lambda: store.alist_namespaces(suffix=namespace),
}[method]
with pytest.raises(InvalidNamespaceError):
await call()
batch.assert_not_called()
abatch.assert_not_called()
def test_search_and_listing_keep_empty_prefixes_and_wildcards() -> None:
store = InMemoryStore()
store.put(("tenant", "a_%"), "key", {"v": 1})
store.put(("tenant", "b", "child"), "key", {"v": 1})
assert len(store.search(())) == 2
assert [item.namespace for item in store.search(("tenant", "a_%"))] == [
("tenant", "a_%")
]
assert sorted(store.list_namespaces(prefix=("tenant", "*"), suffix=("*",))) == [
("tenant", "a_%"),
("tenant", "b", "child"),
]
@pytest.mark.parametrize("batched", [False, True])
async def test_async_search_and_listing_keep_empty_prefixes_and_wildcards(
batched: bool,
) -> None:
store = MockAsyncBatchedStore() if batched else InMemoryStore()
await store.aput(("tenant", "a_%"), "key", {"v": 1})
await store.aput(("tenant", "b", "child"), "key", {"v": 1})
assert len(await store.asearch(())) == 2
assert [item.namespace for item in await store.asearch(("tenant", "a_%"))] == [
("tenant", "a_%")
]
assert sorted(
await store.alist_namespaces(prefix=("tenant", "*"), suffix=("*",))
) == [("tenant", "a_%"), ("tenant", "b", "child")]
@pytest.mark.parametrize("namespace", INVALID_NAMESPACES)
@pytest.mark.parametrize(
"kind", ["get", "put", "delete", "search", "list_prefix", "list_suffix"]
)
def test_validate_op_namespace_rejects_invalid_labels(
namespace: tuple, kind: str
) -> None:
op = {
"get": GetOp(namespace, "key"),
"put": PutOp(namespace, "key", {"v": 1}),
"delete": PutOp(namespace, "key", None),
"search": SearchOp(namespace),
"list_prefix": ListNamespacesOp((MatchCondition("prefix", namespace),)),
"list_suffix": ListNamespacesOp((MatchCondition("suffix", namespace),)),
}[kind]
with pytest.raises(InvalidNamespaceError):
validate_op_namespace(op)
def test_validate_op_namespace_allows_empty_prefix_and_wildcards() -> None:
for op in (
SearchOp(()),
ListNamespacesOp(),
ListNamespacesOp(
(
MatchCondition("prefix", ("tenant", "*")),
MatchCondition("suffix", ("*",)),
)
),
GetOp(("tenant", "a_%"), "key"),
# Write-only rules belong to `put`, not to op validation.
PutOp(("langgraph", "x"), "key", {"v": 1}),
):
validate_op_namespace(op)
async def test_async_batch_store_deduplication(mocker: MockerFixture) -> None:
abatch = mocker.spy(InMemoryStore, "batch")
store = MockAsyncBatchedStore()
+3 -65
View File
@@ -24,6 +24,9 @@ from typing_extensions import TypedDict, deprecated
from langgraph._internal._constants import CONF, CONFIG_KEY_SEND, NS_SEP
from langgraph.graph.state import StateGraph
from langgraph.pregel._messages import (
_messages_delta_reducer as _messages_delta_reducer,
)
from langgraph.warnings import LangGraphDeprecatedSinceV10
__all__ = (
@@ -244,71 +247,6 @@ def add_messages(
return merged
def _messages_delta_reducer(
state: list[AnyMessage], writes: list[list[AnyMessage]]
) -> list[AnyMessage]:
"""**Experimental.** Batch reducer for use with `DeltaChannel`.
Processes all writes in one pass — dedup by ID, `RemoveMessage`
tombstoning — without calling `add_messages`.
This reducer is batching-invariant, as required by `DeltaChannel`:
`reducer(reducer(state, xs), ys) == reducer(state, xs + ys)`.
Raw dict / string / tuple inputs are coerced to typed `BaseMessage`
objects so that HTTP-driven graphs work without a separate coercion
step. This is not full `add_messages` parity — `REMOVE_ALL_MESSAGES`,
unknown-id `RemoveMessage` errors, missing-id UUID assignment, and
`BaseMessageChunk` conversion are not handled here.
Example::
from typing import Annotated
from langgraph.channels.delta import DeltaChannel
from langgraph.graph.message import _messages_delta_reducer
class State(TypedDict):
messages: Annotated[list, DeltaChannel(_messages_delta_reducer)]
"""
# Each write is either a list of message-likes or a single message-like
# (BaseMessage / dict / str / tuple). Only lists flatten; everything
# else is one message.
flat: list[Any] = []
for w in writes:
if isinstance(w, list):
flat.extend(w)
else:
flat.append(w)
# Steady state: the reducer's own output is already typed, so skip
# `convert_to_messages` on state when the first element is a BaseMessage.
# Only raw input (initial dicts, deserialized blobs) hits the slow path.
if state and isinstance(state[0], BaseMessage):
state_msgs = state
else:
state_msgs = cast("list[AnyMessage]", convert_to_messages(state))
msgs = cast("list[AnyMessage]", convert_to_messages(flat))
index: dict[str, int] = {
m.id: i for i, m in enumerate(state_msgs) if m.id is not None
}
result: list[AnyMessage | None] = list(state_msgs)
for msg in msgs:
mid = msg.id
if mid is None:
result.append(msg)
elif isinstance(msg, RemoveMessage):
if mid in index:
result[index[mid]] = None
del index[mid]
elif mid in index:
result[index[mid]] = msg
else:
index[mid] = len(result)
result.append(msg)
return [m for m in result if m is not None]
@deprecated(
"MessageGraph is deprecated in langgraph 1.0.0, to be removed in 2.0.0. Please use StateGraph with a `messages` key instead.",
category=None,
+35 -6
View File
@@ -1,7 +1,7 @@
from __future__ import annotations
import uuid
from collections.abc import Callable, Iterable, Mapping
from collections.abc import Callable, Iterable, Mapping, Sequence
from datetime import datetime, timezone
from inspect import signature
from typing import Any, Literal, cast
@@ -32,6 +32,7 @@ from langgraph._internal._constants import (
)
from langgraph._internal._typing import MISSING
from langgraph.channels.base import BaseChannel
from langgraph.channels.binop import _get_overwrite
from langgraph.channels.delta import DeltaChannel
from langgraph.managed.base import ManagedValueMapping, ManagedValueSpec
@@ -152,6 +153,23 @@ def delta_channels_with_pending_writes(
}
def delta_channels_overwritten(
specs: Mapping[str, Any], writes: Iterable[tuple[str, Any]]
) -> set[str]:
"""Return the names of the DeltaChannels that `writes` set with an `Overwrite`.
`update_state` saves a full snapshot of these channels in the checkpoint it
creates, like the loop does when a node returns an `Overwrite`. Otherwise,
reading the channel later starts from an older snapshot and replays the
writes the `Overwrite` threw away.
"""
return {
ch
for ch, value in writes
if isinstance(specs.get(ch), DeltaChannel) and _get_overwrite(value)[0]
}
def checkpoint_superseded(
saver: BaseCheckpointSaver, config: RunnableConfig, saved: CheckpointTuple
) -> bool:
@@ -256,6 +274,7 @@ def create_checkpoint(
get_next_version: GetNextVersion | None = None,
channels_to_snapshot: set[str] | None = None,
stored_versions: ChannelVersions | None = None,
trigger_to_nodes: Mapping[str, Sequence[str]] | None = None,
) -> Checkpoint:
"""Build a new Checkpoint from the previous one and live channel state.
@@ -314,7 +333,9 @@ def create_checkpoint(
id=id or str(uuid6(clock_seq=step)),
channel_values=values,
channel_versions=channel_versions,
versions_seen=_mark_bumps_seen(checkpoint["versions_seen"], bumped),
versions_seen=_mark_bumps_seen(
checkpoint["versions_seen"], bumped, trigger_to_nodes or {}
),
updated_channels=None if updated_channels is None else sorted(updated_channels),
)
@@ -322,19 +343,27 @@ def create_checkpoint(
def _mark_bumps_seen(
versions_seen: dict[str, ChannelVersions],
bumped: Mapping[str, tuple[Any, Any]],
trigger_to_nodes: Mapping[str, Sequence[str]],
) -> dict[str, ChannelVersions]:
"""Advance whoever had seen a bumped channel's old version to the new one.
A bump that only stores a snapshot is not a write. Left unseen, it would
re-fire `interrupt_before` and rerun the channel's subscribers. For each
entry it advances, `SNAPSHOT_BUMPS` keeps the new version and the one the
node really read, so `versions_seen_without_bumps` can put the read back.
re-fire `interrupt_before` and rerun the channel's subscribers. A channel
bumped from no version was never written, so it also goes to the
subscribers that never ran: they have no entry, and would start on the bump.
For each entry it advances, `SNAPSHOT_BUMPS` keeps the new version and the
one the node really read, so `versions_seen_without_bumps` can put the read
back.
"""
if not bumped:
return versions_seen
out = dict(versions_seen)
for k, (old, _) in bumped.items():
if old is None:
for node in trigger_to_nodes.get(k, ()):
out.setdefault(node, {})
marks = dict(versions_seen.get(SNAPSHOT_BUMPS, {}))
for node, seen in versions_seen.items():
for node, seen in out.items():
if node == SNAPSHOT_BUMPS:
continue
for k, (old, new) in bumped.items():
+100 -60
View File
@@ -19,6 +19,7 @@ from typing import (
TypeVar,
cast,
)
from uuid import UUID, uuid5
from langchain_core.callbacks import AsyncParentRunManager, ParentRunManager
from langchain_core.runnables import RunnableConfig
@@ -160,6 +161,12 @@ def DuplexStream(*streams: StreamProtocol) -> StreamProtocol:
return StreamProtocol(__call__, {mode for s in streams for mode in s.modes})
def _cacheable(writes: WritesT) -> bool:
# A cache hit skips the node, so a cached error would skip it without
# raising: only a node that finished goes to the cache.
return not any(c in (INTERRUPT, ERROR) for c, _ in writes)
class PregelLoop:
config: RunnableConfig
store: BaseStore | None
@@ -191,6 +198,7 @@ class PregelLoop:
Callable[
[
concurrent.futures.Future | None,
Sequence[Any],
RunnableConfig,
Checkpoint,
str,
@@ -204,11 +212,13 @@ class PregelLoop:
submit: Submit
channels: Mapping[str, BaseChannel]
# Futures from `checkpointer.put_writes` calls that produced delta-channel
# writes. `_checkpointer_put_after_previous` drains this list (swap to a
# local `futs` then reset to `[]` and wait/gather) before putting the
# next checkpoint, so a checkpoint never becomes durable before the
# writes that produced it. Initialised to `[]` in both sync and async
# `__enter__`; stays `None` only when no checkpointer.
# writes. `_put_checkpoint` hands this list to the save it submits, which
# waits for them first, so a checkpoint never becomes durable before the
# writes that produced it. If a write or the previous save failed, the
# save fails too: a DeltaChannel is rebuilt from its writes along the
# parent chain, so a checkpoint saved past either gap reads back short
# for good. Initialised to `[]` in both sync and async `__enter__`;
# stays `None` only when no checkpointer.
_delta_write_futs: list[Any] | None = None
# Same pattern as `_delta_write_futs` but for error-handler writes.
@@ -266,6 +276,9 @@ class PregelLoop:
# `_put_exit_delta_writes` uses this to decide between anchoring on
# the existing parent (True) or creating a lazy stub (False).
_has_persisted_parent: bool = False
# True iff `__enter__` loaded the thread's latest checkpoint, not one a
# `checkpoint_id` addressed, so nothing has been built on it yet.
_loaded_latest: bool = False
managed: ManagedValueMapping
checkpoint: Checkpoint
@@ -436,13 +449,18 @@ class PregelLoop:
return None
return self._graph_lifecycle_events.popleft()
def put_writes(self, task_id: str, writes: WritesT) -> None:
def put_writes(
self, task_id: str, writes: WritesT, *, cached: bool = False
) -> None:
"""Put writes for a task, to be read by the next tick."""
if not writes:
return
# deduplicate writes to special channels, last write wins
if all(w[0] in WRITES_IDX_MAP for w in writes):
writes = list({w[0]: w for w in writes}.values())
# Before the writes are saved, kept in `checkpoint_pending_writes` or
# applied, so all of them hold the same message ids.
ensure_message_ids(self.specs, writes)
if task_id == NULL_TASK_ID:
# writes for the null task are accumulated
self.checkpoint_pending_writes = [
@@ -479,14 +497,6 @@ class PregelLoop:
# save writes
self.checkpoint_pending_writes.extend((task_id, c, v) for c, v in writes)
# Assign stable IDs to any id=None BaseMessages in DeltaChannel writes
# before the background thread serialises them. Without this, reducers
# that assign IDs inside apply_writes() race with serialisation and
# store id=None, causing get_state() replays to produce a different UUID
# on every call.
for c, v in writes_to_save:
if isinstance(self.specs.get(c), DeltaChannel):
ensure_message_ids(v)
if self.durability != "exit" and self.checkpointer_put_writes is not None:
config = patch_configurable(
self.checkpoint_config,
@@ -529,7 +539,7 @@ class PregelLoop:
self._error_handler_write_futs.append(fut)
# output writes
if hasattr(self, "tasks"):
self.output_writes(task_id, writes)
self.output_writes(task_id, writes, cached=cached)
def _put_pending_writes(self) -> None:
if self.checkpointer_put_writes is None:
@@ -678,10 +688,6 @@ class PregelLoop:
self.status = "done"
return False
if self.control is not None and self.control.drain_requested:
self.status = "draining"
return False
# if there are pending writes from a previous loop, apply them
if self._reapplies_pending_writes and self.checkpoint_pending_writes:
self._reapply_writes_to_succeeded_nodes(self.tasks)
@@ -694,6 +700,13 @@ class PregelLoop:
self.status = "interrupt_before"
raise GraphInterrupt()
# stop before running the next tasks if a drain was requested. after
# the breakpoint check: a resume passes the next step's breakpoints,
# so draining in front of one would skip it
if self.control is not None and self.control.drain_requested:
self.status = "draining"
return False
# produce debug output
self._emit("tasks", map_debug_tasks, self.tasks.values())
@@ -1005,6 +1018,8 @@ class PregelLoop:
# save writes
for tid, ws in writes.items():
if forks:
# These skip `put_writes`, which gives messages their ids.
ensure_message_ids(self.specs, ws)
carried.extend((tid, c, v) for c, v in ws)
else:
self.put_writes(tid, ws)
@@ -1069,6 +1084,9 @@ class PregelLoop:
)
# map inputs to channel updates
elif input_writes := deque(map_input(input_keys, self.input)):
# Before the input is applied, so this run reads the message ids
# that are saved with it.
ensure_message_ids(self.specs, input_writes)
# discard any unfinished tasks from previous checkpoint
discard_tasks = prepare_next_tasks(
self.checkpoint,
@@ -1114,16 +1132,26 @@ class PregelLoop:
self._exit_delta_writes.append(
(self.step, NULL_TASK_ID, "", c, v)
)
# Persist delta-channel input writes so sub-freq inputs are
# recoverable via ancestor walk (mirrors the Command input path).
# A DeltaChannel reads its input from the writes stored on the
# checkpoint this run starts from, under a task id of their own:
# readers apply a checkpoint's NULL_TASK_ID writes as its own state.
# A new thread has no checkpoint to store them on, and one a
# `checkpoint_id` addressed may have children that would read them,
# so then the input checkpoint snapshots the channel instead.
if self.durability != "exit":
delta_input = [
(c, v)
for c, v in input_writes
if isinstance(self.specs.get(c), DeltaChannel)
]
if delta_input:
self.put_writes(NULL_TASK_ID, delta_input)
if delta_input and self._has_persisted_parent and self._loaded_latest:
self.put_writes(
str(uuid5(UUID(self.checkpoint["id"]), INPUT)), delta_input
)
else:
self._delta_channels_forced_snapshot.update(
c for c, _ in delta_input
)
# save input checkpoint
self.updated_channels = updated_channels
self._put_checkpoint({"source": "input"})
@@ -1253,6 +1281,7 @@ class PregelLoop:
else None,
channels_to_snapshot=channels_to_snapshot,
stored_versions=self.checkpoint_previous_versions,
trigger_to_nodes=self.trigger_to_nodes,
)
for k in channels_to_snapshot:
new_counters[k] = (0, 0)
@@ -1299,12 +1328,17 @@ class PregelLoop:
)
self.checkpoint_previous_versions = channel_versions
# Take this checkpoint's writes now: saves run in the background
# and can start out of order, so a save that took them itself
# could get another checkpoint's writes.
delta_write_futs, self._delta_write_futs = self._delta_write_futs, []
# save it, without blocking
# if there's a previous checkpoint save in progress, wait for it
# ensuring checkpointers receive checkpoints in order
self._put_checkpoint_fut = self.submit(
self._checkpointer_put_after_previous,
getattr(self, "_put_checkpoint_fut", None),
delta_write_futs,
self.checkpoint_config,
copy_checkpoint(self.checkpoint),
self.checkpoint_metadata,
@@ -1337,6 +1371,14 @@ class PregelLoop:
):
return
# A checkpoint a `checkpoint_id` addressed may already have children,
# which read the writes stored on it, so then the channels this run
# wrote snapshot instead.
if self._has_persisted_parent and not self._loaded_latest:
self._delta_channels_forced_snapshot.update(
w[3] for w in self._exit_delta_writes
)
counters = dict(
self.checkpoint_metadata.get("counters_since_delta_snapshot") or {}
)
@@ -1374,6 +1416,7 @@ class PregelLoop:
self._put_checkpoint_fut = self.submit(
self._checkpointer_put_after_previous,
getattr(self, "_put_checkpoint_fut", None),
(),
stub_put_config,
stub_cp,
{"step": -2},
@@ -1641,21 +1684,19 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager):
def _checkpointer_put_after_previous(
self,
prev: concurrent.futures.Future | None,
delta_write_futs: Sequence[concurrent.futures.Future],
config: RunnableConfig,
checkpoint: Checkpoint,
metadata: CheckpointMetadata,
new_versions: ChannelVersions,
) -> RunnableConfig:
if self._delta_write_futs:
futs, self._delta_write_futs = self._delta_write_futs, []
concurrent.futures.wait(futs)
try:
if prev is not None:
prev.result()
finally:
cast(BaseCheckpointSaver, self.checkpointer).put(
config, checkpoint, metadata, new_versions
)
for fut in delta_write_futs:
fut.result()
if prev is not None:
prev.result()
cast(BaseCheckpointSaver, self.checkpointer).put(
config, checkpoint, metadata, new_versions
)
def match_cached_writes(self) -> Sequence[PregelExecutableTask]:
if self.cache is None:
@@ -1677,7 +1718,7 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager):
) -> PregelExecutableTask | None:
if pushed := super().accept_push(task, write_idx, call):
for task in self.match_cached_writes():
self.output_writes(task.id, task.writes, cached=True)
self.put_writes(task.id, task.writes, cached=True)
return pushed
def schedule_error_handler(
@@ -1714,16 +1755,18 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager):
if self._reapplies_pending_writes:
self._reapply_writes_to_succeeded_nodes({handler_task.id: handler_task})
for task in self.match_cached_writes():
self.output_writes(task.id, task.writes, cached=True)
self.put_writes(task.id, task.writes, cached=True)
return handler_task
def put_writes(self, task_id: str, writes: WritesT) -> None:
def put_writes(
self, task_id: str, writes: WritesT, *, cached: bool = False
) -> None:
"""Put writes for a task, to be read by the next tick."""
super().put_writes(task_id, writes)
if not writes or self.cache is None or not hasattr(self, "tasks"):
super().put_writes(task_id, writes, cached=cached)
if cached or not writes or self.cache is None or not hasattr(self, "tasks"):
return
task = self.tasks.get(task_id)
if task is None or task.cache_key is None:
if task is None or task.cache_key is None or not _cacheable(writes):
return
self.submit(
self.cache.set,
@@ -1767,6 +1810,7 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager):
# Normal case: fetch the most recent checkpoint for this
# graph/thread. Returns None on first invocation.
saved = self.checkpointer.get_tuple(self.checkpoint_config)
self._loaded_latest = True
# Capture before the synthetic-empty fallback below overwrites `saved`.
# `_put_exit_delta_writes` uses this on first run (no persisted parent)
@@ -1896,23 +1940,19 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager):
async def _checkpointer_put_after_previous(
self,
prev: asyncio.Task | None,
delta_write_futs: Sequence[asyncio.Future],
config: RunnableConfig,
checkpoint: Checkpoint,
metadata: CheckpointMetadata,
new_versions: ChannelVersions,
) -> RunnableConfig:
# Drain DeltaChannel write futures before committing the checkpoint so
# ancestor walks never see a checkpoint without its backing writes.
if self._delta_write_futs:
futs, self._delta_write_futs = self._delta_write_futs, []
await asyncio.gather(*futs)
try:
if prev is not None:
await prev
finally:
await cast(BaseCheckpointSaver, self.checkpointer).aput(
config, checkpoint, metadata, new_versions
)
if delta_write_futs:
await asyncio.gather(*delta_write_futs)
if prev is not None:
await prev
await cast(BaseCheckpointSaver, self.checkpointer).aput(
config, checkpoint, metadata, new_versions
)
async def amatch_cached_writes(self) -> Sequence[PregelExecutableTask]:
if self.cache is None:
@@ -1934,7 +1974,7 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager):
) -> PregelExecutableTask | None:
if pushed := super().accept_push(task, write_idx, call):
for task in await self.amatch_cached_writes():
self.output_writes(task.id, task.writes, cached=True)
self.put_writes(task.id, task.writes, cached=True)
return pushed
async def aschedule_error_handler(
@@ -1971,19 +2011,18 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager):
if self._reapplies_pending_writes:
self._reapply_writes_to_succeeded_nodes({handler_task.id: handler_task})
for task in await self.amatch_cached_writes():
self.output_writes(task.id, task.writes, cached=True)
self.put_writes(task.id, task.writes, cached=True)
return handler_task
def put_writes(self, task_id: str, writes: WritesT) -> None:
def put_writes(
self, task_id: str, writes: WritesT, *, cached: bool = False
) -> None:
"""Put writes for a task, to be read by the next tick."""
super().put_writes(task_id, writes)
if not writes or self.cache is None or not hasattr(self, "tasks"):
super().put_writes(task_id, writes, cached=cached)
if cached or not writes or self.cache is None or not hasattr(self, "tasks"):
return
task = self.tasks.get(task_id)
if task is None or task.cache_key is None:
return
if writes[0][0] in (INTERRUPT, ERROR):
# only cache successful tasks
if task is None or task.cache_key is None or not _cacheable(writes):
return
self.submit(
self.cache.aset,
@@ -2027,6 +2066,7 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager):
# Normal case: fetch the most recent checkpoint for this
# graph/thread. Returns None on first invocation.
saved = await self.checkpointer.aget_tuple(self.checkpoint_config)
self._loaded_latest = True
# Capture before the synthetic-empty fallback below overwrites `saved`.
# `_put_exit_delta_writes` uses this on first run (no persisted parent)
+176 -35
View File
@@ -1,6 +1,13 @@
from __future__ import annotations
from collections.abc import AsyncIterator, Callable, Iterator, Sequence
from collections.abc import (
AsyncIterator,
Callable,
Iterator,
Mapping,
MutableSequence,
Sequence,
)
from dataclasses import fields, is_dataclass
from typing import (
Any,
@@ -10,16 +17,18 @@ from typing import (
from uuid import UUID, uuid4
from langchain_core.callbacks import BaseCallbackHandler
from langchain_core.messages import BaseMessage, ToolMessage
from langchain_core.messages import AnyMessage, BaseMessage, RemoveMessage, ToolMessage
from langchain_core.messages.utils import convert_to_messages
from langchain_core.outputs import ChatGeneration, ChatGenerationChunk, LLMResult
from pydantic import BaseModel
from langgraph._internal._config import filter_to_user_tags
from langgraph._internal._constants import NS_SEP
from langgraph.channels.binop import _get_overwrite
from langgraph.channels.delta import DeltaChannel
from langgraph.constants import TAG_HIDDEN, TAG_NOSTREAM
from langgraph.pregel.protocol import StreamChunk
from langgraph.types import Command
from langgraph.types import Command, Overwrite
try:
from langchain_core.tracers._streaming import _StreamingCallbackHandler
@@ -408,6 +417,74 @@ class StreamMessagesHandlerV2(StreamMessagesHandler, _V2StreamingCallbackHandler
self.stream((meta[0], "messages", (event, v2_meta)))
# Defined here rather than in `langgraph.graph.message`, which re-exports it:
# `ensure_message_ids` has to recognize it, and importing that module from the
# loop would be an import cycle.
def _messages_delta_reducer(
state: list[AnyMessage], writes: list[list[AnyMessage]]
) -> list[AnyMessage]:
"""**Experimental.** Batch reducer for use with `DeltaChannel`.
Processes all writes in one pass — dedup by ID, `RemoveMessage`
tombstoning — without calling `add_messages`.
This reducer is batching-invariant, as required by `DeltaChannel`:
`reducer(reducer(state, xs), ys) == reducer(state, xs + ys)`.
Raw dict / string / tuple inputs are coerced to typed `BaseMessage`
objects so that HTTP-driven graphs work without a separate coercion
step. This is not full `add_messages` parity — `REMOVE_ALL_MESSAGES`,
unknown-id `RemoveMessage` errors, missing-id UUID assignment, and
`BaseMessageChunk` conversion are not handled here.
Example::
from typing import Annotated
from langgraph.channels.delta import DeltaChannel
from langgraph.graph.message import _messages_delta_reducer
class State(TypedDict):
messages: Annotated[list, DeltaChannel(_messages_delta_reducer)]
"""
# Each write is either a list of message-likes or a single message-like
# (BaseMessage / dict / str / tuple). Only lists flatten; everything
# else is one message.
flat: list[Any] = []
for w in writes:
if isinstance(w, list):
flat.extend(w)
else:
flat.append(w)
# Steady state: the reducer's own output is already typed, so skip
# `convert_to_messages` on state when the first element is a BaseMessage.
# Only raw input (initial dicts, deserialized blobs) hits the slow path.
if state and isinstance(state[0], BaseMessage):
state_msgs = state
else:
state_msgs = cast("list[AnyMessage]", convert_to_messages(state))
msgs = cast("list[AnyMessage]", convert_to_messages(flat))
index: dict[str, int] = {
m.id: i for i, m in enumerate(state_msgs) if m.id is not None
}
result: list[AnyMessage | None] = list(state_msgs)
for msg in msgs:
mid = msg.id
if mid is None:
result.append(msg)
elif isinstance(msg, RemoveMessage):
if mid in index:
result[index[mid]] = None
del index[mid]
elif mid in index:
result[index[mid]] = msg
else:
index[mid] = len(result)
result.append(msg)
return [m for m in result if m is not None]
# Known role values (OpenAI-style) and type values (LangChain serialisation)
# that identify a dict as a message. Checked before coercing to BaseMessage so
# we don't accidentally touch unrelated dicts that happen to have a "role" key.
@@ -423,39 +500,103 @@ def _is_message_dict(item: dict) -> bool:
return item.get("role") in _MESSAGE_ROLES or item.get("type") in _MESSAGE_TYPES
def ensure_message_ids(value: Any) -> None:
"""Coerce message-like write values to typed BaseMessages with stable IDs.
def ensure_message_ids(
channels: Mapping[str, Any], writes: Sequence[tuple[str, Any]]
) -> None:
"""Give the messages written to DeltaChannels stable ids.
Called in put_writes() before DeltaChannel writes are submitted to the
checkpointer. Without this the checkpoint may store raw dicts or id=None
BaseMessages; every get_state() replay then produces a different UUID and
the same message appears with a different ID in each LangSmith trace.
A `DeltaChannel` replays its saved writes on every read, and its reducer
may not assign missing ids (`_messages_delta_reducer` doesn't), so a
message saved without an id comes back without one, or with a new one on
every read, and nothing can target it by id, for example `RemoveMessage`.
Handles three input shapes:
- BaseMessage objects: assign a UUID if id is None.
- Dicts with a known "role" (OpenAI-style) or "type" (LangChain format) at
the root level: stamp "id" into the dict in-place. The reducer's
convert_to_messages call will forward the id to the resulting BaseMessage.
- Lists of the above: apply the same logic to each element, replacing dict
items with coerced BaseMessages so the shared list reference seen by
checkpoint_pending_writes and the background thread both get typed messages.
In a channel whose reducer is `_messages_delta_reducer`, every message-like
it accepts (a message, a dict, a string, a `(role, content)` tuple; alone,
in a list or in an `Overwrite`) becomes a message with an id. Other
DeltaChannels may hold other data, so there only message objects get ids,
plus message dicts outside an `Overwrite`: one in a list becomes a message.
Anything that can't become a message is left as it is, so it fails where
it always has: a message dict in a list raises here, anything else when
the reducer reads it. An empty id counts as missing.
Mutating synchronously here (before the background thread is submitted) is
safe: the serialised bytes always reflect the post-coercion state.
A write that changes is replaced in `writes`: lists and dicts the caller
passed in are never changed, only the ids of message objects.
"""
if isinstance(value, BaseMessage):
if value.id is None:
value.id = str(uuid4())
elif isinstance(value, dict) and _is_message_dict(value):
if not value.get("id"):
value["id"] = str(uuid4())
elif isinstance(value, list):
for i, item in enumerate(value):
if isinstance(item, BaseMessage):
if item.id is None:
item.id = str(uuid4())
elif isinstance(item, dict) and _is_message_dict(item):
msg = convert_to_messages([item])[0]
if msg.id is None:
msg.id = str(uuid4())
value[i] = msg
for i, (ch, value) in enumerate(writes):
spec = channels.get(ch)
if not isinstance(spec, DeltaChannel):
continue
if spec.reducer is _messages_delta_reducer:
new_value = _with_message_ids(value)
else:
new_value = _with_base_message_ids(value)
if new_value is not value:
cast(MutableSequence[tuple[str, Any]], writes)[i] = (ch, new_value)
def _with_message_ids(value: Any) -> Any:
is_overwrite, overwrite_value = _get_overwrite(value)
if not is_overwrite:
if isinstance(value, list):
return _map_messages(value, _list_item_with_id)
return _message_with_id(value)
if isinstance(overwrite_value, (list, tuple)):
items = _map_messages(overwrite_value, _message_with_id)
if items is not overwrite_value:
return Overwrite(items)
return value
def _with_base_message_ids(value: Any) -> Any:
is_overwrite, overwrite_value = _get_overwrite(value)
if is_overwrite:
if isinstance(overwrite_value, (list, tuple)):
for item in overwrite_value:
_stamp_message_id(item)
return value
if isinstance(value, list):
return _map_messages(value, _base_list_item_with_id)
if isinstance(value, dict) and _is_message_dict(value) and not value.get("id"):
return {**value, "id": str(uuid4())}
_stamp_message_id(value)
return value
def _map_messages(items: list | tuple, fn: Callable[[Any], Any]) -> list | tuple:
mapped = [fn(item) for item in items]
if all(new is old for new, old in zip(mapped, items)):
return items
return mapped if isinstance(items, list) else tuple(mapped)
def _list_item_with_id(item: Any) -> Any:
if isinstance(item, dict) and _is_message_dict(item):
return _message_with_id(item, strict=True)
return _message_with_id(item)
def _base_list_item_with_id(item: Any) -> Any:
if isinstance(item, dict) and _is_message_dict(item):
return _message_with_id(item, strict=True)
_stamp_message_id(item)
return item
def _message_with_id(item: Any, *, strict: bool = False) -> Any:
if isinstance(item, BaseMessage):
_stamp_message_id(item)
return item
try:
message = convert_to_messages([item])[0]
except Exception:
if strict:
raise
return item
if not message.id:
message.id = str(uuid4())
return message
def _stamp_message_id(item: Any) -> None:
if isinstance(item, BaseMessage) and item.id is None:
item.id = str(uuid4())
+1 -8
View File
@@ -7,7 +7,7 @@ import sys
import threading
import time
import weakref
from collections.abc import Awaitable, Callable, Sequence
from collections.abc import Callable, Sequence
from contextlib import suppress
from dataclasses import dataclass, replace
from datetime import datetime, timedelta, timezone
@@ -686,8 +686,6 @@ async def arun_with_retry(
task: PregelExecutableTask,
retry_policy: Sequence[RetryPolicy] | None,
stream: bool = False,
match_cached_writes: Callable[[], Awaitable[Sequence[PregelExecutableTask]]]
| None = None,
configurable: dict[str, Any] | None = None,
) -> None:
"""Run a task asynchronously with retries."""
@@ -711,11 +709,6 @@ async def arun_with_retry(
)
},
)
if match_cached_writes is not None and task.cache_key is not None:
for t in await match_cached_writes():
if t is task:
# if the task is already cached, return
return
while True:
runtime = config.get(CONF, {}).get(CONFIG_KEY_RUNTIME)
if isinstance(runtime, Runtime):
+161 -64
View File
@@ -137,6 +137,7 @@ from langgraph.pregel._checkpoint import (
copy_checkpoint,
create_checkpoint,
create_checkpoint_plan_for_update_state_api,
delta_channels_overwritten,
delta_channels_with_pending_writes,
empty_checkpoint,
get_updated_channels_from_tasks,
@@ -152,6 +153,7 @@ from langgraph.pregel._loop import (
from langgraph.pregel._messages import (
StreamMessagesHandler,
StreamMessagesHandlerV2,
ensure_message_ids,
)
from langgraph.pregel._read import DEFAULT_BOUND, PregelNode
from langgraph.pregel._retry import RetryPolicy
@@ -1766,6 +1768,7 @@ class Pregel(
get_next_version=checkpointer.get_next_version,
channels_to_snapshot=channels_to_snapshot,
stored_versions=checkpoint_previous_versions,
trigger_to_nodes=self.trigger_to_nodes,
)
next_config = checkpointer.put(
checkpoint_config,
@@ -1788,6 +1791,22 @@ class Pregel(
)
if input_writes := deque(map_input(self.input_channels, values)):
_store_or_fork_delta_writes(
checkpointer,
config,
saved,
checkpoint_config,
self.channels,
[
(
str(uuid5(UUID(checkpoint["id"]), INPUT)),
input_writes,
None,
)
],
fork_pending,
is_first=is_first,
)
updated_channels = apply_writes(
checkpoint,
channels,
@@ -1795,6 +1814,9 @@ class Pregel(
checkpointer.get_next_version,
self.trigger_to_nodes,
)
fork_pending |= delta_channels_overwritten(
self.channels, input_writes
)
# apply input write to channels
next_step = (
@@ -1822,6 +1844,7 @@ class Pregel(
get_next_version=checkpointer.get_next_version,
channels_to_snapshot=channels_to_snapshot,
stored_versions=checkpoint_previous_versions,
trigger_to_nodes=self.trigger_to_nodes,
)
next_config = checkpointer.put(
checkpoint_config,
@@ -1833,13 +1856,6 @@ class Pregel(
),
)
# store the writes
checkpointer.put_writes(
next_config,
input_writes,
str(uuid5(UUID(checkpoint["id"]), INPUT)),
)
return patch_checkpoint_map(
next_config, saved.metadata if saved else None
)
@@ -2033,30 +2049,22 @@ class Pregel(
),
)
updated_channels = get_updated_channels_from_tasks(run_tasks)
# The base's other children replay whatever is stored on it, so an
# edit of an older checkpoint stores none of its writes there: the
# checkpoint written here carries them, its delta channels
# snapshotted. Later supersteps address the checkpoint just written.
if (
is_first
and saved is not None
and checkpoint_superseded(checkpointer, config, saved)
):
fork_pending.update(
ch
for ch in updated_channels
if isinstance(self.channels.get(ch), DeltaChannel)
)
elif saved is not None:
for task_id, task in zip(run_task_ids, run_tasks):
channel_writes = [w for w in task.writes if w[0] != PUSH]
if channel_writes:
checkpointer.put_writes(
checkpoint_config,
channel_writes,
task_id,
**_task_path_kwarg(checkpointer.put_writes, task),
)
fork_pending |= delta_channels_overwritten(
self.channels, (w for t in run_tasks for w in t.writes)
)
_store_or_fork_delta_writes(
checkpointer,
config,
saved,
checkpoint_config,
self.channels,
[
(task_id, task.writes, task)
for task_id, task in zip(run_task_ids, run_tasks)
],
fork_pending,
is_first=is_first,
)
apply_writes(
checkpoint,
channels,
@@ -2086,6 +2094,7 @@ class Pregel(
else None,
channels_to_snapshot=channels_to_snapshot,
stored_versions=checkpoint_previous_versions,
trigger_to_nodes=self.trigger_to_nodes,
)
next_config = checkpointer.put(
checkpoint_config,
@@ -2258,6 +2267,7 @@ class Pregel(
get_next_version=checkpointer.get_next_version,
channels_to_snapshot=channels_to_snapshot,
stored_versions=checkpoint_previous_versions,
trigger_to_nodes=self.trigger_to_nodes,
)
next_config = await checkpointer.aput(
checkpoint_config,
@@ -2280,6 +2290,22 @@ class Pregel(
)
if input_writes := deque(map_input(self.input_channels, values)):
await _astore_or_fork_delta_writes(
checkpointer,
config,
saved,
checkpoint_config,
self.channels,
[
(
str(uuid5(UUID(checkpoint["id"]), INPUT)),
input_writes,
None,
)
],
fork_pending,
is_first=is_first,
)
updated_channels = apply_writes(
checkpoint,
channels,
@@ -2287,6 +2313,9 @@ class Pregel(
checkpointer.get_next_version,
self.trigger_to_nodes,
)
fork_pending |= delta_channels_overwritten(
self.channels, input_writes
)
# apply input write to channels
next_step = (
@@ -2314,6 +2343,7 @@ class Pregel(
get_next_version=checkpointer.get_next_version,
channels_to_snapshot=channels_to_snapshot,
stored_versions=checkpoint_previous_versions,
trigger_to_nodes=self.trigger_to_nodes,
)
next_config = await checkpointer.aput(
checkpoint_config,
@@ -2325,13 +2355,6 @@ class Pregel(
),
)
# store the writes
await checkpointer.aput_writes(
next_config,
input_writes,
str(uuid5(UUID(checkpoint["id"]), INPUT)),
)
return patch_checkpoint_map(
next_config, saved.metadata if saved else None
)
@@ -2524,30 +2547,22 @@ class Pregel(
),
)
updated_channels = get_updated_channels_from_tasks(run_tasks)
# The base's other children replay whatever is stored on it, so an
# edit of an older checkpoint stores none of its writes there: the
# checkpoint written here carries them, its delta channels
# snapshotted. Later supersteps address the checkpoint just written.
if (
is_first
and saved is not None
and await acheckpoint_superseded(checkpointer, config, saved)
):
fork_pending.update(
ch
for ch in updated_channels
if isinstance(self.channels.get(ch), DeltaChannel)
)
elif saved is not None:
for task_id, task in zip(run_task_ids, run_tasks):
channel_writes = [w for w in task.writes if w[0] != PUSH]
if channel_writes:
await checkpointer.aput_writes(
checkpoint_config,
channel_writes,
task_id,
**_task_path_kwarg(checkpointer.aput_writes, task),
)
fork_pending |= delta_channels_overwritten(
self.channels, (w for t in run_tasks for w in t.writes)
)
await _astore_or_fork_delta_writes(
checkpointer,
config,
saved,
checkpoint_config,
self.channels,
[
(task_id, task.writes, task)
for task_id, task in zip(run_task_ids, run_tasks)
],
fork_pending,
is_first=is_first,
)
apply_writes(
checkpoint,
channels,
@@ -2577,6 +2592,7 @@ class Pregel(
else None,
channels_to_snapshot=channels_to_snapshot,
stored_versions=checkpoint_previous_versions,
trigger_to_nodes=self.trigger_to_nodes,
)
next_config = await checkpointer.aput(
checkpoint_config,
@@ -3038,7 +3054,7 @@ class Pregel(
# with channel updates applied only at the transition between steps.
while loop.tick():
for task in loop.match_cached_writes():
loop.output_writes(task.id, task.writes, cached=True)
loop.put_writes(task.id, task.writes, cached=True)
for _ in runner.tick(
[t for t in loop.tasks.values() if not t.writes],
timeout=self.step_timeout,
@@ -3509,7 +3525,7 @@ class Pregel(
try:
while loop.tick():
for task in await loop.amatch_cached_writes():
loop.output_writes(task.id, task.writes, cached=True)
loop.put_writes(task.id, task.writes, cached=True)
async for _ in runner.atick(
[t for t in loop.tasks.values() if not t.writes],
timeout=self.step_timeout,
@@ -4255,6 +4271,87 @@ def _task_path_kwarg(put_writes: Callable[..., Any], task: PregelTaskWrites) ->
return {"task_path": task_path_str(task.path)}
_UpdateWrites = Sequence[tuple[str, Sequence[tuple[str, Any]], PregelTaskWrites | None]]
def _store_or_fork_delta_writes(
checkpointer: BaseCheckpointSaver,
config: RunnableConfig,
saved: CheckpointTuple | None,
checkpoint_config: RunnableConfig,
channels: Mapping[str, BaseChannel | ManagedValueSpec],
writes: _UpdateWrites,
fork_pending: set[str],
*,
is_first: bool,
) -> None:
"""Save an update's writes where its DeltaChannels will read them.
A DeltaChannel rebuilds its value from the writes saved on a checkpoint's
ancestors, so the writes go on the checkpoint the update builds on. If the
thread already moved past that checkpoint, its other children would read
them too, so the new checkpoint snapshots those channels instead. `writes`
holds `(task_id, writes, task)` per task; `task` is `None` for input. A
task's `PUSH` writes are left out: the update saves them on the checkpoint
it creates. The messages in `writes` get their ids first, in place, as the
loop's `put_writes` does, so pass the writes the update then applies.
"""
for _, task_writes, _ in writes:
ensure_message_ids(channels, task_writes)
if saved is None:
return
if is_first and checkpoint_superseded(checkpointer, config, saved):
fork_pending.update(
ch
for _, task_writes, _ in writes
for ch, _ in task_writes
if isinstance(channels.get(ch), DeltaChannel)
)
return
for task_id, task_writes, task in writes:
if channel_writes := [w for w in task_writes if w[0] != PUSH]:
checkpointer.put_writes(
checkpoint_config,
channel_writes,
task_id,
**(_task_path_kwarg(checkpointer.put_writes, task) if task else {}),
)
async def _astore_or_fork_delta_writes(
checkpointer: BaseCheckpointSaver,
config: RunnableConfig,
saved: CheckpointTuple | None,
checkpoint_config: RunnableConfig,
channels: Mapping[str, BaseChannel | ManagedValueSpec],
writes: _UpdateWrites,
fork_pending: set[str],
*,
is_first: bool,
) -> None:
"""Async `_store_or_fork_delta_writes`."""
for _, task_writes, _ in writes:
ensure_message_ids(channels, task_writes)
if saved is None:
return
if is_first and await acheckpoint_superseded(checkpointer, config, saved):
fork_pending.update(
ch
for _, task_writes, _ in writes
for ch, _ in task_writes
if isinstance(channels.get(ch), DeltaChannel)
)
return
for task_id, task_writes, task in writes:
if channel_writes := [w for w in task_writes if w[0] != PUSH]:
await checkpointer.aput_writes(
checkpoint_config,
channel_writes,
task_id,
**(_task_path_kwarg(checkpointer.aput_writes, task) if task else {}),
)
def _update_task_id(checkpoint_id: str, i: int) -> str:
"""Task id for the `i`th update of a superstep that has no task to reuse.
+1 -1
View File
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
[project]
name = "langgraph"
version = "1.2.14"
version = "1.2.15"
description = "Building stateful, multi-actor applications with LLMs"
authors = []
requires-python = ">=3.10"
@@ -625,3 +625,33 @@ def test_exit_resume_retried_after_its_final_checkpoint_fails_reruns_the_resumed
state = graph.get_state(config)
assert state.values["log"] == state.values["plain"]
assert sorted(state.values["plain"]) == ["ask", "done", "in"]
class _TurnState(_ResumeState):
turn: int
@pytest.mark.parametrize("exit_runs", [1, 2])
def test_exit_runs_from_a_checkpoint_the_thread_moved_past_leave_its_newer_branch_alone(
sync_checkpointer: BaseCheckpointSaver, exit_runs: int
) -> None:
builder = StateGraph(_TurnState)
builder.add_node("a", lambda state: _both(f"a{state['turn']}"))
builder.add_node("b", lambda state: _both(f"b{state['turn']}"))
builder.add_edge(START, "a")
builder.add_edge("a", "b")
graph = builder.compile(checkpointer=sync_checkpointer)
config = {"configurable": {"thread_id": "t"}}
graph.invoke({"turn": 1, **_both("in1")}, config)
end_of_turn_1 = graph.get_state(config).config
graph.invoke({"turn": 2, **_both("in2")}, config)
for turn in range(3, 3 + exit_runs):
graph.invoke(
{"turn": turn, **_both(f"in{turn}")}, end_of_turn_1, durability="exit"
)
history = list(graph.get_state_history(config))
assert [s.values.get("log", []) for s in history] == [
s.values.get("plain", []) for s in history
]
@@ -16,7 +16,11 @@ from langgraph.checkpoint.memory import InMemorySaver
from typing_extensions import TypedDict
from langgraph.channels.delta import DeltaChannel
from langgraph.channels.last_value import LastValue
from langgraph.graph import END, START, StateGraph
from langgraph.graph.message import _messages_delta_reducer
from langgraph.pregel import NodeBuilder, Pregel
from langgraph.types import Command, Durability, Overwrite
pytestmark = pytest.mark.anyio
@@ -34,18 +38,19 @@ def _append_reducer(
return result
def _build_graph(checkpointer: Any) -> Any:
def _build_graph(
checkpointer: Any,
*,
reducer: Any = _append_reducer,
write: Any = lambda: [AIMessage(content="reply", id="ai-1")],
) -> Any:
State = TypedDict( # noqa: UP013
"State",
{
"messages": Annotated[
list, DeltaChannel(_append_reducer, snapshot_frequency=50)
]
},
{"messages": Annotated[list, DeltaChannel(reducer, snapshot_frequency=50)]},
) # type: ignore[call-overload]
def agent(state: dict) -> dict:
return {"messages": [AIMessage(content="reply", id="ai-1")]}
return {"messages": write()}
return (
StateGraph(State)
@@ -138,3 +143,193 @@ def test_delta_channel_dict_style_message_gets_stable_id() -> None:
assert len(set(ids)) == 1, (
f"dict-style HumanMessage id must be stable across get_state() calls; got {ids}"
)
@pytest.mark.parametrize("container", [list, tuple])
def test_delta_channel_message_in_an_overwrite_gets_stable_id(container: type) -> None:
graph = _build_graph(InMemorySaver())
config = {"configurable": {"thread_id": "overwrite-id-stability"}}
graph.invoke(
{"messages": Overwrite(container([HumanMessage(content="hello")]))}, config
)
ids = [
next(
m.id
for m in graph.get_state(config).values["messages"]
if isinstance(m, HumanMessage)
)
for _ in range(3)
]
assert ids[0] is not None, "a message inside an Overwrite should get a UUID"
assert len(set(ids)) == 1, (
f"its id must be stable across get_state() calls; got {ids}"
)
# Writes of "hello" in every shape `_messages_delta_reducer` accepts.
MESSAGE_SHAPES = [
pytest.param(lambda: [HumanMessage("hello")], id="message"),
pytest.param(lambda: HumanMessage("hello"), id="bare message"),
pytest.param(lambda: [{"role": "user", "content": "hello"}], id="dict"),
pytest.param(lambda: {"type": "human", "content": "hello"}, id="bare dict"),
pytest.param(
lambda: [{"role": "user", "content": "hello", "id": ""}], id="dict, empty id"
),
pytest.param(lambda: [("user", "hello")], id="tuple"),
pytest.param(lambda: ("user", "hello"), id="bare tuple"),
pytest.param(lambda: ["hello"], id="string"),
pytest.param(lambda: "hello", id="bare string"),
pytest.param(lambda: Overwrite(["hello"]), id="Overwrite"),
pytest.param(lambda: Overwrite(("hello",)), id="Overwrite of a tuple"),
pytest.param(lambda: {"__overwrite__": [("user", "hello")]}, id="Overwrite dict"),
]
def _hello_ids(graph: Any, config: dict) -> list[str | None]:
return [
next(
m.id
for m in graph.get_state(config).values["messages"]
if m.content == "hello"
)
for _ in range(2)
]
@pytest.mark.parametrize("write", MESSAGE_SHAPES)
def test_input_message_of_any_shape_gets_stable_id(write: Any) -> None:
graph = _build_graph(InMemorySaver(), reducer=_messages_delta_reducer)
config = {"configurable": {"thread_id": "t"}}
graph.invoke({"messages": write()}, config)
first, second = _hello_ids(graph, config)
assert first, "the message should get a UUID"
assert first == second, "its id must be the same on every read"
@pytest.mark.parametrize("write", MESSAGE_SHAPES)
def test_node_message_of_any_shape_gets_stable_id(write: Any) -> None:
graph = _build_graph(InMemorySaver(), reducer=_messages_delta_reducer, write=write)
config = {"configurable": {"thread_id": "t"}}
graph.invoke({"messages": []}, config)
first, second = _hello_ids(graph, config)
assert first, "the message should get a UUID"
assert first == second, "its id must be the same on every read"
def test_a_list_written_twice_adds_its_messages_twice() -> None:
turn = [("user", "hello")]
graph = _build_graph(InMemorySaver(), reducer=_messages_delta_reducer)
config = {"configurable": {"thread_id": "t"}}
graph.invoke({"messages": turn}, config)
graph.invoke({"messages": turn}, config)
hellos = [
m.id for m in graph.get_state(config).values["messages"] if m.content == "hello"
]
assert len(set(hellos)) == 2, f"each write should add a message; got {hellos}"
assert turn == [("user", "hello")], "the caller's list should be left as it was"
def test_a_forked_command_update_gets_a_stable_id() -> None:
graph = _build_graph(InMemorySaver(), reducer=_messages_delta_reducer)
config = {"configurable": {"thread_id": "t"}}
graph.invoke({"messages": [HumanMessage("a", id="a")]}, config)
before_agent = next(
s.config for s in graph.get_state_history(config) if s.next == ("agent",)
)
graph.invoke(Command(update={"messages": "hello"}), before_agent)
first, second = _hello_ids(graph, config)
assert first, "the message should get a UUID"
assert first == second, "its id must be the same on every read"
def test_a_message_dict_the_reducer_rejects_fails_the_run_before_it_is_saved() -> None:
graph = _build_graph(
InMemorySaver(),
reducer=_messages_delta_reducer,
write=lambda: [{"role": "tool", "content": "x"}],
)
config = {"configurable": {"thread_id": "t"}}
with pytest.raises(KeyError):
graph.invoke({"messages": [HumanMessage("a", id="a")]}, config)
assert [m.id for m in graph.get_state(config).values["messages"]] == ["a"]
def _pregel_graph(seen: list[list[str | None]]) -> Any:
def look(state: dict) -> list:
seen.append([m.id for m in state["messages"]])
return []
node = (
NodeBuilder()
.subscribe_to("go", read=False)
.read_from("messages")
.do(look)
.write_to("messages")
)
return Pregel(
nodes={"look": node},
channels={
"messages": DeltaChannel(_messages_delta_reducer),
"go": LastValue(int),
},
input_channels=["messages", "go"],
output_channels=["messages"],
checkpointer=InMemorySaver(),
)
def test_pregel_input_message_gets_stable_id_on_a_new_thread(
durability: Durability,
) -> None:
graph = _pregel_graph([])
config = {"configurable": {"thread_id": "t"}}
graph.invoke(
{"messages": [HumanMessage("hello")], "go": 1}, config, durability=durability
)
first, second = _hello_ids(graph, config)
assert first, "the message should get a UUID"
assert first == second, "its id must be the same on every read"
def test_a_node_reads_the_pregel_input_message_ids_that_are_saved() -> None:
seen: list[list[str | None]] = []
graph = _pregel_graph(seen)
config = {"configurable": {"thread_id": "t"}}
graph.invoke({"messages": [HumanMessage("first", id="first")], "go": 1}, config)
graph.invoke({"messages": [{"role": "user", "content": "hello"}], "go": 1}, config)
saved = [m.id for m in graph.get_state(config).values["messages"]]
assert None not in saved
assert seen[-1] == saved, "the node must see the ids that are saved"
@pytest.mark.parametrize(
"value",
[
pytest.param({"type": "tool", "name": "search"}, id="not a message"),
pytest.param({"role": "user", "content": "x"}, id="message-shaped"),
],
)
def test_other_delta_channels_keep_the_dicts_in_an_overwrite(value: dict) -> None:
graph = _build_graph(InMemorySaver(), write=lambda: Overwrite([dict(value)]))
config = {"configurable": {"thread_id": "t"}}
graph.invoke({"messages": []}, config)
assert graph.get_state(config).values["messages"] == [value]
@@ -0,0 +1,122 @@
"""A node served from the node cache saves its writes like a node that ran,
without writing them back to the cache, and a node that fails isn't cached."""
import operator
from typing import Annotated, Any
import pytest
from langgraph.cache.memory import InMemoryCache
from langgraph.checkpoint.memory import InMemorySaver
from typing_extensions import TypedDict
from langgraph.channels.delta import DeltaChannel
from langgraph.graph import START, StateGraph
from langgraph.types import CachePolicy, Durability
pytestmark = pytest.mark.anyio
INPUT = {"log": [], "plain": []}
def _append(current: list, writes: list) -> list:
return [*current, *(item for write in writes for item in write)]
class _State(TypedDict):
log: Annotated[list, DeltaChannel(_append)]
plain: Annotated[list, operator.add]
class _CountsSets(InMemoryCache):
sets = 0
def set(self, keys: Any) -> None:
self.sets += 1
super().set(keys)
def _a_then_cached_b_then_c(runs: list[str], cache: InMemoryCache) -> Any:
def node(name: str) -> Any:
def run(state: _State) -> dict:
runs.append(name)
return {"log": [name], "plain": [name]}
return run
builder = StateGraph(_State)
builder.add_node("a", node("a"))
builder.add_node("b", node("b"), cache_policy=CachePolicy())
builder.add_node("c", node("c"))
builder.add_edge(START, "a")
builder.add_edge("a", "b")
builder.add_edge("b", "c")
return builder.compile(checkpointer=InMemorySaver(), cache=cache)
def test_a_cache_hit_saves_its_writes_without_caching_them_again(
durability: Durability,
) -> None:
runs: list[str] = []
cache = _CountsSets()
graph = _a_then_cached_b_then_c(runs, cache)
graph.invoke(INPUT, {"configurable": {"thread_id": "1"}}, durability=durability)
config = {"configurable": {"thread_id": "2"}}
graph.invoke(INPUT, config, durability=durability)
assert runs == ["a", "b", "c", "a", "c"]
assert cache.sets == 1, "the cache hit was written back to the cache"
for state in graph.get_state_history(config):
assert state.values.get("log", []) == state.values.get("plain", [])
async def test_a_cache_hit_saves_its_writes_without_caching_them_again_async(
durability: Durability,
) -> None:
runs: list[str] = []
cache = _CountsSets()
graph = _a_then_cached_b_then_c(runs, cache)
await graph.ainvoke(
INPUT, {"configurable": {"thread_id": "1"}}, durability=durability
)
config = {"configurable": {"thread_id": "2"}}
await graph.ainvoke(INPUT, config, durability=durability)
assert runs == ["a", "b", "c", "a", "c"]
assert cache.sets == 1, "the cache hit was written back to the cache"
async for state in graph.aget_state_history(config):
assert state.values.get("log", []) == state.values.get("plain", [])
def _cached_node_that_fails(runs: list[str]) -> Any:
def fail(state: _State) -> dict:
runs.append("b")
raise ValueError("b failed")
builder = StateGraph(_State)
builder.add_node("b", fail, cache_policy=CachePolicy())
builder.add_edge(START, "b")
return builder.compile(checkpointer=InMemorySaver(), cache=InMemoryCache())
def test_a_node_that_fails_is_not_cached() -> None:
runs: list[str] = []
graph = _cached_node_that_fails(runs)
for thread in ("1", "2"):
with pytest.raises(ValueError, match="b failed"):
graph.invoke(INPUT, {"configurable": {"thread_id": thread}})
assert runs == ["b", "b"]
async def test_a_node_that_fails_is_not_cached_async() -> None:
runs: list[str] = []
graph = _cached_node_that_fails(runs)
for thread in ("1", "2"):
with pytest.raises(ValueError, match="b failed"):
await graph.ainvoke(INPUT, {"configurable": {"thread_id": thread}})
assert runs == ["b", "b"]
@@ -0,0 +1,103 @@
"""A run's input to a DeltaChannel input channel reads back on the checkpoints
built from it, and on no others."""
import pytest
from langgraph.checkpoint.memory import InMemorySaver
from langgraph.channels.binop import BinaryOperatorAggregate
from langgraph.channels.delta import DeltaChannel
from langgraph.channels.last_value import LastValue
from langgraph.pregel import NodeBuilder, Pregel
from langgraph.types import Durability
pytestmark = pytest.mark.anyio
def _sorted_extend(current: list, writes: list) -> list:
return sorted([*current, *(item for write in writes for item in write)])
def _delta_input_graph() -> Pregel:
node = NodeBuilder().subscribe_only("go").do(lambda _: [2]).write_to("log", "plain")
return Pregel(
nodes={"n": node},
channels={
"log": DeltaChannel(_sorted_extend),
"plain": BinaryOperatorAggregate(list, lambda a, b: sorted(a + b)),
"go": LastValue(int),
},
input_channels=["log", "plain", "go"],
output_channels=["log", "plain"],
checkpointer=InMemorySaver(),
)
def test_each_run_input_reads_back_on_its_own_checkpoints(
durability: Durability,
) -> None:
graph = _delta_input_graph()
config = {"configurable": {"thread_id": "t"}}
graph.invoke({"log": [0], "plain": [0], "go": 1}, config, durability=durability)
graph.invoke({"log": [5], "plain": [5], "go": 1}, config, durability=durability)
for state in graph.get_state_history(config):
assert state.values.get("log", []) == state.values.get("plain", [])
async def test_each_run_input_reads_back_on_its_own_checkpoints_async(
durability: Durability,
) -> None:
graph = _delta_input_graph()
config = {"configurable": {"thread_id": "t"}}
await graph.ainvoke(
{"log": [0], "plain": [0], "go": 1}, config, durability=durability
)
await graph.ainvoke(
{"log": [5], "plain": [5], "go": 1}, config, durability=durability
)
async for state in graph.aget_state_history(config):
assert state.values.get("log", []) == state.values.get("plain", [])
OTHER_BRANCH_INPUTS = pytest.mark.parametrize(
"other_branch_input",
[{"go": 1}, {"log": [5], "plain": [5], "go": 1}],
ids=["other-branch-without-delta-input", "other-branch-with-delta-input"],
)
@OTHER_BRANCH_INPUTS
def test_run_input_from_an_older_checkpoint_stays_out_of_its_other_branch(
durability: Durability, other_branch_input: dict
) -> None:
graph = _delta_input_graph()
config = {"configurable": {"thread_id": "t"}}
graph.invoke({"go": 1}, config, durability=durability)
older = graph.get_state(config).config
graph.invoke(other_branch_input, config, durability=durability)
graph.invoke({"log": [7], "plain": [7], "go": 1}, older, durability=durability)
for state in graph.get_state_history(config):
assert state.values.get("log", []) == state.values.get("plain", [])
@OTHER_BRANCH_INPUTS
async def test_run_input_from_an_older_checkpoint_stays_out_of_its_other_branch_async(
durability: Durability, other_branch_input: dict
) -> None:
graph = _delta_input_graph()
config = {"configurable": {"thread_id": "t"}}
await graph.ainvoke({"go": 1}, config, durability=durability)
older = (await graph.aget_state(config)).config
await graph.ainvoke(other_branch_input, config, durability=durability)
await graph.ainvoke(
{"log": [7], "plain": [7], "go": 1}, older, durability=durability
)
async for state in graph.aget_state_history(config):
assert state.values.get("log", []) == state.values.get("plain", [])
@@ -0,0 +1,78 @@
import pytest
from langgraph.checkpoint.memory import InMemorySaver
from langgraph.channels.delta import DeltaChannel
from langgraph.channels.last_value import LastValue
from langgraph.pregel import NodeBuilder, Pregel
pytestmark = pytest.mark.anyio
CONFIG = {"configurable": {"thread_id": "t"}}
def _extend(current: list, writes: list) -> list:
return [*current, *(item for write in writes for item in write)]
def _graph(reads: list) -> Pregel:
writer = NodeBuilder().subscribe_only("a").do(lambda _: [1]).write_to("d")
reader = NodeBuilder().subscribe_only("d").do(lambda d: reads.append(list(d)))
return Pregel(
nodes={"writer": writer, "reader": reader},
channels={"a": LastValue(str), "d": DeltaChannel(_extend)},
input_channels="a",
output_channels=["d"],
checkpointer=InMemorySaver(),
)
def test_an_input_update_from_before_the_first_write_starts_only_the_writer() -> None:
reads: list = []
graph = _graph(reads)
graph.invoke("go", CONFIG)
first = next(
s.config for s in graph.get_state_history(CONFIG) if s.metadata["step"] == -1
)
fork = graph.update_state(first, {"a": "go"}, as_node="__input__")
assert graph.get_state(fork).next == ("writer",)
graph.invoke(None, fork)
assert reads == [[1], [1]]
def test_a_replay_from_before_the_first_write_forks_with_only_the_writer_next() -> None:
reads: list = []
graph = _graph(reads)
graph.invoke("go", CONFIG)
first = next(
s.config for s in graph.get_state_history(CONFIG) if s.metadata["step"] == -1
)
graph.invoke(None, first, durability="sync")
fork = next(
s for s in graph.get_state_history(CONFIG) if s.metadata["source"] == "fork"
)
assert fork.next == ("writer",)
graph.invoke(None, fork.config)
assert reads == [[1], [1], [1]]
async def test_an_ainput_update_from_before_the_first_write_starts_only_the_writer() -> (
None
):
reads: list = []
graph = _graph(reads)
await graph.ainvoke("go", CONFIG)
first = [
s.config
async for s in graph.aget_state_history(CONFIG)
if s.metadata["step"] == -1
][0]
fork = await graph.aupdate_state(first, {"a": "go"}, as_node="__input__")
assert (await graph.aget_state(fork)).next == ("writer",)
await graph.ainvoke(None, fork)
assert reads == [[1], [1]]
@@ -0,0 +1,110 @@
"""An `Overwrite` through `update_state` snapshots its DeltaChannel on the
checkpoint the update saves, as a node's `Overwrite` does on the loop's."""
from typing import Annotated, Any
import pytest
from langchain_core.messages import HumanMessage
from langgraph.checkpoint.memory import InMemorySaver
from langgraph.checkpoint.serde.types import _DeltaSnapshot
from typing_extensions import TypedDict
from langgraph.channels.delta import DeltaChannel
from langgraph.channels.last_value import LastValue
from langgraph.graph import START, StateGraph
from langgraph.graph.message import _messages_delta_reducer
from langgraph.pregel import NodeBuilder, Pregel
from langgraph.types import Overwrite
pytestmark = pytest.mark.anyio
class _State(TypedDict):
messages: Annotated[list, DeltaChannel(_messages_delta_reducer)]
def _messages_graph(saver: InMemorySaver) -> Any:
builder = StateGraph(_State)
builder.add_node("model", lambda state: {})
builder.add_edge(START, "model")
return builder.compile(checkpointer=saver)
def _extend(current: list, writes: list) -> list:
return [*current, *(item for write in writes for item in write)]
def _delta_input_graph(saver: InMemorySaver) -> Pregel:
node = NodeBuilder().subscribe_only("go").do(lambda _: [2]).write_to("log")
return Pregel(
nodes={"n": node},
channels={"log": DeltaChannel(_extend), "go": LastValue(int)},
input_channels=["log", "go"],
output_channels=["log"],
checkpointer=saver,
)
def test_update_state_with_an_overwrite_snapshots_the_channel() -> None:
saver = InMemorySaver()
graph = _messages_graph(saver)
config = {"configurable": {"thread_id": "t"}}
graph.invoke({"messages": [HumanMessage(content="a", id="1")]}, config)
graph.update_state(
config,
{"messages": Overwrite([HumanMessage(content="b", id="2")])},
as_node="model",
)
head = saver.get_tuple(config)
assert head is not None
assert isinstance(head.checkpoint["channel_values"].get("messages"), _DeltaSnapshot)
assert [m.content for m in graph.get_state(config).values["messages"]] == ["b"]
async def test_aupdate_state_with_an_overwrite_snapshots_the_channel() -> None:
saver = InMemorySaver()
graph = _messages_graph(saver)
config = {"configurable": {"thread_id": "t"}}
await graph.ainvoke({"messages": [HumanMessage(content="a", id="1")]}, config)
await graph.aupdate_state(
config,
{"messages": Overwrite([HumanMessage(content="b", id="2")])},
as_node="model",
)
head = await saver.aget_tuple(config)
assert head is not None
assert isinstance(head.checkpoint["channel_values"].get("messages"), _DeltaSnapshot)
values = (await graph.aget_state(config)).values
assert [m.content for m in values["messages"]] == ["b"]
def test_update_state_as_input_with_an_overwrite_snapshots_the_channel() -> None:
saver = InMemorySaver()
graph = _delta_input_graph(saver)
config = {"configurable": {"thread_id": "t"}}
graph.invoke({"log": [0], "go": 1}, config)
graph.update_state(config, {"log": Overwrite([1])}, as_node="__input__")
head = saver.get_tuple(config)
assert head is not None
assert isinstance(head.checkpoint["channel_values"].get("log"), _DeltaSnapshot)
assert graph.get_state(config).values["log"] == [1]
async def test_aupdate_state_as_input_with_an_overwrite_snapshots_the_channel() -> None:
saver = InMemorySaver()
graph = _delta_input_graph(saver)
config = {"configurable": {"thread_id": "t"}}
await graph.ainvoke({"log": [0], "go": 1}, config)
await graph.aupdate_state(config, {"log": Overwrite([1])}, as_node="__input__")
head = await saver.aget_tuple(config)
assert head is not None
assert isinstance(head.checkpoint["channel_values"].get("log"), _DeltaSnapshot)
assert (await graph.aget_state(config)).values["log"] == [1]
@@ -25,10 +25,13 @@ from langgraph.checkpoint.memory import InMemorySaver
from langgraph.checkpoint.serde.types import _DeltaSnapshot
from typing_extensions import TypedDict
from langgraph.channels.binop import BinaryOperatorAggregate
from langgraph.channels.delta import DeltaChannel
from langgraph.channels.last_value import LastValue
from langgraph.graph import START, StateGraph
from langgraph.graph.message import _messages_delta_reducer
from langgraph.types import StateSnapshot, StateUpdate
from langgraph.pregel import NodeBuilder, Pregel
from langgraph.types import Overwrite, StateSnapshot, StateUpdate
pytestmark = pytest.mark.anyio
@@ -538,3 +541,207 @@ def test_update_state_that_snapshots_keeps_a_deferred_node_pending() -> None:
assert [m.content for m in final["messages"]] == ["s", "a", "u", "b"]
assert graph.get_state(config).next == ()
def _sorted_extend(current: list, writes: list) -> list:
return sorted([*current, *(item for write in writes for item in write)])
def _delta_input_graph(snapshot_frequency: int = 1000) -> Any:
node = NodeBuilder().subscribe_only("go").do(lambda _: [2]).write_to("log", "plain")
return Pregel(
nodes={"n": node},
channels={
"log": DeltaChannel(_sorted_extend, snapshot_frequency=snapshot_frequency),
"plain": BinaryOperatorAggregate(list, lambda a, b: sorted(a + b)),
"go": LastValue(int),
},
input_channels=["log", "plain", "go"],
output_channels=["log", "plain"],
checkpointer=InMemorySaver(),
)
@pytest.mark.parametrize("snapshot_frequency", [1, 2])
def test_update_as_input_reads_back_on_its_checkpoint_and_after_the_next_run(
snapshot_frequency: int,
) -> None:
graph = _delta_input_graph(snapshot_frequency)
config = {"configurable": {"thread_id": "t"}}
graph.invoke({"log": [0], "plain": [0], "go": 1}, config)
graph.update_state(config, {"log": [1], "plain": [1], "go": 1}, as_node="__input__")
after_update = graph.get_state(config).values
graph.invoke(None, config)
after_run = graph.get_state(config).values
assert after_update["log"] == after_update["plain"]
assert after_run["log"] == after_run["plain"]
@pytest.mark.parametrize("snapshot_frequency", [1, 2])
async def test_aupdate_as_input_reads_back_on_its_checkpoint_and_after_the_next_run(
snapshot_frequency: int,
) -> None:
graph = _delta_input_graph(snapshot_frequency)
config = {"configurable": {"thread_id": "t"}}
await graph.ainvoke({"log": [0], "plain": [0], "go": 1}, config)
await graph.aupdate_state(
config, {"log": [1], "plain": [1], "go": 1}, as_node="__input__"
)
after_update = (await graph.aget_state(config)).values
await graph.ainvoke(None, config)
after_run = (await graph.aget_state(config)).values
assert after_update["log"] == after_update["plain"]
assert after_run["log"] == after_run["plain"]
def test_update_as_input_to_an_older_checkpoint_stays_out_of_its_other_branch() -> None:
graph = _delta_input_graph()
config = {"configurable": {"thread_id": "t"}}
graph.invoke({"go": 1}, config)
older = graph.get_state(config).config
graph.invoke({"go": 1}, config)
other_branch = graph.get_state(config)
edited = graph.update_state(
older, {"log": [1], "plain": [1], "go": 1}, as_node="__input__"
)
values = graph.get_state(edited).values
assert values["log"] == values["plain"]
assert graph.get_state(other_branch.config).values == other_branch.values
async def test_aupdate_as_input_to_an_older_checkpoint_stays_out_of_its_other_branch() -> (
None
):
graph = _delta_input_graph()
config = {"configurable": {"thread_id": "t"}}
await graph.ainvoke({"go": 1}, config)
older = (await graph.aget_state(config)).config
await graph.ainvoke({"go": 1}, config)
other_branch = await graph.aget_state(config)
edited = await graph.aupdate_state(
older, {"log": [1], "plain": [1], "go": 1}, as_node="__input__"
)
values = (await graph.aget_state(edited)).values
assert values["log"] == values["plain"]
assert (await graph.aget_state(other_branch.config)).values == other_branch.values
def _message_ids(graph: Any, config: dict) -> list[str | None]:
return [m.id for m in graph.get_state(config).values["messages"]]
async def _amessage_ids(graph: Any, config: dict) -> list[str | None]:
return [m.id for m in (await graph.aget_state(config)).values["messages"]]
def _messages_input_graph() -> Any:
node = (
NodeBuilder()
.subscribe_only("go")
.do(lambda _: [HumanMessage("n", id="n")])
.write_to("messages")
)
return Pregel(
nodes={"n": node},
channels={
"messages": DeltaChannel(_messages_delta_reducer),
"go": LastValue(int),
},
input_channels=["messages", "go"],
output_channels=["messages"],
checkpointer=InMemorySaver(),
)
# Writes of "b" in shapes `_messages_delta_reducer` accepts.
MESSAGE_SHAPES = [
pytest.param(lambda: [HumanMessage("b")], id="message"),
pytest.param(lambda: ["b"], id="string"),
pytest.param(lambda: "b", id="bare string"),
pytest.param(lambda: [("user", "b")], id="tuple"),
pytest.param(lambda: {"role": "user", "content": "b"}, id="bare dict"),
pytest.param(lambda: Overwrite([HumanMessage("b")]), id="Overwrite"),
]
@pytest.mark.parametrize("write", MESSAGE_SHAPES)
def test_update_state_gives_a_message_of_any_shape_an_id_that_every_read_keeps(
write: Any,
) -> None:
graph = _build_graph(InMemorySaver())
config = {"configurable": {"thread_id": "t"}}
graph.invoke({"messages": [HumanMessage("a", id="a")]}, config)
graph.update_state(config, {"messages": write()})
first, second = _message_ids(graph, config), _message_ids(graph, config)
assert first[-1] is not None
assert first == second
@pytest.mark.parametrize("write", MESSAGE_SHAPES)
async def test_aupdate_state_gives_a_message_of_any_shape_an_id_that_every_read_keeps(
write: Any,
) -> None:
graph = _build_graph(InMemorySaver())
config = {"configurable": {"thread_id": "t"}}
await graph.ainvoke({"messages": [HumanMessage("a", id="a")]}, config)
await graph.aupdate_state(config, {"messages": write()})
first, second = (
await _amessage_ids(graph, config),
await _amessage_ids(graph, config),
)
assert first[-1] is not None
assert first == second
@pytest.mark.parametrize("write", MESSAGE_SHAPES)
def test_update_as_input_gives_a_message_of_any_shape_an_id(write: Any) -> None:
graph = _messages_input_graph()
config = {"configurable": {"thread_id": "t"}}
graph.invoke({"messages": [HumanMessage("a", id="a")], "go": 1}, config)
graph.update_state(config, {"messages": write()}, as_node="__input__")
first, second = _message_ids(graph, config), _message_ids(graph, config)
assert first[-1] is not None
assert first == second
def test_update_state_on_an_older_checkpoint_gives_a_message_an_id() -> None:
graph = _build_graph(InMemorySaver())
config = {"configurable": {"thread_id": "t"}}
graph.invoke({"messages": [HumanMessage("a", id="a")]}, config)
older = graph.get_state(config).config
graph.invoke({"messages": [HumanMessage("c", id="c")]}, config)
branch = graph.update_state(older, {"messages": [HumanMessage("b")]})
assert _message_ids(graph, branch)[-1] is not None
async def test_aupdate_as_input_gives_a_message_an_id_that_every_read_keeps() -> None:
graph = _messages_input_graph()
config = {"configurable": {"thread_id": "t"}}
await graph.ainvoke({"messages": [HumanMessage("a", id="a")], "go": 1}, config)
await graph.aupdate_state(
config, {"messages": [HumanMessage("b")]}, as_node="__input__"
)
first, second = (
await _amessage_ids(graph, config),
await _amessage_ids(graph, config),
)
assert first[-1] is not None
assert first == second
@@ -0,0 +1,168 @@
"""A checkpoint must never be saved without the `DeltaChannel` writes it reads."""
import operator
import threading
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 START, StateGraph
from langgraph.types import Durability
pytestmark = pytest.mark.anyio
INPUT = {"log": [], "plain": []}
FINAL = {"log": ["a", "b", "c"], "plain": ["a", "b", "c"]}
# Exit mode saves nothing before the failed write, so its retry starts over.
RETRIES = [
pytest.param("sync", None, id="sync"),
pytest.param("async", None, id="async"),
pytest.param("exit", INPUT, id="exit"),
]
def _append(current: list, writes: list) -> list:
return [*current, *(item for write in writes for item in write)]
class _State(TypedDict):
log: Annotated[list, DeltaChannel(_append)]
plain: Annotated[list, operator.add]
class _FailsTheWriteOfBOnce(InMemorySaver):
failed = False
def _fail_once(self, writes: Any) -> None:
if not self.failed and ("log", ["b"]) in writes:
self.failed = True
raise ConnectionError("b's write was not saved")
def put_writes(
self, config: Any, writes: Any, task_id: str, task_path: str = ""
) -> None:
self._fail_once(writes)
super().put_writes(config, writes, task_id, task_path)
async def aput_writes(
self, config: Any, writes: Any, task_id: str, task_path: str = ""
) -> None:
self._fail_once(writes)
await super().aput_writes(config, writes, task_id, task_path)
class _FailsTheSaveAfterAOnce(InMemorySaver):
failed = False
def _fail_once(self, checkpoint: Any) -> None:
if not self.failed and checkpoint["channel_values"].get("plain") == ["a"]:
self.failed = True
raise ConnectionError("the checkpoint after a was not saved")
def put(
self, config: Any, checkpoint: Any, metadata: Any, new_versions: Any
) -> Any:
self._fail_once(checkpoint)
return super().put(config, checkpoint, metadata, new_versions)
async def aput(
self, config: Any, checkpoint: Any, metadata: Any, new_versions: Any
) -> Any:
self._fail_once(checkpoint)
return await super().aput(config, checkpoint, metadata, new_versions)
def _a_then_b_then_c(saver: InMemorySaver) -> Any:
builder = StateGraph(_State)
for name in "abc":
builder.add_node(
name, lambda state, name=name: {"log": [name], "plain": [name]}
)
builder.add_edge(START, "a")
builder.add_edge("a", "b")
builder.add_edge("b", "c")
return builder.compile(checkpointer=saver)
@pytest.mark.parametrize(("durability", "retry_input"), RETRIES)
def test_a_failed_delta_write_is_rerun_not_lost(
durability: Durability, retry_input: dict | None
) -> None:
graph = _a_then_b_then_c(_FailsTheWriteOfBOnce())
config = {"configurable": {"thread_id": "t"}}
with pytest.raises(ConnectionError):
graph.invoke(INPUT, config, durability=durability)
for state in graph.get_state_history(config):
assert state.values.get("log", []) == state.values.get("plain", [])
graph.invoke(retry_input, config, durability=durability)
assert graph.get_state(config).values == FINAL
@pytest.mark.parametrize(("durability", "retry_input"), RETRIES)
async def test_a_failed_delta_write_is_rerun_not_lost_async(
durability: Durability, retry_input: dict | None
) -> None:
graph = _a_then_b_then_c(_FailsTheWriteOfBOnce())
config = {"configurable": {"thread_id": "t"}}
with pytest.raises(ConnectionError):
await graph.ainvoke(INPUT, config, durability=durability)
async for state in graph.aget_state_history(config):
assert state.values.get("log", []) == state.values.get("plain", [])
await graph.ainvoke(retry_input, config, durability=durability)
assert (await graph.aget_state(config)).values == FINAL
@pytest.mark.parametrize("durability", ["sync", "async"])
def test_a_failed_checkpoint_save_is_rerun_not_built_on(
durability: Durability,
) -> None:
graph = _a_then_b_then_c(_FailsTheSaveAfterAOnce())
config = {"configurable": {"thread_id": "t"}}
with pytest.raises(ConnectionError):
graph.invoke(INPUT, config, durability=durability)
for state in graph.get_state_history(config):
assert state.values.get("log", []) == state.values.get("plain", [])
graph.invoke(None, config, durability=durability)
assert graph.get_state(config).values == FINAL
@pytest.mark.parametrize("durability", ["sync", "async"])
async def test_a_failed_checkpoint_save_is_rerun_not_built_on_async(
durability: Durability,
) -> None:
graph = _a_then_b_then_c(_FailsTheSaveAfterAOnce())
config = {"configurable": {"thread_id": "t"}}
with pytest.raises(ConnectionError):
await graph.ainvoke(INPUT, config, durability=durability)
async for state in graph.aget_state_history(config):
assert state.values.get("log", []) == state.values.get("plain", [])
await graph.ainvoke(None, config, durability=durability)
assert (await graph.aget_state(config)).values == FINAL
def test_a_delta_graph_finishes_on_a_single_background_thread() -> None:
graph = _a_then_b_then_c(InMemorySaver())
config = {"configurable": {"thread_id": "t"}, "max_concurrency": 1}
result: dict = {}
run = threading.Thread(
target=lambda: result.update(graph.invoke(INPUT, config, durability="async")),
daemon=True,
)
run.start()
run.join(timeout=10)
assert not run.is_alive(), "invoke hung"
assert result == FINAL
+2 -2
View File
@@ -5932,9 +5932,9 @@ def test_no_redundant_put_writes_for_cached_task(
put_writes_task_ids: list[str] = []
orig = PregelLoop.put_writes
def spy(self, task_id, writes):
def spy(self, task_id, writes, **kwargs):
put_writes_task_ids.append(task_id)
return orig(self, task_id, writes)
return orig(self, task_id, writes, **kwargs)
with patch.object(PregelLoop, "put_writes", spy):
result = workflow.invoke(Command(resume="ans"), config=config)
+2 -2
View File
@@ -8157,9 +8157,9 @@ async def test_no_redundant_put_writes_for_cached_task(
put_writes_task_ids: list[str] = []
orig = PregelLoop.put_writes
def spy(self, task_id, writes):
def spy(self, task_id, writes, **kwargs):
put_writes_task_ids.append(task_id)
return orig(self, task_id, writes)
return orig(self, task_id, writes, **kwargs)
with patch.object(PregelLoop, "put_writes", spy):
result = await workflow.ainvoke(Command(resume="ans"), config=config)
+89
View File
@@ -248,6 +248,95 @@ def test_drain_from_subgraph_can_resume_parent() -> None:
}
class _BreakpointState(TypedDict, total=False):
first: str
second: str
def _drain_before_breakpoint_builder(
second_runs: list[str],
) -> StateGraph[_BreakpointState]:
def first(state: _BreakpointState, runtime: Runtime) -> _BreakpointState:
runtime.control.request_drain("rollout")
return {"first": "done"}
def second(state: _BreakpointState) -> _BreakpointState:
second_runs.append("second")
return {"second": "done"}
builder = StateGraph(_BreakpointState)
builder.add_node("first", first)
builder.add_node("second", second)
builder.add_edge(START, "first")
builder.add_edge("first", "second")
return builder
@pytest.mark.parametrize("at_compile", [True, False])
def test_drain_before_interrupt_before_stops_at_breakpoint(at_compile: bool) -> None:
# A resume passes the next step's breakpoints, so a drain right before a
# breakpoint has to stop there instead, or the resumed run skips it.
second_runs: list[str] = []
breakpoints = {"interrupt_before": ["second"]}
compiled = _drain_before_breakpoint_builder(second_runs).compile(
checkpointer=MemorySaver(), **(breakpoints if at_compile else {})
)
run_kwargs = {} if at_compile else breakpoints
config = {"configurable": {"thread_id": "drain-breakpoint"}}
assert compiled.invoke({}, config, control=RunControl(), **run_kwargs) == {
"first": "done"
}
assert compiled.get_state(config).next == ("second",)
assert second_runs == []
assert compiled.invoke(None, config, **run_kwargs) == {
"first": "done",
"second": "done",
}
assert second_runs == ["second"]
def test_drain_before_subgraph_interrupt_before_stops_at_breakpoint() -> None:
second_runs: list[str] = []
child = _drain_before_breakpoint_builder(second_runs).compile(
interrupt_before=["second"]
)
parent = StateGraph(_BreakpointState)
parent.add_node("child", child)
parent.add_edge(START, "child")
compiled = parent.compile(checkpointer=MemorySaver())
config = {"configurable": {"thread_id": "drain-subgraph-breakpoint"}}
compiled.invoke({}, config, control=RunControl())
state = compiled.get_state(config, subgraphs=True)
assert state.next == ("child",)
assert state.tasks[0].state.next == ("second",)
assert second_runs == []
assert compiled.invoke(None, config) == {"first": "done", "second": "done"}
assert second_runs == ["second"]
@pytest.mark.anyio
async def test_drain_before_interrupt_before_stops_at_breakpoint_async() -> None:
second_runs: list[str] = []
compiled = _drain_before_breakpoint_builder(second_runs).compile(
checkpointer=MemorySaver(), interrupt_before=["second"]
)
config = {"configurable": {"thread_id": "drain-breakpoint-async"}}
assert await compiled.ainvoke({}, config, control=RunControl()) == {"first": "done"}
assert (await compiled.aget_state(config)).next == ("second",)
assert second_runs == []
assert await compiled.ainvoke(None, config) == {
"first": "done",
"second": "done",
}
assert second_runs == ["second"]
@pytest.mark.anyio
async def test_drain_requested_in_terminal_step_finishes_normally_async() -> None:
class State(TypedDict, total=False):
+3 -3
View File
@@ -1451,7 +1451,7 @@ wheels = [
[[package]]
name = "langgraph"
version = "1.2.14"
version = "1.2.15"
source = { editable = "." }
dependencies = [
{ name = "langchain-core" },
@@ -1685,7 +1685,7 @@ test = [
[[package]]
name = "langgraph-checkpoint-postgres"
version = "3.1.2"
version = "3.1.3"
source = { editable = "../checkpoint-postgres" }
dependencies = [
{ name = "langgraph-checkpoint" },
@@ -1734,7 +1734,7 @@ test = [
[[package]]
name = "langgraph-checkpoint-sqlite"
version = "3.1.1"
version = "3.1.2"
source = { editable = "../checkpoint-sqlite" }
dependencies = [
{ name = "aiosqlite" },
+3 -3
View File
@@ -286,7 +286,7 @@ wheels = [
[[package]]
name = "langgraph"
version = "1.2.14"
version = "1.2.15"
source = { editable = "../langgraph" }
dependencies = [
{ name = "langchain-core" },
@@ -418,7 +418,7 @@ test = [
[[package]]
name = "langgraph-checkpoint-postgres"
version = "3.1.2"
version = "3.1.3"
source = { editable = "../checkpoint-postgres" }
dependencies = [
{ name = "langgraph-checkpoint" },
@@ -467,7 +467,7 @@ test = [
[[package]]
name = "langgraph-checkpoint-sqlite"
version = "3.1.1"
version = "3.1.2"
source = { editable = "../checkpoint-sqlite" }
dependencies = [
{ name = "aiosqlite" },
+1 -1
View File
@@ -3,7 +3,7 @@ from langgraph_sdk.client import get_client, get_sync_client
from langgraph_sdk.encryption import Encryption
from langgraph_sdk.encryption.types import DecryptResult, EncryptionContext
__version__ = "0.4.6"
__version__ = "0.4.7"
__all__ = [
"Auth",
+1 -1
View File
@@ -301,7 +301,7 @@ wheels = [
[[package]]
name = "langgraph"
version = "1.2.14"
version = "1.2.15"
source = { editable = "../langgraph" }
dependencies = [
{ name = "langchain-core" },