Compare commits

..
Author SHA1 Message Date
raghavverma-langchainandGitHub 4ac46ba186 release(cli): 0.4.36 (#9330)
## Summary

<!-- Write the summary yourself -->

## Notes for reviewers
- Includes [push through the active Docker
context](https://github.com/langchain-ai/langgraph/pull/9321) and
[accept serverless and dedicated deployment
aliases](https://github.com/langchain-ai/langgraph/pull/9325), the two
CLI changes merged since 0.4.35.
- Merging doesn't publish by itself: after it merges, run the `release`
workflow with `working-directory: libs/cli`, as was done for 0.4.35.

## Test Plan
- [x] Only `__version__` in `libs/cli/langgraph_cli/__init__.py` changes
(0.4.35 → 0.4.36), like the 0.4.35 release PR
- [ ] After merge, the `release` workflow for `libs/cli` succeeds and
PyPI shows `langgraph-cli` 0.4.36
- [ ] `uvx --from langgraph-cli==0.4.36 langgraph deploy` pushes on
Docker Desktop without `DOCKER_HOST` set

## Demo

<!-- Drag a screen recording or screenshot here -->
2026-10-16 12:29:42 -04:00
Elior Nataf LackritzandGitHub e14f4277dc fix(langgraph): apply a resumed Command's update once (#9322)
`Command(resume=..., update=...)` applied its update twice once the
resumed node called `interrupt()` again, with sync or async durability
(`invoke`'s default). A node that asks three times, resumed with
`Command(resume=str(i), update={"log": [f"u{i}"]})`, ends with `["u1",
"u1", "u2", "u3", "a:123"]`.

## Cause

The null task's writes are saved as one list, and savers key a regular
write by its position in that list. The Command's resume came first, so
its update was stored at index 1. The next Command saved the null task's
earlier writes again without the resume, which put the update at index
0, so it was stored twice and every later resume applied both copies.
`durability="exit"` saves after the resume is consumed, so it wasn't
affected.

## Fix

The null task's special writes (a resume) now go after its regular ones,
so each regular write keeps its index as the list grows.

`test_resume_with_update_applies_the_update_once` fails on `main` for
sync and async durability.

Found by @hntrl, and flagged by @soarez on #9314.
2026-10-16 12:07:59 -04:00
hari-dhanushkodiandGitHub 4e2dfe4a24 fix(cli): accept serverless and dedicated deployment aliases (#9325) 2026-10-16 08:02:43 -04:00
raghavverma-langchainandGitHub e77a5cb08c fix(cli): push through the active Docker context in langgraph deploy (#9321)
## Summary

<!-- Write the summary yourself -->

## Notes for reviewers
- Known gap: `--host` carries only the context's endpoint, not its TLS
settings, so a remote context that uses client certificates still can't
push. That case was broken before this change too, and it's rare for
local builds.
- A context with no Docker endpoint (Docker prints `<no value>`) or a
missing `docker` binary falls back to today's behavior.

## Test Plan
- [x] `make lint` and `make test` pass in `libs/cli` (556 tests)
- [x] New tests: `login` and `push` get the active context's endpoint as
`--host`; with `DOCKER_HOST` set there's no lookup and no `--host`; when
the context lookup fails, or the context has no Docker endpoint, there's
no `--host` and the deploy still succeeds. The first fails without the
change.
- [x] On a Mac running Docker Desktop with no `/var/run/docker.sock`:
`docker --config <temp dir> version` fails with `failed to connect to
the docker API at unix:///var/run/docker.sock`, and adding `--host
<active context endpoint>` reaches the daemon
- [x] Reproduced end to end on LangSmith AWS: `langgraph deploy`
(0.4.35) failed at the push with that error, then passed with
`DOCKER_HOST` set to the context's endpoint, which this change does
automatically
- [ ] After release, `langgraph deploy` pushes on Docker Desktop without
`DOCKER_HOST` set

## Demo

<!-- Drag a screen recording or screenshot here -->
2026-10-15 16:14:01 -04:00
b0482d0627 fix(langgraph): stop a subgraph from rerunning its finished tasks on resume (#9315)
Fixes #9106

When a subgraph pauses in several parallel interrupts and they're answered one at a time, each resume ran the nodes that were already answered again. Their saved answers carried them straight through `interrupt()`, so whatever they do after it (an API write, say) ran once more per later answer, while the final state still looked right. A `@task` that finished before an interrupt in a nested graph reran the same way, which is the rerun reported in #6792.

## Cause

Every task config sets `checkpoint_id` to None, and `PregelLoop` treats a config with that key as a replay, so a nested graph replayed on every resume. A replaying loop doesn't hand saved writes back to the tasks that made them, so finished tasks ran again.

## Fix

`PregelLoop` now notes whether the run points it at a checkpoint (a truthy `checkpoint_id`, its own `checkpoint_map` entry, or a parent that is replaying), the test `__enter__` already uses to pick the checkpoint to load. `_first` then decides once whether finished tasks rerun: only on time travel, or on a resume at a checkpoint the thread has moved past. A plain resume of a subgraph is neither, so it hands finished tasks their writes back the way a top-level resume does.

The same rule fixes the top level, which @hntrl pointed out: a resume sent with the head's `checkpoint_id`, as Studio sends it, or with `checkpoint_id: None` reran finished siblings too. A resume at an older checkpoint still reruns them, as the checkpoint migration tests expect. Telling the two apart costs one more checkpoint read for a run sent with a `checkpoint_id`. `test_addressed_resume_keeps_a_rerun_tasks_write_to_a_new_channel` relied on a resume at the head rerunning a task, so it now resumes a checkpoint the thread has moved past.

Most subgraph time travel forks first, which gives tasks new ids, so saved writes can't match them anyway. `test_subgraph_replay_reasks_answered_parallel_interrupt` replays an update checkpoint, which runs in place: it passes on `main` and fails if subgraphs never replay.

The counters added to `test_parallel_interrupts_double` and the three new tests in `test_interrupt_state.py` fail on `main`.

## Not in this PR

- With `durability="exit"` and three or more parallel interrupts, a subgraph now hits a bug top-level graphs already have: a finished task's output is applied twice. #9314 fixes it.
- #6792 also reports the interrupt of a `@task` inside a nested `@entrypoint` coming back twice. That's a separate bug.

Thanks @connorimes for the report, the root cause and the fix, which this applies as proposed. The same change was opened as #9113 by @AnvitDevadiga, and @morehardy found the root cause for #6792 in #6793.

JS: langchain-ai/langgraphjs#3028

Co-authored-by: Connor Imes <5701850+connorimes@users.noreply.github.com>
2026-10-15 15:51:21 -04:00
Elior Nataf LackritzandGitHub 889a5ff16f fix(langgraph): save a resumed exit run's pending writes once (#9314)
With `durability="exit"`, resuming three or more parallel interrupts one
at a time applied the earlier answers twice. Nodes `a`, `b` and `c` each
return `{"log": [f"{name}:{interrupt(name)}"]}`, and the run ends with
`{"log": ["a:yes", "a:yes", "b:yes", "c:yes"]}`.

## Cause

An exit-mode run saves its pending writes when it ends, and that save
took all of them, including the ones loaded with the checkpoint. Savers
key a task's writes by their position in the list they're given. A task
that finished in an earlier run comes back with its old interrupt next
to its output, so saving its list again put the output under a new
index, and the next load applied both copies.

## Fix

The save at exit keeps only the writes of tasks that wrote during this
run. The rest are already stored on the checkpoint, which is what sync
and async durability do.

`test_finished_task_output_applies_once` fails on `main` in exit mode
for every saver and passes for sync and async.

The #9106 fix needs this too: once a subgraph stops rerunning its
finished tasks, it hits this bug in exit mode.

JS: langchain-ai/langgraphjs#3027
2026-10-15 15:26:01 -04:00
Baihao WangandGitHub 5eb60aef6f release(prebuilt): 1.1.1 (#9316)
Bumps langgraph-prebuilt from 1.1.0 to 1.1.1.

Includes:

- Stop `ToolNode` from swallowing a resume value that fails an
interrupt's `response_schema` (#9232). This needs langgraph 1.2.15 or
later, which marks those errors; with an older langgraph, `ToolNode`
keeps today's behavior.
- Keep tool results out of v3 messages (#7838)

Managed Deep Agents (langchain-ai/managed-deepagents-sdk#710) will raise
its `langgraph-prebuilt` minimum to this release, so an invalid answer
to an interrupt raised inside a tool, such as a subagent, fails the
resume instead of reaching the model.

After merge, run the Release workflow with
`working-directory=libs/prebuilt`.

---

This PR was written with help from an AI coding agent (Claude Code); a
human reviewed and directed the changes.

🤖 Generated with [Claude Code](https://claude.com/claude-code)
2026-10-14 18:31:09 -07:00
Elior Nataf LackritzandGitHub d4149cc1f5 release(checkpoint-postgres): 3.2.0 (#9313)
Version bump only. No library code changes in this PR.

Unreleased since `checkpointpostgres==3.1.3`:

- #9290 - add `prune` / `aprune` to `PostgresSaver` and
`AsyncPostgresSaver`

Minor version for the new public API. It needs nothing beyond the
`langgraph-checkpoint>=4.3.0` floor 3.1.3 already set.
2026-10-14 15:41:43 -04:00
Elior Nataf LackritzandGitHub 36bf3887f0 feat(checkpoint-postgres): add prune to PostgresSaver and AsyncPostgresSaver (#9290)
Closes #8531.

`PostgresSaver` and `AsyncPostgresSaver` inherited `prune` / `aprune`
from the base class, which raises `NotImplementedError`, so long-lived
threads had no supported way to drop old checkpoints. This adds both
strategies. Thanks to @wangjiawei-vegetable for proposing this scope in
#8531 and #9247.

- **`delete`** removes every checkpoint, write and blob of the threads,
in one transaction.
- **`keep_latest`** keeps the latest checkpoint of each namespace (the
one `get_tuple` loads), with its pending writes and the blobs it
references, and deletes the rest, including blobs only the deleted
checkpoints used. A kept checkpoint whose parent is deleted is left with
no parent. A subgraph call's namespace goes with the root checkpoint the
call started from; a `checkpointer=True` subgraph keeps its latest
checkpoint, since its state outlives each call.
- **Each thread is pruned on one snapshot.** Its transaction runs at
`REPEATABLE READ` and deletes everything that snapshot sees except the
kept checkpoints, so a run appending to the thread at the same time is
safe, whatever its checkpoint IDs.
- **`keep_latest` refuses threads that use a `DeltaChannel`**
(`NotImplementedError` naming them, nothing pruned). A `DeltaChannel`
rebuilds its latest value from older checkpoints, so deleting them would
empty it. The saver can't see channel types, so it detects one by what
only a `DeltaChannel` leaves: `counters_since_delta_snapshot` in
metadata, a snapshot (an inline `true` with a blob behind it), or
`delta_writes_versioned`, which LangGraph.js sets on threads its
`DeltaChannel` graphs start. The check runs for the whole batch first
and again in each thread's transaction, so a thread that starts using a
`DeltaChannel` mid-batch is refused at its turn.

Unlike the `prune` in #6883, `keep_latest` here deletes the blobs older
checkpoints no longer share with the kept ones (most of the storage on a
long thread), and guards `DeltaChannel` threads.

Measured locally, with or without planner statistics: a 6,000-checkpoint
thread prunes in 100 to 200 ms, and one with 1,500 finished subgraph
calls in about 110 ms.

## Known limitations

- Two prunes of the same thread at once can fail one of them with a
serialization failure or a deadlock; retrying is safe.
- Forking or replaying from a checkpoint while it's being pruned is not
safe.
- Threads written by older releases can escape the `DeltaChannel` check:
before #9142 (Python 1.2.13) and langchain-ai/langgraphjs#2973 (JS
1.5.0), some `update_state` paths saved checkpoints with neither
counters nor a snapshot.
- Postgres reuses the freed space but doesn't shrink the table files;
that takes `VACUUM FULL` or `pg_repack`.
- `DeltaChannel`-aware pruning (keeping ancestors back to the last
snapshot) is a follow-up.

## Test plan

- [x] `libs/checkpoint-postgres/tests/test_prune.py`: 71 tests over the
sync and async savers (single connection, pool, pipeline), including the
conformance `prune` capability. Passes on Postgres 15 and 16.
- [x] `libs/langgraph/tests/test_prune_postgres.py`: 39 graph tests. A
pruned thread loads the same state and keeps running like an unpruned
twin (sync, async and exit durability, interrupts, a failed parallel
step, interrupted and `checkpointer=True` subgraphs), plus deterministic
races: a turn saved mid-prune, and threads switching to a `DeltaChannel`
before and after their own check.
- [x] Each guard (the snapshot, the writes and blob scopes, the subgraph
call rules, the parent reset, the `DeltaChannel` checks, the
single-string check, the pre-v4 pending sends, the pipeline sync) broken
one at a time: every break fails at least one test.
- [x] A thread pruned in a loop while it runs, in each durability mode,
ends identical to its twin; threads written by LangGraph.js 1.5.0 and
pruned from Python behave like their twins.
2026-10-14 15:09:51 -04:00
19 changed files with 1728 additions and 37 deletions
@@ -27,10 +27,15 @@ from psycopg_pool import ConnectionPool
from langgraph.checkpoint.postgres import _internal
from langgraph.checkpoint.postgres.base import (
_DELTA_PAGE_SIZE,
PRUNE_DELETE_THREADS_SQL,
PRUNE_STRATEGIES,
BasePostgresSaver,
_build_delta_stage1_sql,
_build_delta_stage2_sql,
_DeltaStage2Row,
_execute,
_keep_latest,
_refuse_delta_threads,
)
from langgraph.checkpoint.postgres.shallow import ShallowPostgresSaver
@@ -401,6 +406,73 @@ class PostgresSaver(BasePostgresSaver):
(str(thread_id),),
)
def prune(
self,
thread_ids: Sequence[str],
*,
strategy: str = "keep_latest",
) -> None:
"""Prune checkpoints for the given threads.
`"delete"` removes every checkpoint, write and blob of the threads, in
one transaction.
`"keep_latest"` keeps the latest checkpoint of each namespace (the one
`get_tuple` loads, plus its parent when it predates v4 and has pending
sends), with its pending writes and the blobs it references, and
deletes the rest. A kept checkpoint whose parent it deletes is left with
no parent. A subgraph call's namespace is deleted with the root
checkpoint the call started from, except a `checkpointer=True`
subgraph's, which keeps its latest checkpoint. Each thread is pruned in
its own transaction, on one snapshot of it, so a run appending to the
thread meanwhile is safe; forking or replaying from a checkpoint being
deleted is not. A thread that fails because another transaction changed
it leaves the threads before it pruned, and retrying is safe.
`"keep_latest"` refuses threads that use a `DeltaChannel`, since it
rebuilds the latest value from older checkpoints: if any requested
thread does, nothing is pruned. A thread that starts using one before
its turn in the batch is refused then, after the threads before it.
Args:
thread_ids: The thread IDs to prune.
strategy: `"keep_latest"` or `"delete"`.
Raises:
TypeError: If `thread_ids` is a single string.
ValueError: If `strategy` is unknown.
NotImplementedError: If `strategy` is `"keep_latest"` and a thread
uses a `DeltaChannel`.
psycopg.errors.SerializationFailure: If another transaction, such
as a second prune, changes a thread's rows while `"keep_latest"`
prunes it.
psycopg.errors.DeadlockDetected: If a second prune of a thread runs
at the same time.
"""
if isinstance(thread_ids, str):
raise TypeError("thread_ids must be a sequence of thread IDs, not a str")
if strategy not in PRUNE_STRATEGIES:
raise ValueError(f"Unknown prune strategy: {strategy!r}")
thread_ids = [str(t) for t in thread_ids]
if not thread_ids:
return
if strategy == "delete":
with self._cursor() as cur, cur.connection.transaction():
for sql in PRUNE_DELETE_THREADS_SQL:
cur.execute(sql, (thread_ids,))
return
with self._cursor() as cur:
_execute(cur, _refuse_delta_threads(thread_ids))
for thread_id in thread_ids:
with self._cursor() as cur:
if self.pipe:
# The transaction has to start with its isolation level, and
# in pipeline mode it would take in what ran since the last
# sync.
self.pipe.sync()
with cur.connection.transaction():
_execute(cur, _keep_latest(thread_id))
@contextmanager
def _cursor(self, *, pipeline: bool = False) -> Iterator[Cursor[DictRow]]:
"""Create a database cursor as a context manager.
@@ -27,10 +27,15 @@ from psycopg_pool import AsyncConnectionPool
from langgraph.checkpoint.postgres import _ainternal
from langgraph.checkpoint.postgres.base import (
_DELTA_PAGE_SIZE,
PRUNE_DELETE_THREADS_SQL,
PRUNE_STRATEGIES,
BasePostgresSaver,
_aexecute,
_build_delta_stage1_sql,
_build_delta_stage2_sql,
_DeltaStage2Row,
_keep_latest,
_refuse_delta_threads,
)
from langgraph.checkpoint.postgres.shallow import AsyncShallowPostgresSaver
@@ -360,6 +365,55 @@ class AsyncPostgresSaver(BasePostgresSaver):
(str(thread_id),),
)
async def aprune(
self,
thread_ids: Sequence[str],
*,
strategy: str = "keep_latest",
) -> None:
"""Prune checkpoints for the given threads.
See `PostgresSaver.prune` for what each strategy keeps.
Args:
thread_ids: The thread IDs to prune.
strategy: `"keep_latest"` or `"delete"`.
Raises:
TypeError: If `thread_ids` is a single string.
ValueError: If `strategy` is unknown.
NotImplementedError: If `strategy` is `"keep_latest"` and a thread
uses a `DeltaChannel`.
psycopg.errors.SerializationFailure: If another transaction, such
as a second prune, changes a thread's rows while `"keep_latest"`
prunes it.
psycopg.errors.DeadlockDetected: If a second prune of a thread runs
at the same time.
"""
if isinstance(thread_ids, str):
raise TypeError("thread_ids must be a sequence of thread IDs, not a str")
if strategy not in PRUNE_STRATEGIES:
raise ValueError(f"Unknown prune strategy: {strategy!r}")
thread_ids = [str(t) for t in thread_ids]
if not thread_ids:
return
if strategy == "delete":
async with self._cursor() as cur, cur.connection.transaction():
for sql in PRUNE_DELETE_THREADS_SQL:
await cur.execute(sql, (thread_ids,))
return
async with self._cursor() as cur:
await _aexecute(cur, _refuse_delta_threads(thread_ids))
for thread_id in thread_ids:
async with self._cursor() as cur:
if self.pipe:
# The transaction has to start with its isolation level, and
# in pipeline mode it would take in what ran since the last
# sync.
await self.pipe.sync()
async with cur.connection.transaction():
await _aexecute(cur, _keep_latest(thread_id))
@asynccontextmanager
async def _cursor(
self, *, pipeline: bool = False
@@ -688,6 +742,33 @@ class AsyncPostgresSaver(BasePostgresSaver):
self.adelete_thread(thread_id), self.loop
).result()
def prune(
self,
thread_ids: Sequence[str],
*,
strategy: str = "keep_latest",
) -> None:
"""Prune checkpoints for the given threads.
See `PostgresSaver.prune` for what each strategy keeps.
Args:
thread_ids: The thread IDs to prune.
strategy: `"keep_latest"` or `"delete"`.
"""
try:
if asyncio.get_running_loop() is self.loop:
raise asyncio.InvalidStateError(
"Synchronous calls to AsyncPostgresSaver are only allowed from a "
"different thread. From the main thread, use the async interface. "
"For example, use `await checkpointer.aprune(...)`."
)
except RuntimeError:
pass
return asyncio.run_coroutine_threadsafe(
self.aprune(thread_ids, strategy=strategy), self.loop
).result()
def get_delta_channel_history(
self, *, config: RunnableConfig, channels: Sequence[str]
) -> Mapping[str, DeltaChannelHistory]:
@@ -2,7 +2,7 @@ from __future__ import annotations
import random
import warnings
from collections.abc import Mapping, Sequence
from collections.abc import Generator, Mapping, Sequence
from importlib.metadata import version as get_version
from typing import Any, TypedDict, cast
@@ -17,6 +17,8 @@ from langgraph.checkpoint.base import (
writes_sort_key,
)
from langgraph.checkpoint.serde.types import TASKS
from psycopg import AsyncCursor, Cursor
from psycopg.rows import DictRow
from psycopg.types.json import Jsonb
# Page size for stage-1 paged scan in `get_delta_channel_history`. Internal
@@ -160,6 +162,218 @@ INSERT_CHECKPOINT_WRITES_SQL = """
"""
PRUNE_STRATEGIES = ("keep_latest", "delete")
PRUNE_DELETE_THREADS_SQL = (
"DELETE FROM checkpoints WHERE thread_id = ANY(%s)",
"DELETE FROM checkpoint_blobs WHERE thread_id = ANY(%s)",
"DELETE FROM checkpoint_writes WHERE thread_id = ANY(%s)",
)
# The prune statements below never join two sets of a thread's rows: a join is
# planned from per-thread row estimates, which are worst for the few huge
# threads worth pruning, and can go quadratic. They look rows up by primary key
# instead (`LATERAL ... LIMIT 1`, `= ANY`).
# A `DeltaChannel` leaves update counters in checkpoint metadata, or a
# snapshot: an inline `true` marker with a blob behind it, which a plain `True`
# value never has. LangGraph.js inlines no values, but marks the threads its
# DeltaChannel graphs start with `delta_writes_versioned`.
PRUNE_DELTA_THREADS_SQL = """
SELECT thread_id FROM checkpoints
WHERE thread_id = ANY(%s) AND (
metadata ? 'counters_since_delta_snapshot'
OR metadata ? 'delta_writes_versioned'
)
UNION
SELECT c.thread_id
FROM checkpoints c
CROSS JOIN LATERAL jsonb_each(c.checkpoint -> 'channel_values') AS v
CROSS JOIN LATERAL (
SELECT 1 FROM checkpoint_blobs b
WHERE b.thread_id = c.thread_id
AND b.checkpoint_ns = c.checkpoint_ns
AND b.channel = v.key
AND b.version = c.checkpoint -> 'channel_versions' ->> v.key
LIMIT 1
) AS snapshot
WHERE c.thread_id = ANY(%s) AND v.value = 'true'::jsonb
"""
# Per namespace: the latest checkpoint, the parent it reads its pending sends
# from when it predates v4, and the root graph's checkpoint its subgraph call
# started from, with whether that one is saved.
PRUNE_HEADS_SQL = f"""
SELECT
h.checkpoint_ns,
h.checkpoint_id AS head,
CASE WHEN (c.checkpoint ->> 'v')::int < 4 AND EXISTS (
SELECT 1 FROM checkpoint_writes w
WHERE w.thread_id = %s
AND w.checkpoint_ns = h.checkpoint_ns
AND w.checkpoint_id = c.parent_checkpoint_id
AND w.channel = '{TASKS}'
) THEN c.parent_checkpoint_id END AS sends_parent,
c.metadata -> 'parents' ->> '' AS root,
EXISTS (
SELECT 1 FROM checkpoints r
WHERE r.thread_id = %s AND r.checkpoint_ns = ''
AND r.checkpoint_id = c.metadata -> 'parents' ->> ''
) AS root_saved
FROM (
SELECT DISTINCT ON (checkpoint_ns) checkpoint_ns, checkpoint_id
FROM checkpoints
WHERE thread_id = %s
ORDER BY checkpoint_ns DESC, checkpoint_id DESC
) AS h
CROSS JOIN LATERAL (
SELECT checkpoint, parent_checkpoint_id, metadata FROM checkpoints
WHERE thread_id = %s
AND checkpoint_ns = h.checkpoint_ns
AND checkpoint_id = h.checkpoint_id
) AS c
"""
PRUNE_SNAPSHOT_SQL = "SET TRANSACTION ISOLATION LEVEL REPEATABLE READ"
PRUNE_UNKEPT_SQL = """
SELECT checkpoint_id FROM checkpoints
WHERE thread_id = %s AND checkpoint_ns = %s AND checkpoint_id <> ALL(%s)
"""
# A run saves a superstep's writes in the background, where they can commit
# before their checkpoint does, so they go by the checkpoints deleted.
PRUNE_WRITES_SQL = """
DELETE FROM checkpoint_writes
WHERE thread_id = %s AND checkpoint_ns = %s AND checkpoint_id = ANY(%s)
"""
PRUNE_CHECKPOINTS_SQL = """
WITH deleted AS (
DELETE FROM checkpoints
WHERE thread_id = %s AND checkpoint_ns = %s AND checkpoint_id <> ALL(%s)
RETURNING checkpoint -> 'channel_versions' AS versions
)
SELECT DISTINCT v.key AS channel, v.value AS version
FROM deleted
CROSS JOIN LATERAL jsonb_each_text(deleted.versions) AS v
"""
PRUNE_PARENTS_SQL = """
UPDATE checkpoints SET parent_checkpoint_id = NULL
WHERE thread_id = %s AND checkpoint_ns = %s
AND checkpoint_id = ANY(%s) AND parent_checkpoint_id <> ALL(%s)
"""
PRUNE_REFERENCED_BLOBS_SQL = """
SELECT DISTINCT v.key AS channel, v.value AS version
FROM checkpoints c
CROSS JOIN LATERAL jsonb_each_text(c.checkpoint -> 'channel_versions') AS v
WHERE c.thread_id = %s AND c.checkpoint_ns = %s
"""
PRUNE_BLOBS_SQL = """
DELETE FROM checkpoint_blobs
WHERE thread_id = %s AND checkpoint_ns = %s AND channel = %s AND version = ANY(%s)
"""
PRUNE_SUBGRAPH_CALLS_SQL = (
"DELETE FROM checkpoints WHERE thread_id = %s AND checkpoint_ns = ANY(%s)",
"DELETE FROM checkpoint_writes WHERE thread_id = %s AND checkpoint_ns = ANY(%s)",
"DELETE FROM checkpoint_blobs WHERE thread_id = %s AND checkpoint_ns = ANY(%s)",
)
def _garbage_blobs(
freed: Sequence[Mapping[str, Any]], referenced: Sequence[Mapping[str, Any]]
) -> dict[str, list[str]]:
"""Blobs a deleted checkpoint used and no kept one does, by channel."""
kept = {(r["channel"], r["version"]) for r in referenced}
garbage: dict[str, list[str]] = {}
for r in freed:
if (r["channel"], r["version"]) not in kept:
garbage.setdefault(r["channel"], []).append(r["version"])
return garbage
_Statements = Generator[tuple[str, Sequence[Any], bool], list[DictRow], None]
def _refuse_delta_threads(thread_ids: Sequence[str]) -> _Statements:
rows = yield PRUNE_DELTA_THREADS_SQL, (thread_ids, thread_ids), True
if delta := sorted(r["thread_id"] for r in rows):
raise NotImplementedError(
f"keep_latest can't prune threads that use a DeltaChannel: {', '.join(delta)}"
)
def _keep_latest(thread_id: str) -> _Statements:
"""The statements of one thread's keep_latest transaction, for `_execute`
and `_aexecute` to run. Each comes with whether its rows are sent back.
They all read one snapshot, so what a run saves meanwhile is neither seen
nor deleted, whatever its checkpoint ID.
"""
yield PRUNE_SNAPSHOT_SQL, (), False
heads = yield PRUNE_HEADS_SQL, (thread_id,) * 4, True
# Catches a thread that started using a DeltaChannel after the batch check.
yield from _refuse_delta_threads([thread_id])
kept = {
h["checkpoint_ns"]: [i for i in (h["head"], h["sends_parent"]) if i]
for h in heads
}
# A subgraph call's namespace has its task ID after a ":", unlike a
# `checkpointer=True` subgraph's, whose state outlives each call. A call
# goes with the root checkpoint it started from, but only once that one is
# saved: runs save in the background, so a call can commit first.
calls = {
h["checkpoint_ns"]
for h in heads
if ":" in h["checkpoint_ns"]
and h["root_saved"]
and h["root"] not in kept.get("", ())
}
if calls:
for sql in PRUNE_SUBGRAPH_CALLS_SQL:
yield sql, (thread_id, list(calls)), False
for ns, ids in kept.items():
if ns in calls:
continue
unkept = yield PRUNE_UNKEPT_SQL, (thread_id, ns, ids), True
if unkept:
deleted = [r["checkpoint_id"] for r in unkept]
yield PRUNE_WRITES_SQL, (thread_id, ns, deleted), False
freed = yield PRUNE_CHECKPOINTS_SQL, (thread_id, ns, ids), True
yield PRUNE_PARENTS_SQL, (thread_id, ns, ids, ids), False
if freed:
referenced = yield PRUNE_REFERENCED_BLOBS_SQL, (thread_id, ns), True
for channel, versions in _garbage_blobs(freed, referenced).items():
yield PRUNE_BLOBS_SQL, (thread_id, ns, channel, versions), False
# Unprepared, so Postgres plans each statement knowing its arrays: a cached
# generic plan checks a row against an `= ANY(...)` array one element at a time.
def _execute(cur: Cursor[DictRow], statements: _Statements) -> None:
try:
sql, params, fetch = next(statements)
while True:
cur.execute(sql, params, prepare=False)
sql, params, fetch = statements.send(cur.fetchall() if fetch else [])
except StopIteration:
pass
async def _aexecute(cur: AsyncCursor[DictRow], statements: _Statements) -> None:
try:
sql, params, fetch = next(statements)
while True:
await cur.execute(sql, params, prepare=False)
rows = await cur.fetchall() if fetch else []
sql, params, fetch = statements.send(rows)
except StopIteration:
pass
class _DeltaStage2Row(TypedDict, total=False):
"""One row from `_build_delta_stage2_sql` (a UNION ALL of writes and blobs)."""
+1 -1
View File
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
[project]
name = "langgraph-checkpoint-postgres"
version = "3.1.3"
version = "3.2.0"
description = "Library with a Postgres implementation of LangGraph checkpoint saver."
authors = []
requires-python = ">=3.10"
@@ -0,0 +1,540 @@
import asyncio
from collections.abc import Awaitable, Callable
from contextlib import asynccontextmanager
from typing import Any
from uuid import uuid4
import pytest
from langgraph.checkpoint.base import empty_checkpoint
from langgraph.checkpoint.base.id import uuid6
from langgraph.checkpoint.conformance import validate
from langgraph.checkpoint.conformance.initializer import checkpointer_test
from langgraph.checkpoint.serde.types import TASKS, _DeltaSnapshot
from psycopg import AsyncConnection
from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
from langgraph.checkpoint.postgres.base import (
PRUNE_DELTA_THREADS_SQL,
PRUNE_HEADS_SQL,
PRUNE_REFERENCED_BLOBS_SQL,
)
from tests.conftest import DEFAULT_POSTGRES_URI, DEFAULT_URI
from tests.test_async import _saver as _async_saver
from tests.test_sync import _saver as _sync_saver
SAVERS = ["base", "pool", "pipe", "aio_base", "aio_pool", "aio_pipe"]
@asynccontextmanager
async def _saver(name: str):
if name.startswith("aio_"):
async with _async_saver(name.removeprefix("aio_")) as saver:
yield saver
else:
with _sync_saver(name) as saver:
yield saver
async def _call(saver: Any, method: str, *args: Any, **kwargs: Any) -> Any:
if isinstance(saver, AsyncPostgresSaver):
return await getattr(saver, f"a{method}")(*args, **kwargs)
return getattr(saver, method)(*args, **kwargs)
async def _rows(saver: Any, thread_id: str) -> dict[str, list[tuple]]:
queries = {
"checkpoints": "SELECT checkpoint_ns, checkpoint_id FROM checkpoints"
" WHERE thread_id = %s ORDER BY 1, 2",
"writes": "SELECT checkpoint_ns, checkpoint_id, task_id, idx"
" FROM checkpoint_writes WHERE thread_id = %s ORDER BY 1, 2, 3, 4",
"blobs": "SELECT checkpoint_ns, channel, version FROM checkpoint_blobs"
" WHERE thread_id = %s ORDER BY 1, 2, 3",
}
rows = {}
for table, sql in queries.items():
if isinstance(saver, AsyncPostgresSaver):
async with saver._cursor() as cur:
await cur.execute(sql, (thread_id,))
rows[table] = [tuple(r.values()) for r in await cur.fetchall()]
else:
with saver._cursor() as cur:
cur.execute(sql, (thread_id,))
rows[table] = [tuple(r.values()) for r in cur.fetchall()]
return rows
async def _put(
saver: Any,
thread_id: str,
parent: dict | None,
step: int,
values: dict[str, Any],
versions: dict[str, str],
*,
ns: str = "",
metadata: dict[str, Any] | None = None,
v: int = 4,
checkpoint_id: str | None = None,
) -> dict:
config = {"configurable": {"thread_id": thread_id, "checkpoint_ns": ns}}
if parent is not None:
config["configurable"]["checkpoint_id"] = parent["configurable"][
"checkpoint_id"
]
checkpoint = empty_checkpoint()
checkpoint.update(
v=v,
id=checkpoint_id or str(uuid6(clock_seq=step)),
channel_values=values,
channel_versions=versions,
)
new_versions = {k: versions[k] for k in values}
return await _call(
saver,
"put",
config,
checkpoint,
{"source": "loop", "step": step, "parents": {}, **(metadata or {})},
new_versions,
)
async def _three_checkpoints(
saver: Any,
thread_id: str,
ns: str = "",
metadata: dict[str, Any] | None = None,
) -> dict:
parent = None
steps = [
({"a": [1], "s": ["shared"]}, {"a": "1", "s": "1"}),
({"a": [1, 2]}, {"a": "2", "s": "1"}),
({"a": [1, 2, 3]}, {"a": "3", "s": "1"}),
]
for step, (values, versions) in enumerate(steps):
parent = await _put(
saver, thread_id, parent, step, values, versions, ns=ns, metadata=metadata
)
await _call(saver, "put_writes", parent, [("a", step)], str(uuid4()))
return parent
@asynccontextmanager
async def _second_saver(saver: AsyncPostgresSaver):
assert isinstance(saver.conn, AsyncConnection)
uri = DEFAULT_POSTGRES_URI + saver.conn.info.dbname
async with AsyncPostgresSaver.from_conn_string(uri) as other:
yield other
def _run_once_after(
saver: AsyncPostgresSaver,
sql: str,
action: Callable[[], Awaitable[Any]],
*,
nth: int = 1,
) -> None:
seen = 0
cursor = saver._cursor
class Hooked:
def __init__(self, cur: Any) -> None:
self.cur = cur
def __getattr__(self, name: str) -> Any:
return getattr(self.cur, name)
async def execute(self, query: Any, params: Any = None, **kwargs: Any) -> None:
nonlocal seen
await self.cur.execute(query, params, **kwargs)
if query is sql:
seen += 1
if seen == nth:
await action()
@asynccontextmanager
async def hooked(**kwargs: Any):
async with cursor(**kwargs) as cur:
yield Hooked(cur)
saver._cursor = hooked
async def _delta_thread(saver: Any, thread_id: str, evidence: str) -> None:
if evidence == "counters":
first = await _put(saver, thread_id, None, 0, {"m": [1]}, {"m": "1"})
await _put(
saver,
thread_id,
first,
1,
{},
{"m": "2"},
metadata={"counters_since_delta_snapshot": {"m": [1, 1]}},
)
elif evidence == "snapshot":
first = await _put(
saver, thread_id, None, 0, {"m": _DeltaSnapshot([1])}, {"m": "1"}
)
await _put(saver, thread_id, first, 1, {}, {"m": "2"})
else:
marker = {"delta_writes_versioned": True}
first = await _put(
saver, thread_id, None, 0, {"m": [1]}, {"m": "1"}, metadata=marker
)
await _put(
saver, thread_id, first, 1, {"m": [1, 2]}, {"m": "2"}, metadata=marker
)
@pytest.mark.parametrize("name", SAVERS)
async def test_keep_latest_keeps_each_namespace_head_with_its_writes_and_blobs(
name: str,
) -> None:
async with _saver(name) as saver:
thread_id, other_id = str(uuid4()), str(uuid4())
heads = {
ns: await _three_checkpoints(saver, thread_id, ns) for ns in ("", "sub:1")
}
await _three_checkpoints(saver, other_id)
before = {ns: await _call(saver, "get_tuple", h) for ns, h in heads.items()}
other_before = await _rows(saver, other_id)
await _call(saver, "prune", [thread_id])
rows = await _rows(saver, thread_id)
assert rows["checkpoints"] == [
(ns, heads[ns]["configurable"]["checkpoint_id"]) for ns in ("", "sub:1")
]
assert [(ns, cid) for ns, cid, *_ in rows["writes"]] == rows["checkpoints"]
assert rows["blobs"] == [
(ns, ch, ver)
for ns in ("", "sub:1")
for ch, ver in (("a", "3"), ("s", "1"))
]
for ns, head in heads.items():
assert await _call(saver, "get_tuple", head) == before[ns]._replace(
parent_config=None
)
assert await _rows(saver, other_id) == other_before
@pytest.mark.parametrize("evidence", ["counters", "snapshot", "versioned"])
@pytest.mark.parametrize("name", SAVERS)
async def test_keep_latest_refuses_a_batch_with_a_delta_channel_thread(
name: str, evidence: str
) -> None:
async with _saver(name) as saver:
ordinary_id, delta_id = str(uuid4()), str(uuid4())
await _three_checkpoints(saver, ordinary_id)
await _delta_thread(saver, delta_id, evidence)
before = [await _rows(saver, t) for t in (ordinary_id, delta_id)]
with pytest.raises(NotImplementedError, match=delta_id):
await _call(saver, "prune", [ordinary_id, delta_id])
assert [await _rows(saver, t) for t in (ordinary_id, delta_id)] == before
@pytest.mark.parametrize("name", SAVERS)
async def test_keep_latest_prunes_a_thread_storing_a_plain_true(name: str) -> None:
async with _saver(name) as saver:
thread_id = str(uuid4())
first = await _put(
saver, thread_id, None, 0, {"flag": True, "a": [1]}, {"flag": "1", "a": "1"}
)
head = await _put(
saver, thread_id, first, 1, {"flag": True}, {"flag": "2", "a": "1"}
)
await _call(saver, "prune", [thread_id])
assert (await _rows(saver, thread_id))["checkpoints"] == [
("", head["configurable"]["checkpoint_id"])
]
@pytest.mark.parametrize("name", SAVERS)
async def test_keep_latest_keeps_the_parent_holding_a_pre_v4_heads_sends(
name: str,
) -> None:
async with _saver(name) as saver:
thread_id = str(uuid4())
first = await _put(saver, thread_id, None, 0, {}, {}, v=1)
await _call(saver, "put_writes", first, [(TASKS, "send-1")], str(uuid4()))
second = await _put(saver, thread_id, first, 1, {}, {}, v=1)
head = await _put(saver, thread_id, second, 2, {}, {}, v=1)
await _call(saver, "put_writes", second, [(TASKS, "send-2")], str(uuid4()))
await _call(saver, "prune", [thread_id])
assert (await _rows(saver, thread_id))["checkpoints"] == [
("", second["configurable"]["checkpoint_id"]),
("", head["configurable"]["checkpoint_id"]),
]
loaded = await _call(saver, "get_tuple", head)
assert loaded.checkpoint["channel_values"][TASKS] == ["send-2"]
assert loaded.parent_config == second
assert (await _call(saver, "get_tuple", second)).parent_config is None
@pytest.mark.parametrize("name", SAVERS)
async def test_delete_removes_every_row_of_the_threads(name: str) -> None:
async with _saver(name) as saver:
ordinary_id, delta_id, other_id = str(uuid4()), str(uuid4()), str(uuid4())
await _three_checkpoints(saver, ordinary_id)
await _delta_thread(saver, delta_id, "counters")
await _three_checkpoints(saver, other_id)
other_before = await _rows(saver, other_id)
await _call(saver, "prune", [ordinary_id, delta_id], strategy="delete")
empty = {"checkpoints": [], "writes": [], "blobs": []}
assert await _rows(saver, ordinary_id) == empty
assert await _rows(saver, delta_id) == empty
assert await _rows(saver, other_id) == other_before
@pytest.mark.parametrize("name", SAVERS)
async def test_unknown_strategy_changes_nothing(name: str) -> None:
async with _saver(name) as saver:
thread_id = str(uuid4())
await _three_checkpoints(saver, thread_id)
before = await _rows(saver, thread_id)
with pytest.raises(ValueError, match="keep_all"):
await _call(saver, "prune", [thread_id], strategy="keep_all")
assert await _rows(saver, thread_id) == before
async def test_async_saver_prune_runs_from_another_thread_only() -> None:
async with _saver("aio_base") as saver:
thread_id = str(uuid4())
head = await _three_checkpoints(saver, thread_id)
with pytest.raises(asyncio.InvalidStateError):
saver.prune([thread_id])
await asyncio.to_thread(saver.prune, [thread_id])
assert (await _rows(saver, thread_id))["checkpoints"] == [
("", head["configurable"]["checkpoint_id"])
]
async def test_prune_conformance() -> None:
@checkpointer_test(name="AsyncPostgresSaver")
async def postgres_saver():
async with AsyncPostgresSaver.from_conn_string(DEFAULT_URI) as saver:
await saver.setup()
yield saver
report = await validate(postgres_saver, capabilities={"prune"})
prune = report.results["prune"]
if not prune.passed:
details = "\n".join(prune.failures)
pytest.fail(f"Capability prune failed (passed={prune.passed}):\n{details}")
@pytest.mark.parametrize("strategy", ["keep_latest", "delete"])
@pytest.mark.parametrize("name", ["base", "aio_base"])
async def test_prune_rejects_a_single_thread_id_string(
name: str, strategy: str
) -> None:
async with _saver(name) as saver:
for thread_id in ("ab", "a", "b"):
await _three_checkpoints(saver, thread_id)
before = [await _rows(saver, t) for t in ("ab", "a", "b")]
with pytest.raises(TypeError, match="not a str"):
await _call(saver, "prune", "ab", strategy=strategy)
assert [await _rows(saver, t) for t in ("ab", "a", "b")] == before
@pytest.mark.parametrize(
"after", [PRUNE_HEADS_SQL, PRUNE_REFERENCED_BLOBS_SQL], ids=["heads", "blobs"]
)
async def test_keep_latest_keeps_a_checkpoint_saved_while_it_runs(after: str) -> None:
async with _saver("aio_base") as saver, _second_saver(saver) as other:
thread_id = str(uuid4())
head = await _three_checkpoints(saver, thread_id)
saved: list[dict] = []
async def save_a_child() -> None:
saved.append(
await _put(
other, thread_id, head, 3, {"a": [1, 2, 3, 4]}, {"a": "4", "s": "1"}
)
)
_run_once_after(saver, after, save_a_child)
await saver.aprune([thread_id])
loaded = await saver.aget_tuple(saved[0])
assert loaded is not None, "the checkpoint saved during the prune was deleted"
assert loaded.checkpoint["channel_values"] == {
"a": [1, 2, 3, 4],
"s": ["shared"],
}
assert loaded.parent_config == head
async def test_keep_latest_refuses_a_thread_that_turns_delta_after_the_batch_check() -> (
None
):
async with _saver("aio_base") as saver, _second_saver(saver) as other:
thread_id = str(uuid4())
head = await _three_checkpoints(saver, thread_id)
async def save_a_delta_checkpoint() -> None:
await _put(
other,
thread_id,
head,
3,
{},
{"a": "4", "s": "1"},
metadata={"counters_since_delta_snapshot": {"a": [1, 1]}},
)
_run_once_after(saver, PRUNE_DELTA_THREADS_SQL, save_a_delta_checkpoint)
with pytest.raises(NotImplementedError, match=thread_id):
await saver.aprune([thread_id])
assert len((await _rows(saver, thread_id))["checkpoints"]) == 4
async def test_keep_latest_keeps_the_history_of_a_delta_checkpoint_saved_after_the_thread_check() -> (
None
):
async with _saver("aio_base") as saver, _second_saver(saver) as other:
thread_id = str(uuid4())
head = await _three_checkpoints(saver, thread_id)
saved: list[dict] = []
async def save_a_delta_checkpoint() -> None:
saved.append(
await _put(
other,
thread_id,
head,
3,
{},
{"a": "4", "s": "1"},
metadata={"counters_since_delta_snapshot": {"a": [1, 1]}},
)
)
_run_once_after(saver, PRUNE_DELTA_THREADS_SQL, save_a_delta_checkpoint, nth=2)
await saver.aprune([thread_id])
assert (await _rows(saver, thread_id))["checkpoints"] == [
("", head["configurable"]["checkpoint_id"]),
("", saved[0]["configurable"]["checkpoint_id"]),
]
history = await saver.aget_delta_channel_history(
config=saved[0], channels=["a"]
)
assert history["a"]["seed"] == [1, 2, 3]
assert [value for _, _, value in history["a"]["writes"]] == [2]
@pytest.mark.parametrize("name", SAVERS)
async def test_keep_latest_deletes_a_subgraph_call_with_the_checkpoint_it_started_from(
name: str,
) -> None:
async with _saver(name) as saver:
thread_id = str(uuid4())
older = await _put(saver, thread_id, None, 0, {}, {})
head = await _put(saver, thread_id, older, 1, {}, {})
unsaved = str(uuid6(clock_seq=2))
started_from = {
"finished:1": older["configurable"]["checkpoint_id"],
"interrupted:2": head["configurable"]["checkpoint_id"],
"running:3": unsaved,
"checkpointer_true": older["configurable"]["checkpoint_id"],
}
for ns, root in started_from.items():
await _three_checkpoints(
saver, thread_id, ns, metadata={"parents": {"": root}}
)
await _call(saver, "prune", [thread_id])
rows = await _rows(saver, thread_id)
assert {ns for ns, *_ in rows["checkpoints"]} == {
"",
"interrupted:2",
"running:3",
"checkpointer_true",
}
assert {ns for ns, *_ in rows["writes"]} == {
"interrupted:2",
"running:3",
"checkpointer_true",
}
assert {ns for ns, *_ in rows["blobs"]} == {
"interrupted:2",
"running:3",
"checkpointer_true",
}
async def test_keep_latest_keeps_a_checkpoint_saved_meanwhile_whatever_its_id() -> None:
async with _saver("aio_base") as saver, _second_saver(saver) as other:
thread_id = str(uuid4())
head = await _three_checkpoints(saver, thread_id)
lower_id = "0" + head["configurable"]["checkpoint_id"][1:]
saved: list[dict] = []
async def save_a_child_with_a_lower_id() -> None:
child = await _put(
other,
thread_id,
head,
3,
{"a": [1, 2, 3, 4]},
{"a": "4", "s": "1"},
checkpoint_id=lower_id,
)
await _call(other, "put_writes", child, [("a", 4)], str(uuid4()))
saved.append(child)
_run_once_after(saver, PRUNE_HEADS_SQL, save_a_child_with_a_lower_id)
await saver.aprune([thread_id])
loaded = await saver.aget_tuple(saved[0])
assert loaded is not None, "the checkpoint saved during the prune was deleted"
assert loaded.checkpoint["channel_values"] == {
"a": [1, 2, 3, 4],
"s": ["shared"],
}
assert [value for _, _, value in loaded.pending_writes or []] == [4]
@pytest.mark.parametrize("name", SAVERS)
async def test_keep_latest_keeps_writes_saved_before_their_checkpoint(
name: str,
) -> None:
async with _saver(name) as saver:
thread_id = str(uuid4())
head = await _three_checkpoints(saver, thread_id)
next_id = str(uuid6(clock_seq=3))
next_config = {
"configurable": {
"thread_id": thread_id,
"checkpoint_ns": "",
"checkpoint_id": next_id,
}
}
await _call(saver, "put_writes", next_config, [("a", 4)], str(uuid4()))
await _call(saver, "prune", [thread_id])
await _put(
saver, thread_id, head, 3, {}, {"a": "3", "s": "1"}, checkpoint_id=next_id
)
loaded = await _call(saver, "get_tuple", next_config)
assert [value for _, _, value in loaded.pending_writes] == [4]
+1 -1
View File
@@ -353,7 +353,7 @@ test = [
[[package]]
name = "langgraph-checkpoint-postgres"
version = "3.1.3"
version = "3.2.0"
source = { editable = "." }
dependencies = [
{ name = "langgraph-checkpoint" },
+1 -1
View File
@@ -1 +1 @@
__version__ = "0.4.35"
__version__ = "0.4.36"
+44 -12
View File
@@ -1209,6 +1209,33 @@ class DockerBuildCommand:
return cls(("docker", "buildx", "build"), flags)
def _active_docker_host(runner: CommandRunner) -> str | None:
"""Return the daemon endpoint of the user's active Docker context.
None when `DOCKER_HOST` is set, which Docker honors with any config
directory. Also None when the lookup fails (for example, `DOCKER_CONTEXT`
names a context that doesn't exist) or the context has no Docker endpoint
(Docker prints `<no value>`); that keeps Docker's own default in place.
"""
if os.environ.get("DOCKER_HOST"):
return None
try:
stdout, _ = runner.run(
subp_exec(
"docker",
"context",
"inspect",
"--format",
"{{.Endpoints.docker.Host}}",
collect=True,
)
)
except (click.exceptions.Exit, OSError):
return None
host = (stdout or "").strip()
return host if "://" in host else None
def _docker_argv(docker_config_dir: str | None) -> tuple[str, ...]:
if docker_config_dir is None:
return ("docker",)
@@ -1242,17 +1269,13 @@ def _push_image(
runner: CommandRunner,
image: str,
*,
docker_config_dir: str | None,
docker: tuple[str, ...],
verbose: bool,
) -> None:
for attempt in range(1, _PUSH_ATTEMPTS + 1):
try:
with Progress(message="Pushing...", elapsed=not verbose):
runner.run(
subp_exec(
*_docker_argv(docker_config_dir), "push", image, verbose=verbose
)
)
runner.run(subp_exec(*docker, "push", image, verbose=verbose))
return
except click.exceptions.Exit:
if attempt == _PUSH_ATTEMPTS:
@@ -1337,6 +1360,11 @@ def _run_local_build(
registry_host = repository.split("/")[0]
with _docker_config_for_token(registry_host, deployment_token, username) as cfg:
# A config directory has no current context, so Docker would fall
# back to /var/run/docker.sock; keep using the daemon the build used.
docker_host = _active_docker_host(runner)
host_args = ("--host", docker_host) if docker_host else ()
docker = (*_docker_argv(cfg), *host_args)
_log_deploy_step(step, f"Logging into {registry_host}")
token_input = (
deployment_token
@@ -1345,7 +1373,7 @@ def _run_local_build(
)
runner.run(
subp_exec(
*_docker_argv(cfg),
*docker,
"login",
"-u",
username,
@@ -1361,7 +1389,7 @@ def _run_local_build(
runner.run(
subp_exec("docker", "tag", image_to_push, remote_image, verbose=verbose)
)
_push_image(runner, remote_image, docker_config_dir=cfg, verbose=verbose)
_push_image(runner, remote_image, docker=docker, verbose=verbose)
step += 1
resolved_image = _resolve_pushed_image_digest(
@@ -1826,7 +1854,7 @@ class CustomerRegistrySource:
_build_image(runner, ctx.spec, image, verbose=ctx.verbose)
step += 1
_log_deploy_step(step, f"Pushing image {image}")
_push_image(runner, image, docker_config_dir=None, verbose=ctx.verbose)
_push_image(runner, image, docker=_docker_argv(None), verbose=ctx.verbose)
step += 1
digest = _resolve_pushed_image_digest(
runner, remote_image=image, docker_config_dir=None, verbose=ctx.verbose
@@ -2185,11 +2213,13 @@ def _deploy_base_options(
),
click.option(
"--deployment-type",
type=click.Choice(["dev", "prod"]),
type=click.Choice(["serverless", "dedicated", "dev", "prod"]),
default="dev",
show_default=True,
help=(
"Deployment type (used when creating a new deployment). "
"Deployment type (used when creating a new deployment): "
"serverless (scale to zero) or dedicated (prod); "
"legacy dev and prod are also accepted. "
"Ignored with --push-to."
),
),
@@ -2480,7 +2510,9 @@ def _deploy_cmd(
),
verbose=verbose,
selector=selector,
deployment_type=deployment_type,
deployment_type={"serverless": "dev_zero", "dedicated": "prod"}.get(
deployment_type, deployment_type
),
secrets=secrets,
tracked_packages=tracked_packages,
)
@@ -52,6 +52,8 @@ ARCHIVE = ("/tmp/src.tgz", 2048, "langgraph.json")
OBJECT_PATH = "tarballs/src.tgz"
PLATFORM_FORMAT = "{{.Os}}/{{.Architecture}}"
DIGESTS_FORMAT = "{{json .RepoDigests}}"
CONTEXT_HOST_FORMAT = "{{.Endpoints.docker.Host}}"
CONTEXT_HOST = "unix:///Users/dev/.docker/run/docker.sock"
NOT_A_CLI_DEPLOYMENT = (
"push token is only available for 'internal_docker' source deployments"
)
@@ -192,6 +194,8 @@ class DockerCommand:
class DockerDouble:
timeline: list[str]
failing_pushes: int = 0
context_lookup_fails: bool = False
context_host: str = CONTEXT_HOST
builds: list[dict] = field(default_factory=list)
commands: list[DockerCommand] = field(default_factory=list)
@@ -235,6 +239,10 @@ class DockerDouble:
if "push" in args and self.failing_pushes > 0:
self.failing_pushes -= 1
raise click.exceptions.Exit(1)
if CONTEXT_HOST_FORMAT in args:
if self.context_lookup_fails:
raise click.exceptions.Exit(1)
return f"{self.context_host}\n", None
if PLATFORM_FORMAT in args:
return "linux/amd64\n", None
if DIGESTS_FORMAT in args:
@@ -244,6 +252,8 @@ class DockerDouble:
@staticmethod
def _verb(args: tuple[str, ...]) -> str:
if CONTEXT_HOST_FORMAT in args:
return "inspect-context"
if PLATFORM_FORMAT in args:
return "inspect-platform"
if DIGESTS_FORMAT in args:
@@ -292,7 +302,12 @@ def deploy_project(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> DeployPro
json.dumps({"graphs": {"agent": "agent.py:graph"}, "dependencies": ["."]})
)
monkeypatch.chdir(tmp_path)
for name in ("LANGSMITH_TENANT_ID", "LANGSMITH_ENDPOINT", "LANGGRAPH_HOST_URL"):
for name in (
"LANGSMITH_TENANT_ID",
"LANGSMITH_ENDPOINT",
"LANGGRAPH_HOST_URL",
"DOCKER_HOST",
):
monkeypatch.delenv(name, raising=False)
timeline: list[str] = []
@@ -340,6 +355,7 @@ def test_first_local_deploy_creates_then_builds_pushes_and_updates_in_order(
CREATE_DEPLOYMENT,
"docker build",
_push_token(CREATED_ID),
"docker inspect-context",
"docker login",
"docker tag",
"docker push",
@@ -349,6 +365,27 @@ def test_first_local_deploy_creates_then_builds_pushes_and_updates_in_order(
assert "Deployment updated" in result.output
@pytest.mark.parametrize(
("deployment_type", "wire_type"),
[
("serverless", "dev_zero"),
("dedicated", "prod"),
("dev", "dev"),
("prod", "prod"),
],
)
@pytest.mark.parametrize("build_mode", ["--no-remote", "--remote"])
def test_deployment_type_aliases(
deploy_project: DeployProject, deployment_type: str, wire_type: str, build_mode: str
) -> None:
result = deploy_project.run(build_mode, "--deployment-type", deployment_type)
assert result.exit_code == 0, result.output
assert deploy_project.control_plane.bodies[CREATE_DEPLOYMENT]["source_config"] == {
"deployment_type": wire_type
}
def test_first_local_deploy_creates_an_internal_docker_deployment(
deploy_project: DeployProject,
) -> None:
@@ -406,7 +443,7 @@ def test_local_deploy_logs_in_with_the_control_plane_push_token(
login = deploy_project.docker.command("login")
assert login.args[:2] == ("docker", "--config")
assert login.args[3:] == (
assert login.args[-5:] == (
"login",
"-u",
"oauth2accesstoken",
@@ -416,6 +453,53 @@ def test_local_deploy_logs_in_with_the_control_plane_push_token(
assert login.kwargs["input"] == f"{PUSH_TOKEN}\n"
def test_local_deploy_logs_in_and_pushes_through_the_active_docker_context(
deploy_project: DeployProject,
) -> None:
"""The temporary config has no current context, so without the active
context's endpoint Docker would push through /var/run/docker.sock."""
result = deploy_project.run("--no-remote")
assert result.exit_code == 0, result.output
for verb in ("login", "push"):
args = deploy_project.docker.command(verb).args
assert args[3:5] == ("--host", CONTEXT_HOST), verb
def test_local_deploy_leaves_an_explicit_docker_host_to_docker(
deploy_project: DeployProject, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setenv("DOCKER_HOST", "tcp://127.0.0.1:2375")
result = deploy_project.run("--no-remote")
assert result.exit_code == 0, result.output
for verb in ("login", "push"):
assert "--host" not in deploy_project.docker.command(verb).args, verb
def test_local_deploy_keeps_dockers_default_when_the_context_lookup_fails(
deploy_project: DeployProject,
) -> None:
deploy_project.docker.context_lookup_fails = True
result = deploy_project.run("--no-remote")
assert result.exit_code == 0, result.output
assert "--host" not in deploy_project.docker.command("push").args
def test_local_deploy_ignores_a_context_without_a_docker_endpoint(
deploy_project: DeployProject,
) -> None:
deploy_project.docker.context_host = "<no value>"
result = deploy_project.run("--no-remote")
assert result.exit_code == 0, result.output
assert "--host" not in deploy_project.docker.command("push").args
def test_local_deploy_tells_the_control_plane_it_honors_registry_usernames(
deploy_project: DeployProject,
) -> None:
@@ -438,7 +522,7 @@ def test_local_deploy_logs_in_with_the_username_the_control_plane_returns(
deploy_project.run("--no-remote")
login = deploy_project.docker.command("login")
assert login.args[3:] == (
assert login.args[-5:] == (
"login",
"-u",
"AWS",
@@ -506,6 +590,7 @@ def test_prebuilt_image_is_validated_and_pushed_without_a_build(
assert deploy_project.docker.builds == []
assert deploy_project.docker.verbs() == [
"docker inspect-platform",
"docker inspect-context",
"docker login",
"docker tag",
"docker push",
+32 -10
View File
@@ -257,8 +257,10 @@ class PregelLoop:
_resume_loaded_writes: Sequence[PendingWrite] = ()
# Tasks `_reapply_writes_to_succeeded_nodes` handed loaded writes back to.
_reapplied_task_ids: set[str]
# An update or fork checkpoint addressed by `checkpoint_id` that the thread
# has moved past, so it may already have children.
# Tasks that wrote during this run, the only ones exit mode saves at exit.
_written_task_ids: set[str]
# A checkpoint addressed by `checkpoint_id` that the thread has moved past,
# so it may already have children.
_addressed_checkpoint_superseded: bool = False
# The checkpoint_config that points at the parent loaded at `__enter__`
@@ -349,7 +351,6 @@ class PregelLoop:
self.interrupt_before = interrupt_before
self.manager = manager
self.is_nested = CONFIG_KEY_TASK_ID in self.config.get(CONF, {})
self.is_replaying = CONFIG_KEY_CHECKPOINT_ID in config[CONF]
self._migrate_checkpoint = migrate_checkpoint
self.trigger_to_nodes = trigger_to_nodes
self.retry_policy = retry_policy
@@ -407,6 +408,14 @@ class PregelLoop:
else ()
)
self.prev_checkpoint_config = None
# Whether the run points this loop at a checkpoint: a truthy
# `checkpoint_id`, this graph's `checkpoint_map` entry, or a replaying
# parent. Every task config carries `checkpoint_id: None`, so a key
# check would count every subgraph.
self._addressed = bool(
self.checkpoint_config[CONF].get(CONFIG_KEY_CHECKPOINT_ID)
or self.config[CONF].get(CONFIG_KEY_REPLAY_STATE)
)
runtime = self.config[CONF].get(CONFIG_KEY_RUNTIME)
self.control = runtime.control if isinstance(runtime, Runtime) else None
@@ -455,6 +464,7 @@ class PregelLoop:
"""Put writes for a task, to be read by the next tick."""
if not writes:
return
self._written_task_ids.add(task_id)
# 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())
@@ -468,6 +478,11 @@ class PregelLoop:
for w in self.checkpoint_pending_writes
if w[0] != task_id or w[1] not in WRITES_IDX_MAP
]
# Savers key a regular write by its position in the list, and the
# special ones dropped above won't be there next time, so they go
# last: a resume saved ahead of an update would move it to a new
# slot once the list is saved again, and it would apply twice.
writes = sorted(writes, key=lambda w: w[0] in WRITES_IDX_MAP)
writes_to_save: WritesT = [
w[1:] for w in self.checkpoint_pending_writes if w[0] == task_id
] + list(writes)
@@ -556,10 +571,12 @@ class PregelLoop:
CONFIG_KEY_CHECKPOINT_ID: self.checkpoint["id"],
},
)
# group by task id
# Writes loaded with the checkpoint are stored on it already. Saving them
# again can file them under new indexes, and then they apply twice.
by_task = defaultdict(list)
for task_id, channel, value in self.checkpoint_pending_writes:
by_task[task_id].append((channel, value))
if task_id in self._written_task_ids:
by_task[task_id].append((channel, value))
# submit writes to checkpointer
for task_id, writes in by_task.items():
if self.checkpointer_put_writes_accepts_task_path and hasattr(
@@ -904,6 +921,7 @@ class PregelLoop:
if tid != NULL_TASK_ID and isinstance(self.specs.get(ch), DeltaChannel)
}
self._exit_command_writes = []
self._written_task_ids = set()
# Resuming from a previous checkpoint requires two things:
# 1. A prior checkpoint exists (channel_versions is non-empty)
# 2. The input signals continuation (not a fresh run with new input)
@@ -931,7 +949,7 @@ class PregelLoop:
# writes so that interrupt() calls re-fire instead of returning
# stale values. But if we're actively resuming, keep them —
# multi-interrupt scenarios need previously resolved values preserved.
is_time_traveling = self.is_replaying and (
is_time_traveling = self._addressed and (
# Time-travel to a subgraph checkpoint: the parent sets
# RESUMING=True (it can't distinguish time-travel from resume),
# so we check if this subgraph's own ns is in checkpoint_map.
@@ -950,6 +968,10 @@ class PregelLoop:
or configurable.get(CONFIG_KEY_RESUMING, False)
)
)
# Finished tasks rerun only on time travel, or on a resume at a
# checkpoint the thread has moved past. A resume addressed at the
# thread's latest checkpoint hands their writes back, as without an id.
self.is_replaying = is_time_traveling or self._addressed_checkpoint_superseded
if is_time_traveling:
self.checkpoint_pending_writes = [
w for w in self.checkpoint_pending_writes if w[1] != RESUME
@@ -1165,11 +1187,11 @@ class PregelLoop:
# checkpoint ID since the fork was created after the subgraph's
# checkpoints from the original execution.
#
# Only gate on is_time_traveling (not is_replaying). When the
# Only gate on is_time_traveling (not _addressed). When the
# client resumes with an explicit checkpoint_id that happens to
# point at the current head (e.g. LangGraph Studio sending
# `checkpoint: {checkpoint_id}` alongside Command(resume=...)),
# is_replaying is True but is_time_traveling is False. In that
# the run is addressed but not time traveling. In that
# case subgraphs should load their latest checkpoint normally,
# not go through ReplayState's before-bound lookup which would
# miss subgraph checkpoints created during processing of the
@@ -1789,7 +1811,7 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager):
# This covers both normal replay and subgraphs resolved via
# checkpoint_map during time-travel.
saved = self.checkpointer.get_tuple(self.checkpoint_config)
if saved is not None and saved.metadata.get("source") in ("update", "fork"):
if saved is not None:
self._addressed_checkpoint_superseded = checkpoint_superseded(
self.checkpointer, self.checkpoint_config, saved
)
@@ -2045,7 +2067,7 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager):
# This covers both normal replay and subgraphs resolved via
# checkpoint_map during time-travel.
saved = await self.checkpointer.aget_tuple(self.checkpoint_config)
if saved is not None and saved.metadata.get("source") in ("update", "fork"):
if saved is not None:
self._addressed_checkpoint_superseded = await acheckpoint_superseded(
self.checkpointer, self.checkpoint_config, saved
)
@@ -521,10 +521,11 @@ def test_addressed_resume_keeps_a_rerun_tasks_write_to_a_new_channel(
graph = builder.compile(checkpointer=sync_checkpointer)
config = {"configurable": {"thread_id": "t"}}
graph.invoke(_both("in"), config, durability=durability)
head = graph.get_state(config).config
paused = graph.get_state(config).config
graph.update_state(config, None)
live = graph.invoke(
Command(resume="yes", update={"flag": True}), head, durability=durability
Command(resume="yes", update={"flag": True}), paused, durability=durability
)
state = graph.get_state(config)
@@ -152,6 +152,83 @@ def test_finished_task_does_not_report_answered_interrupt(
assert asked == [["A1", "B"]]
def test_finished_task_output_applies_once(
sync_checkpointer: BaseCheckpointSaver, durability: Durability
) -> None:
builder = StateGraph(State)
builder.add_node("a", lambda s: {"log": [f"a:{interrupt('A')}"]})
builder.add_node("b", lambda s: {"log": [f"b:{interrupt('B')}"]})
builder.add_node("c", lambda s: {"log": [f"c:{interrupt('C')}"]})
builder.add_edge(START, "a")
builder.add_edge(START, "b")
builder.add_edge(START, "c")
graph = builder.compile(checkpointer=sync_checkpointer)
config = _config()
result = graph.invoke({"log": []}, config, durability=durability)
while pending := result.get("__interrupt__"):
result = graph.invoke(
Command(resume={pending[0].id: "yes"}), config, durability=durability
)
assert sorted(result["log"]) == ["a:yes", "b:yes", "c:yes"]
def test_resume_with_update_applies_the_update_once(
sync_checkpointer: BaseCheckpointSaver, durability: Durability
) -> None:
def ask_three_times(state: State) -> State:
answers = [interrupt(str(question)) for question in range(1, 4)]
return {"log": ["a:" + "".join(answers)]}
graph = (
StateGraph(State)
.add_node("a", ask_three_times)
.add_edge(START, "a")
.compile(checkpointer=sync_checkpointer)
)
config = _config()
result = graph.invoke({"log": []}, config, durability=durability)
for i in range(1, 4):
result = graph.invoke(
Command(resume=str(i), update={"log": [f"u{i}"]}),
config,
durability=durability,
)
assert sorted(result["log"]) == ["a:123", "u1", "u2", "u3"]
assert sorted(graph.get_state(config).values["log"]) == sorted(result["log"])
@pytest.mark.parametrize("checkpoint_id", ["none", "head"])
def test_resume_at_the_latest_checkpoint_does_not_rerun_finished_task(
sync_checkpointer: BaseCheckpointSaver, checkpoint_id: str
) -> None:
calls: Counter[str] = Counter()
graph = _build_parallel(sync_checkpointer, calls)
config = _config()
graph.invoke({"log": []}, config)
def resume(question: str, answer: str) -> Any:
snapshot = graph.get_state(config)
target = (
snapshot.config
if checkpoint_id == "head"
else {"configurable": {**config["configurable"], "checkpoint_id": None}}
)
return graph.invoke(
Command(resume={_interrupt_by_value(snapshot, question).id: answer}),
target,
)
resume("A1", "yes")
result = resume("B", "ok")
assert sorted(result["log"]) == ["a:yes", "b:ok"]
assert calls == {"a": 2, "b": 3}, "resuming B ran the finished A again"
@NEEDS_CONTEXTVARS
async def test_finished_task_does_not_report_answered_interrupt_async(
async_checkpointer: BaseCheckpointSaver,
@@ -474,6 +551,64 @@ def test_functional_task_finished_with_none_is_not_rerun(
assert calls == {"a": 2, "b": 3}
def test_subgraph_finished_task_is_not_rerun(
sync_checkpointer: BaseCheckpointSaver, durability: Durability
) -> None:
calls: Counter[str] = Counter()
builder = StateGraph(State)
builder.add_node("child", _build_parallel(None, calls))
builder.add_edge(START, "child")
graph = builder.compile(checkpointer=sync_checkpointer)
config = _config()
graph.invoke({"log": []}, config, durability=durability)
snapshot = graph.get_state(config)
graph.invoke(
Command(resume={_interrupt_by_value(snapshot, "A1").id: "yes"}),
config,
durability=durability,
)
snapshot = graph.get_state(config)
result = graph.invoke(
Command(resume={_interrupt_by_value(snapshot, "B").id: "ok"}),
config,
durability=durability,
)
assert sorted(result["log"]) == ["a:yes", "b:ok"]
assert calls == {"a": 2, "b": 3}, "resuming B ran the finished A again"
def test_nested_entrypoint_finished_task_is_not_rerun(
sync_checkpointer: BaseCheckpointSaver,
) -> None:
calls: Counter[str] = Counter()
@task
def step_1(query: str) -> str:
calls["step_1"] += 1
return f"{query} bar"
@task
def human_feedback(query: str) -> str:
calls["human_feedback"] += 1
return f"{query} {interrupt('feedback?')}"
@entrypoint()
def child(query: str) -> str:
return human_feedback(step_1(query).result()).result()
@entrypoint(checkpointer=sync_checkpointer)
def parent(query: str) -> str:
return child.invoke(query)
config = _config()
parent.invoke("foo", config)
assert parent.invoke(Command(resume="baz"), config) == "foo bar baz"
assert calls == {"step_1": 1, "human_feedback": 2}
# --- Reading task status from recorded writes ---
+5
View File
@@ -7737,6 +7737,8 @@ def test_parallel_interrupts(sync_checkpointer: BaseCheckpointSaver) -> None:
def test_parallel_interrupts_double(sync_checkpointer: BaseCheckpointSaver) -> None:
# --- CHILD GRAPH ---
answered: Counter[str] = Counter()
class ChildState(BaseModel):
prompt: str = Field(..., description="What is going to be asked to the user?")
human_input: str | None = Field(None, description="What the human said")
@@ -7746,6 +7748,7 @@ def test_parallel_interrupts_double(sync_checkpointer: BaseCheckpointSaver) -> N
def get_human_input(state: ChildState):
human_input = interrupt(state.prompt)
answered["human"] += 1
return dict(
human_inputs=[human_input], # update parent state
@@ -7753,6 +7756,7 @@ def test_parallel_interrupts_double(sync_checkpointer: BaseCheckpointSaver) -> N
def get_dolphin_input(state: ChildState):
human_input = interrupt(state.prompt)
answered["dolphin"] += 1
return dict(
human_inputs=[human_input], # update parent state
@@ -7846,6 +7850,7 @@ def test_parallel_interrupts_double(sync_checkpointer: BaseCheckpointSaver) -> N
assert invokes == 5
assert len(events) == 5
assert answered == {"human": 2, "dolphin": 2}
def test_pregel_loop_refcount():
+467
View File
@@ -0,0 +1,467 @@
import operator
import sys
from collections.abc import Callable
from contextlib import contextmanager
from typing import Annotated, Any
import pytest
from langchain_core.messages import AIMessage, AnyMessage, HumanMessage
from langchain_core.runnables import RunnableConfig
from langgraph.checkpoint.postgres import PostgresSaver
from langgraph.checkpoint.postgres.base import PRUNE_DELTA_THREADS_SQL, PRUNE_HEADS_SQL
from typing_extensions import TypedDict
from langgraph.channels.delta import DeltaChannel
from langgraph.graph import START, StateGraph
from langgraph.graph.message import _messages_delta_reducer, add_messages
from langgraph.types import Command, StateSnapshot, interrupt
from tests.conftest import NO_DOCKER
from tests.conftest_checkpointer import (
DEFAULT_POSTGRES_URI,
_checkpointer_postgres,
_checkpointer_postgres_aio,
_checkpointer_postgres_aio_pipe,
_checkpointer_postgres_aio_pool,
_checkpointer_postgres_pipe,
_checkpointer_postgres_pool,
)
pytestmark = [
pytest.mark.anyio,
pytest.mark.skipif(NO_DOCKER, reason="needs Postgres"),
]
NEEDS_CONTEXTVARS = pytest.mark.skipif(
sys.version_info < (3, 11),
reason="Python 3.11+ is required for async contextvars support",
)
SYNC = [
_checkpointer_postgres,
_checkpointer_postgres_pipe,
_checkpointer_postgres_pool,
]
ASYNC = [
_checkpointer_postgres_aio,
_checkpointer_postgres_aio_pipe,
_checkpointer_postgres_aio_pool,
]
class State(TypedDict):
messages: Annotated[list[AnyMessage], add_messages]
notes: Annotated[list[str], operator.add]
class DeltaState(TypedDict):
messages: Annotated[list[AnyMessage], DeltaChannel(_messages_delta_reducer)]
def _reply(ask: bool):
def reply(state: dict) -> dict:
n = len(state["messages"])
answer = interrupt("approve?") if ask else "ok"
return {"messages": [AIMessage(f"{answer} {n}", id=f"reply-{n}")]}
return reply
def _graph(
checkpointer: Any,
*,
ask: bool = False,
fail: set[str] | None = None,
calls: dict[str, int] | None = None,
) -> Any:
def note(name: str):
def node(state: State) -> dict:
if calls is not None:
calls[name] = calls.get(name, 0) + 1
if fail is not None and name in fail:
fail.discard(name)
raise ValueError(f"{name} failed")
return {"notes": [name]}
return node
builder = StateGraph(State)
builder.add_node("left", note("left"))
builder.add_node("right", note("right"))
builder.add_node("reply", _reply(ask))
builder.add_edge(START, "left")
builder.add_edge(START, "right")
builder.add_edge(["left", "right"], "reply")
return builder.compile(checkpointer=checkpointer)
def _counting_subgraph(name: str, *, checkpointer: bool | None) -> Any:
class Calls(TypedDict):
messages: Annotated[list[AnyMessage], add_messages]
calls: Annotated[list[str], operator.add]
def count(state: Calls) -> dict:
n = len(state["messages"])
reply = AIMessage(f"{name} call {len(state['calls']) + 1}", id=f"{name}-{n}")
return {"calls": [name], "messages": [reply]}
builder = StateGraph(Calls)
builder.add_node("count", count)
builder.add_edge(START, "count")
return builder.compile(checkpointer=checkpointer)
def _delta_graph(checkpointer: Any) -> Any:
builder = StateGraph(DeltaState)
builder.add_node("reply", _reply(ask=False))
builder.add_edge(START, "reply")
return builder.compile(checkpointer=checkpointer)
def _config(thread_id: str) -> RunnableConfig:
return {"configurable": {"thread_id": thread_id}}
def _say(i: int) -> dict:
return {"messages": [HumanMessage(f"hi {i}", id=f"hi-{i}")]}
def _checkpoint_ids(checkpointer: Any, thread_id: str) -> list[str]:
return [c.checkpoint["id"] for c in checkpointer.list(_config(thread_id))]
def _namespaces(checkpointer: Any, thread_id: str) -> list[str]:
return sorted(
c.config["configurable"]["checkpoint_ns"]
for c in checkpointer.list(_config(thread_id))
)
async def _acheckpoint_ids(checkpointer: Any, thread_id: str) -> list[str]:
return [c.checkpoint["id"] async for c in checkpointer.alist(_config(thread_id))]
def _parentless(snapshot: StateSnapshot) -> StateSnapshot:
return snapshot._replace(
parent_config=None,
tasks=tuple(
t._replace(state=_parentless(t.state))
if isinstance(t.state, StateSnapshot)
else t
for t in snapshot.tasks
),
)
def _run_once_after(
saver: PostgresSaver, sql: str, action: Callable[[], Any], *, nth: int = 1
) -> None:
seen = 0
cursor = saver._cursor
class Hooked:
def __init__(self, cur: Any) -> None:
self.cur = cur
def __getattr__(self, name: str) -> Any:
return getattr(self.cur, name)
def execute(self, query: Any, params: Any = None, **kwargs: Any) -> None:
nonlocal seen
self.cur.execute(query, params, **kwargs)
if query is sql:
seen += 1
if seen == nth:
action()
@contextmanager
def hooked(**kwargs: Any):
with cursor(**kwargs) as cur:
yield Hooked(cur)
saver._cursor = hooked
@pytest.mark.parametrize("durability", ["sync", "async", "exit"])
@pytest.mark.parametrize("make_checkpointer", SYNC)
def test_keep_latest_keeps_the_state_and_the_thread_runs_on_like_its_twin(
make_checkpointer: Any, durability: str
) -> None:
with make_checkpointer() as checkpointer:
graph = _graph(checkpointer)
for i in range(3):
for thread_id in ("pruned", "twin"):
graph.invoke(_say(i), _config(thread_id), durability=durability)
before = graph.get_state(_config("pruned"))
checkpointer.prune(["pruned"])
assert _checkpoint_ids(checkpointer, "pruned") == [
before.config["configurable"]["checkpoint_id"]
]
assert graph.get_state(_config("pruned")) == _parentless(before)
for thread_id in ("pruned", "twin"):
graph.invoke(_say(3), _config(thread_id), durability=durability)
assert (
graph.get_state(_config("pruned")).values
== graph.get_state(_config("twin")).values
)
@pytest.mark.parametrize("make_checkpointer", SYNC)
def test_keep_latest_keeps_a_pending_interrupt_resumable(
make_checkpointer: Any,
) -> None:
with make_checkpointer() as checkpointer:
graph = _graph(checkpointer, ask=True)
for thread_id in ("pruned", "twin"):
for i in range(2):
graph.invoke(_say(i), _config(thread_id))
graph.invoke(Command(resume=f"yes {i}"), _config(thread_id))
graph.invoke(_say(2), _config(thread_id))
before = graph.get_state(_config("pruned"))
assert before.interrupts
checkpointer.prune(["pruned"])
assert graph.get_state(_config("pruned")) == _parentless(before)
for thread_id in ("pruned", "twin"):
graph.invoke(Command(resume="yes 2"), _config(thread_id))
assert (
graph.get_state(_config("pruned")).values
== graph.get_state(_config("twin")).values
)
@pytest.mark.parametrize("make_checkpointer", SYNC)
def test_keep_latest_keeps_writes_of_tasks_that_finished_before_a_failure(
make_checkpointer: Any,
) -> None:
with make_checkpointer() as checkpointer:
calls: dict[str, int] = {}
fail: set[str] = set()
graph = _graph(checkpointer, fail=fail, calls=calls)
for thread_id in ("pruned", "twin"):
for i in range(2):
graph.invoke(_say(i), _config(thread_id))
fail.add("right")
with pytest.raises(ValueError, match="right failed"):
graph.invoke(_say(2), _config(thread_id))
before = graph.get_state(_config("pruned"))
checkpointer.prune(["pruned"])
assert graph.get_state(_config("pruned")) == _parentless(before)
left_runs = calls["left"]
graph.invoke(None, _config("pruned"))
assert calls["left"] == left_runs, "a task that had finished ran again"
graph.invoke(None, _config("twin"))
assert (
graph.get_state(_config("pruned")).values
== graph.get_state(_config("twin")).values
)
@pytest.mark.parametrize("make_checkpointer", SYNC)
def test_keep_latest_keeps_an_interrupted_subgraph_resumable(
make_checkpointer: Any,
) -> None:
with make_checkpointer() as checkpointer:
sub = StateGraph(State)
sub.add_node("reply", _reply(ask=True))
sub.add_edge(START, "reply")
parent = StateGraph(State)
parent.add_node("sub", sub.compile())
parent.add_edge(START, "sub")
graph = parent.compile(checkpointer=checkpointer)
for thread_id in ("pruned", "twin"):
for i in range(2):
graph.invoke(_say(i), _config(thread_id))
graph.invoke(Command(resume=f"yes {i}"), _config(thread_id))
graph.invoke(_say(2), _config(thread_id))
before = graph.get_state(_config("pruned"), subgraphs=True)
interrupted = before.tasks[0].state.config["configurable"]["checkpoint_ns"]
checkpointer.prune(["pruned"])
assert _namespaces(checkpointer, "pruned") == ["", interrupted]
assert graph.get_state(_config("pruned"), subgraphs=True) == _parentless(before)
for thread_id in ("pruned", "twin"):
graph.invoke(Command(resume="yes 2"), _config(thread_id))
assert (
graph.get_state(_config("pruned")).values
== graph.get_state(_config("twin")).values
)
@pytest.mark.parametrize("durability", ["sync", "async", "exit"])
@pytest.mark.parametrize("make_checkpointer", SYNC)
def test_keep_latest_keeps_a_checkpointer_true_subgraph_but_not_finished_calls(
make_checkpointer: Any, durability: str
) -> None:
with make_checkpointer() as checkpointer:
parent = StateGraph(State)
parent.add_node("memory", _counting_subgraph("memory", checkpointer=True))
parent.add_node("call", _counting_subgraph("call", checkpointer=None))
parent.add_edge(START, "memory")
parent.add_edge("memory", "call")
graph = parent.compile(checkpointer=checkpointer)
for i in range(3):
for thread_id in ("pruned", "twin"):
graph.invoke(_say(i), _config(thread_id), durability=durability)
checkpointer.prune(["pruned"])
assert _namespaces(checkpointer, "pruned") == ["", "memory"]
for thread_id in ("pruned", "twin"):
graph.invoke(_say(3), _config(thread_id), durability=durability)
assert (
graph.get_state(_config("pruned")).values
== graph.get_state(_config("twin")).values
)
def test_keep_latest_keeps_a_turn_saved_after_it_read_the_heads() -> None:
with (
_checkpointer_postgres() as checkpointer,
PostgresSaver.from_conn_string(
DEFAULT_POSTGRES_URI + checkpointer.conn.info.dbname
) as pruner,
):
graph = _graph(checkpointer)
for i in range(3):
for thread_id in ("pruned", "twin"):
graph.invoke(_say(i), _config(thread_id))
_run_once_after(
pruner, PRUNE_HEADS_SQL, lambda: graph.invoke(_say(3), _config("pruned"))
)
pruner.prune(["pruned"])
graph.invoke(_say(3), _config("twin"))
assert (
graph.get_state(_config("pruned")).values
== graph.get_state(_config("twin")).values
)
def test_keep_latest_refuses_a_thread_that_switches_to_a_delta_channel_mid_batch() -> (
None
):
with (
_checkpointer_postgres() as checkpointer,
PostgresSaver.from_conn_string(
DEFAULT_POSTGRES_URI + checkpointer.conn.info.dbname
) as pruner,
):
graph, delta_graph = _graph(checkpointer), _delta_graph(checkpointer)
for i in range(3):
for thread_id in ("switched", "twin"):
graph.invoke(_say(i), _config(thread_id))
_run_once_after(
pruner,
PRUNE_DELTA_THREADS_SQL,
lambda: delta_graph.invoke(_say(3), _config("switched")),
)
with pytest.raises(NotImplementedError, match="switched"):
pruner.prune(["switched"])
delta_graph.invoke(_say(3), _config("twin"))
assert (
delta_graph.get_state(_config("switched")).values
== delta_graph.get_state(_config("twin")).values
)
def test_keep_latest_keeps_a_delta_channel_turn_saved_after_the_thread_check() -> None:
with (
_checkpointer_postgres() as checkpointer,
PostgresSaver.from_conn_string(
DEFAULT_POSTGRES_URI + checkpointer.conn.info.dbname
) as pruner,
):
graph, delta_graph = _graph(checkpointer), _delta_graph(checkpointer)
for i in range(3):
for thread_id in ("switched", "twin"):
graph.invoke(_say(i), _config(thread_id))
head = _checkpoint_ids(checkpointer, "switched")[0]
_run_once_after(
pruner,
PRUNE_DELTA_THREADS_SQL,
lambda: delta_graph.invoke(_say(3), _config("switched")),
nth=2,
)
pruner.prune(["switched"])
assert _checkpoint_ids(checkpointer, "switched")[-1] == head
delta_graph.invoke(_say(3), _config("twin"))
assert (
delta_graph.get_state(_config("switched")).values
== delta_graph.get_state(_config("twin")).values
)
@pytest.mark.parametrize("make_checkpointer", SYNC)
def test_keep_latest_refuses_a_batch_with_a_delta_channel_thread(
make_checkpointer: Any,
) -> None:
with make_checkpointer() as checkpointer:
graph, delta_graph = _graph(checkpointer), _delta_graph(checkpointer)
for i in range(3):
graph.invoke(_say(i), _config("ordinary"))
delta_graph.invoke(_say(i), _config("delta"))
before = {t: _checkpoint_ids(checkpointer, t) for t in ("ordinary", "delta")}
delta_before = delta_graph.get_state(_config("delta"))
with pytest.raises(NotImplementedError, match="delta"):
checkpointer.prune(["ordinary", "delta"])
assert {t: _checkpoint_ids(checkpointer, t) for t in before} == before
assert delta_graph.get_state(_config("delta")) == delta_before
@NEEDS_CONTEXTVARS
@pytest.mark.parametrize("make_checkpointer", ASYNC)
async def test_akeep_latest_keeps_a_pending_interrupt_resumable(
make_checkpointer: Any,
) -> None:
async with make_checkpointer() as checkpointer:
graph = _graph(checkpointer, ask=True)
for thread_id in ("pruned", "twin"):
for i in range(2):
await graph.ainvoke(_say(i), _config(thread_id))
await graph.ainvoke(Command(resume=f"yes {i}"), _config(thread_id))
await graph.ainvoke(_say(2), _config(thread_id))
before = await graph.aget_state(_config("pruned"))
await checkpointer.aprune(["pruned"])
assert await _acheckpoint_ids(checkpointer, "pruned") == [
before.config["configurable"]["checkpoint_id"]
]
assert await graph.aget_state(_config("pruned")) == _parentless(before)
for thread_id in ("pruned", "twin"):
await graph.ainvoke(Command(resume="yes 2"), _config(thread_id))
assert (await graph.aget_state(_config("pruned"))).values == (
await graph.aget_state(_config("twin"))
).values
@pytest.mark.parametrize("make_checkpointer", ASYNC)
async def test_akeep_latest_refuses_a_batch_with_a_delta_channel_thread(
make_checkpointer: Any,
) -> None:
async with make_checkpointer() as checkpointer:
graph, delta_graph = _graph(checkpointer), _delta_graph(checkpointer)
for i in range(3):
await graph.ainvoke(_say(i), _config("ordinary"))
await delta_graph.ainvoke(_say(i), _config("delta"))
before = {
t: await _acheckpoint_ids(checkpointer, t) for t in ("ordinary", "delta")
}
with pytest.raises(NotImplementedError, match="delta"):
await checkpointer.aprune(["ordinary", "delta"])
assert {t: await _acheckpoint_ids(checkpointer, t) for t in before} == before
+37
View File
@@ -1177,6 +1177,43 @@ def test_subgraph_interrupt_resume_with_explicit_head_checkpoint_id(
assert result["value"] == ["sub_a", "human:answer", "sub_b"]
def test_subgraph_replay_reasks_answered_parallel_interrupt(
sync_checkpointer: BaseCheckpointSaver,
) -> None:
"""Replaying a subgraph checkpoint asks its step's questions again, including
one already answered. An update checkpoint replays in place instead of
forking, so its tasks keep their ids and its saved answers could be reused."""
subgraph = (
StateGraph(State)
.add_node("a", lambda s: {"value": [f"a:{interrupt('a')}"]})
.add_node("b", lambda s: {"value": [f"b:{interrupt('b')}"]})
.add_edge(START, "a")
.add_edge(START, "b")
.compile(checkpointer=True)
)
graph = (
StateGraph(State)
.add_node("subgraph_node", subgraph)
.add_edge(START, "subgraph_node")
.compile(checkpointer=sync_checkpointer)
)
config = {"configurable": {"thread_id": "1"}}
def subgraph_head() -> RunnableConfig:
return graph.get_state(config, subgraphs=True).tasks[0].state.config
graph.invoke({"value": []}, config)
graph.update_state(subgraph_head(), {"value": ["u"]}, as_node="__start__")
result = graph.invoke(None, config)
question_a = next(i for i in result["__interrupt__"] if i.value == "a")
graph.invoke(Command(resume={question_a.id: "yes"}), config)
result = graph.invoke(None, subgraph_head())
assert sorted(i.value for i in result["__interrupt__"]) == ["a", "b"]
def test_subgraph_replay_loads_accumulated_state_then_resume(
sync_checkpointer: BaseCheckpointSaver,
) -> None:
+2 -2
View File
@@ -1685,7 +1685,7 @@ test = [
[[package]]
name = "langgraph-checkpoint-postgres"
version = "3.1.3"
version = "3.2.0"
source = { editable = "../checkpoint-postgres" }
dependencies = [
{ name = "langgraph-checkpoint" },
@@ -1835,7 +1835,7 @@ test = [
[[package]]
name = "langgraph-prebuilt"
version = "1.1.0"
version = "1.1.1"
source = { editable = "../prebuilt" }
dependencies = [
{ name = "langchain-core" },
+1 -1
View File
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
[project]
name = "langgraph-prebuilt"
version = "1.1.0"
version = "1.1.1"
description = "Library with high-level APIs for creating and executing LangGraph agents and tools."
authors = []
requires-python = ">=3.10"
+2 -2
View File
@@ -418,7 +418,7 @@ test = [
[[package]]
name = "langgraph-checkpoint-postgres"
version = "3.1.3"
version = "3.2.0"
source = { editable = "../checkpoint-postgres" }
dependencies = [
{ name = "langgraph-checkpoint" },
@@ -512,7 +512,7 @@ test = [
[[package]]
name = "langgraph-prebuilt"
version = "1.1.0"
version = "1.1.1"
source = { editable = "." }
dependencies = [
{ name = "langchain-core" },
+1 -1
View File
@@ -433,7 +433,7 @@ test = [
[[package]]
name = "langgraph-prebuilt"
version = "1.1.0"
version = "1.1.1"
source = { editable = "../prebuilt" }
dependencies = [
{ name = "langchain-core" },