Compare commits

..
Author SHA1 Message Date
Sydney RunkleandGitHub 71e4eafcb8 Merge branch 'main' into delta-channel-writes-based 2026-04-29 17:15:31 -04:00
Will Fu-Hinthorn af82510cf9 cleanup 2026-04-29 13:56:25 -07:00
Will Fu-Hinthorn e8e730c861 fixup exit mode 2026-04-29 13:30:59 -07:00
Sydney Runkle e735645264 style: apply ruff format fixes 2026-04-29 14:15:02 -04:00
Sydney Runkle 68f7a3acc4 test: add add_messages migration test; move all imports to module level
- Add test_add_messages_to_delta_migration_preserves_message_history (sync
  + async) covering the primary real-world BinaryOperatorAggregate →
  DeltaChannel migration path with real Message objects and IDs
- Hoist all in-function imports to module level in test_channels.py and
  fix _delta_channel_with_type helper accordingly
- Add section headers in test_channels.py for better navigation
2026-04-29 14:11:17 -04:00
Sydney Runkle be7101b0ff nits 2026-04-29 14:05:06 -04:00
Sydney Runkle ee5fd582b8 refactor: remove InMemorySaver.prune — out of scope for DeltaChannel PR
prune was not previously implemented on InMemorySaver (raised
NotImplementedError); adding a DeltaChannel-aware implementation is a
follow-up concern, not required for the core feature.
2026-04-29 14:01:41 -04:00
Sydney Runkle 9a5f844e1b refactor: remove unnecessary variable extractions from checkpoint load paths
Revert pure-style refactors (local variable hoisting, redundant null
guards, Sequence/list annotation change) that cluttered the DeltaChannel
PR diff without any semantic change.
2026-04-29 13:53:53 -04:00
Sydney RunkleandClaude Sonnet 4.6 c0c5479722 fix(pregel): revert unnecessary default on increment channel param
channel: None = None -> channel: None; all call sites pass None explicitly
and the default was never needed. Restores the original signature.

Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>
2026-04-29 13:42:46 -04:00
Sydney RunkleandClaude Sonnet 4.6 0e5d61692e refactor(channels): DeltaChannel batch reducer interface + _messages_delta_reducer
Renames `operator` → `reducer` and flips arg order to `(reducer, typ=None)`,
matching the new batch contract: `reducer(state, list[writes]) -> state`. The
reducer receives all writes for a step in one call instead of being folded
pairwise, enabling single-pass implementations that avoid O(N²) reprocessing.

`typ` is now optional — `_is_field_channel` in `graph/state.py` always
overwrites it from the `Annotated[T, ...]` outer type, so users can write
`DeltaChannel(my_reducer)` rather than `DeltaChannel(list, my_reducer)`.

Adds `_messages_delta_reducer` to `langgraph.graph.message` (experimental):
a single-pass bulk reducer for message lists that deduplicates by ID and
handles `RemoveMessage` tombstoning without calling `add_messages`, avoiding
repeated dedup passes that `add_messages` would incur in a fold.

Also fixes the `_delta_write_futs` mypy error in `AsyncPregelLoop` by moving
the type annotation to the class body, and unignores `new_pr_desc.md` from
the repo via `.gitignore`.

Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>
2026-04-29 12:58:49 -04:00
Sydney RunkleandClaude Sonnet 4.6 959c8c8618 fix(pregel): async write-ordering safety for DeltaChannel via _delta_write_futs
In durability="async" mode (the default), put_writes calls are
fire-and-forget coroutines — a process crash between write submission and
checkpoint commit leaves a DELTA_SENTINEL blob with no backing writes,
causing silent data loss on replay.

AsyncPregelLoop now maintains _delta_write_futs: any write to a
DeltaChannel channel appends its asyncio.Future to this list in
accept_writes. _checkpointer_put_after_previous drains the list with
await asyncio.gather() before calling aput(), guaranteeing
checkpoint_writes are durable before the sentinel blob is committed.

The sync loop is unchanged: BackgroundExecutor.__exit__ already ensures
all background tasks complete before invoke() returns.

Also fixes DeltaChannel(list, add_messages) constructor call in
checkpoint-postgres async test (missing typ arg).

Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>
2026-04-29 10:51:18 -04:00
Sydney RunkleandClaude Sonnet 4.6 7439ab2e5b feat(serde): DELTA_SENTINEL via msgpack ext 8; remove "delta" type tag
DELTA_SENTINEL is now serialized as a msgpack ext code (EXT_DELTA_SENTINEL=8)
alongside _DeltaSnapshot (ext 7), keeping both sentinel types in the same
codec path. The dedicated "delta" string type tag and its special-case in
dumps_typed/loads_typed are removed — no migration needed since this is
introduced fresh.

InMemorySaver.prune() updated to deserialize blobs and check `is DELTA_SENTINEL`
rather than comparing the raw type tag string, making it codec-agnostic.

Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>
2026-04-29 09:59:56 -04:00
Sydney RunkleandGitHub c2c1e4412d Merge branch 'main' into delta-channel-writes-based 2026-04-29 09:43:48 -04:00
Sydney RunkleandClaude Sonnet 4.6 e6e44d5512 refactor(channels): DeltaChannel takes typ as first arg, matching BinOpChannel
Previously DeltaChannel.__init__ hardcoded typ=list and _is_field_channel
patched item.typ/item.value after construction. This mirrors BinaryOperatorAggregate:

- DeltaChannel(typ, operator, *, snapshot_frequency=None) — typ is now
  a required first argument; __init__ strips abstract/parameterized types
  to their concrete counterparts (same logic as BinaryOperatorAggregate)
- _is_field_channel reconstructs the channel via its constructor instead
  of patching typ and value externally
- copy() and from_checkpoint() use self.__class__(self.typ, self.operator, ...)
  — no post-construction attribute hacking needed
- _empty() helper removed; self.typ() is always a concrete callable
- All call sites updated: DeltaChannel(list, op), DeltaChannel(dict, op), etc.

Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>
2026-04-29 09:42:45 -04:00
Sydney RunkleandClaude Sonnet 4.6 16c09c1bad refactor(channels): inline _clone_empty into copy and from_checkpoint
The helper was three lines called from exactly two places — inlining
it removes indirection without adding duplication.

Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>
2026-04-29 09:34:48 -04:00
Sydney RunkleandClaude Sonnet 4.6 a98ffc95fd fix(checkpoint-postgres): type blob_values as Sequence[tuple[bytes, bytes, bytes]]
Replace Any with the concrete element type — each row is (key, type_tag,
blob) all as bytes — matching how _load_blobs unpacks and decodes them.

Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>
2026-04-29 09:31:31 -04:00
Sydney RunkleandClaude Sonnet 4.6 08304d5c24 revert(_checkpoint-postgres): restore _load_blobs comprehension
The loop refactor added no clarity — revert to the original one-liner.

Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>
2026-04-29 09:30:33 -04:00
Sydney RunkleandClaude Sonnet 4.6 9b4bbd0649 fix(_checkpoint): remove assert, use cast for DeltaChannel narrowing
_needs_replay already gates on isinstance(spec, DeltaChannel), so the
assert/isinstance checks inside the replay branch were unreachable.
Replace with cast(DeltaChannel, spec) for zero-cost type narrowing that
survives -O and avoids any runtime check.

Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>
2026-04-29 09:28:13 -04:00
Sydney RunkleandClaude Sonnet 4.6 f80a18b376 fix(channels): address code review feedback on DeltaChannel
delta.py — value consistency:
- `__init__` now starts with `value=MISSING` (was `[]`); both fresh
  construction and clones are consistently uninitialised until
  `from_checkpoint()` or `copy()` sets the real value
- `_clone_empty` drops `__new__` in favour of the normal constructor;
  `typ` and `key` are restored explicitly afterwards (`typ` may differ
  from `list` when set via Annotated injection; `key` is injected by
  the graph builder after construction)

delta.py — snapshot cadence:
- `is_snapshot_step` now guards `step > 0`; snapshots fire at steps N,
  2N, 3N, … instead of also at step 0 where `0 % N == 0` always held

_checkpoint.py — runtime guards:
- replace both `assert isinstance(spec, DeltaChannel)` with proper
  `if not isinstance: raise TypeError`; `assert` is stripped by `-O`
  and is wrong for production invariant checks

binop.py — avoid unnecessary allocation:
- `_get_overwrite`: replace `set(value.keys()) == {OVERWRITE}` with
  `len(value) == 1 and OVERWRITE in value` to avoid allocating a
  throwaway set on every call

checkpoint-postgres — typed rows:
- add `_DeltaCombinedRow(TypedDict, total=False)` documenting the nine
  columns emitted by `SELECT_DELTA_COMBINED_SQL`'s UNION ALL; change
  `_build_delta_channel_writes_history` parameter from `Sequence[Any]`
  to `Sequence[_DeltaCombinedRow]`; call sites cast the psycopg
  `DictRow` result accordingly

Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>
2026-04-29 09:22:05 -04:00
Sydney RunkleandClaude Sonnet 4.6 7df8b54ba5 refactor(channels): extract _operators_equal helper, deduplicate __eq__ logic
Both BinaryOperatorAggregate and DeltaChannel had identical inline logic for
comparing operators that may be lambdas. Extract _operators_equal into
binop.py (alongside _get_overwrite) and use it in both __eq__ methods.

Also removes the duplicate _get_overwrite definition from delta.py — it was
identical to binop.py's and is now imported from there instead, along with
the now-unused OVERWRITE constant and Overwrite imports.

Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>
2026-04-28 18:02:18 -04:00
Sydney RunkleandClaude Sonnet 4.6 5787258a63 fix(checkpoint-postgres): remove unused _DeltaSnapshot import
Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>
2026-04-28 17:58:49 -04:00
Sydney RunkleandClaude Sonnet 4.6 e37e68d631 fix(checkpoint): remove unused _DeltaSnapshot import
Dropped after removing the _DeltaSnapshot special-case in the seed-terminator
logic — write-collection ordering fix handles both blob types uniformly.

Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>
2026-04-28 17:58:46 -04:00
Sydney RunkleandClaude Sonnet 4.6 702895e484 fix(checkpoint-postgres): single-roundtrip CTE query + write-ordering fix
Replaces the three sequential SELECT roundtrips in _get_channel_writes_history
(checkpoints, checkpoint_writes, checkpoint_blobs) with one combined
UNION ALL query tagged by a _kind discriminator column. Both sync and async
paths now do one execute + one fetchall regardless of pipeline mode.

_build_delta_channel_writes_history is updated to accept the single tagged
rows list and dispatch on _kind while building its lookup dicts; the three
old SQL constants are removed.

Also fixes write-collection ordering in _build_delta_channel_writes_history:
the seed-terminator blob check previously fired before collecting that
ancestor's writes, silently dropping the transition writes needed to
reconstruct the child's state. Writes are now collected first.

Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>
2026-04-28 17:55:29 -04:00
Sydney RunkleandClaude Sonnet 4.6 f879e49a96 fix(checkpoint): diamond pattern replaces ContextVar re-entrancy guard
Adds _get_tuple_raw / _aget_tuple_raw as the pure-storage-read layer that
_get_channel_writes_history calls instead of get_tuple. Default
implementation delegates to get_tuple for full backward compatibility — no
changes needed for existing savers whose get_tuple is a plain storage query.

Savers that perform channel hydration inside get_tuple can override
_get_tuple_raw with the raw read to structurally break any possible cycle;
a Python RecursionError surfaces the problem if they don't, rather than the
previous silent data corruption (returning empty writes).

Also fixes write-collection ordering in the reference implementation: pending
writes from the seed-terminator ancestor were silently dropped because the
terminator check fired before the collection loop. Writes are now collected
first so the seed ancestor's transition writes are included in reconstruction.

Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>
2026-04-28 17:54:58 -04:00
Sydney RunkleandClaude Sonnet 4.6 5b3d4a218f fix(channels): DeltaChannel subclass safety and order-independent Overwrite
copy() and from_checkpoint() hardcoded DeltaChannel instead of
self.__class__, breaking subclasses. Now mirrors the BinaryOperatorAggregate
pattern: self.__class__(self.operator) with explicit typ/key assignment.

update() applied non-overwrite values that arrived before an Overwrite in
the sequence, then discarded them when the Overwrite fired — order-dependent
behaviour in a method whose contract says order is arbitrary. Now pre-scans
for an Overwrite and applies only it (or folds all values normally if none).

Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>
2026-04-28 17:52:26 -04:00
Sydney RunkleandClaude Sonnet 4.6 bfc192c02e fix(channels): correct _strip_extras dead branch for Required/NotRequired
The second `if hasattr(t, "__origin__")` block was unreachable — the first
branch always returned, so Required[T] / NotRequired[T] resolved to the bare
class instead of the inner type. Check Required/NotRequired before the generic
__origin__ fallback.

Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>
2026-04-28 17:51:11 -04:00
9abee46990 feat(langgraph): DeltaChannel snapshot_frequency — bounded read depth with write-count snapshotting (#7634)
## Summary

Builds on #7586. Adds `snapshot_frequency: int | None` to
`DeltaChannel`, letting users trade storage for bounded read depth. Also
promotes `channels/_delta.py` from private to public
(`channels/delta.py`).

### How it works

Every Nth **pregel step**, `create_checkpoint` writes a `_DeltaSnapshot`
blob instead of `DELTA_SENTINEL`. The ancestor walk in
`_get_channel_writes_history` terminates at the snapshot rather than
walking the full chain, bounding replay to at most N steps.

Snapshots are **eager**: fired even on steps where the channel had no
write (via a `get_next_version` version bump), so the depth bound holds
unconditionally — no risk of the cadence drifting if a channel happens
to be silent at a snapshot step.

### Storage formula

| Mode | Blob storage | Read depth |
|------|-------------|------------|
| `snapshot_frequency=None` (pure delta) | O(N) — sentinels only | O(N)
steps |
| `snapshot_frequency=K` | O(N²/K) — periodic snapshots of growing size
| O(K) steps |
| add_messages / BinOp | O(N²) — full blob every step | O(1) |

At N turns with ~400 char/msg messages, total snapshot storage ≈ N²/(2K)
× avg_msg_size, since each snapshot blob grows linearly with accumulated
messages.

### Key design decisions

- **Step-based**: `snapshot_frequency=K` means "snapshot every K pregel
steps." `create_checkpoint` has the step number; the channel itself
doesn't need to track writes.
- **Eager**: version-bumped via `get_next_version` even on non-write
steps so `put()` always stores the blob.
- **`_DeltaSnapshot` NamedTuple + msgpack ext type**
(`EXT_DELTA_SNAPSHOT = 7`): serde type tag dispatches in
`from_checkpoint` — no dict key inspection, no collision risk.
- **`from_checkpoint` semantics**: `_DeltaSnapshot` → restore value
directly (no replay needed); `DELTA_SENTINEL` / `MISSING` → replay from
ancestor writes; plain value → pre-migration BinOp blob.
- **InMemorySaver and PostgresSaver updated**:
`_get_channel_writes_history` collects the snapshot ancestor's
pending_writes before terminating (they encode the *next* step's
transition, unlike pre-delta migration blobs which subsume their own
writes).
- **`snapshot_frequency=None`** is the pure-delta default (replaces
`math.inf`).

### Benchmark results (InMemory, ~400 char/msg)

**Storage**

| turns | ctx | freq=1 | freq=5 | freq=10 | freq=50 | freq=inf |
|------:|----:|-------:|-------:|--------:|--------:|---------:|
| 50 | ~10K tok | 5.9 MB | 1.2 MB | 601.3 KB | 119.8 KB | 29.5 KB |
| 100 | ~20K tok | 23.7 MB | 4.8 MB | 2.4 MB | 475.8 KB | 58.4 KB |
| 200 | ~40K tok | 94.6 MB | 19.0 MB | 9.5 MB | 1.9 MB | 116.4 KB |
| 500 | ~100K tok | 591.5 MB | 118.4 MB | 59.2 MB | 11.8 MB | 290.3 KB |

**Read latency** (avg of 5 `get_state` calls)

| turns | ctx | freq=1 | freq=5 | freq=10 | freq=50 | freq=inf |
|------:|----:|-------:|-------:|--------:|--------:|---------:|
| 50 | ~10K tok | 0.4ms | 0.4ms | 0.7ms | 0.9ms | 1.8ms |
| 100 | ~20K tok | 0.7ms | 0.9ms | 1.0ms | 1.7ms | 5.7ms |
| 200 | ~40K tok | 1.5ms | 1.7ms | 4.5ms | 3.7ms | 20.1ms |
| 500 | ~100K tok | 3.6ms | 4.2ms | 4.4ms | 9.0ms | 110.3ms |

**Per-invoke write latency**

| turns | ctx | freq=1 | freq=5 | freq=10 | freq=50 | freq=inf |
|------:|----:|-------:|-------:|--------:|--------:|---------:|
| 50 | ~10K tok | 1.5ms | 1.1ms | 1.1ms | 1.3ms | 1.7ms |
| 100 | ~20K tok | 2.5ms | 1.6ms | 1.5ms | 1.7ms | 3.3ms |
| 200 | ~40K tok | 3.4ms | 2.3ms | 2.2ms | 2.5ms | 8.3ms |
| 500 | ~100K tok | 6.2ms | 4.2ms | 3.6ms | 4.1ms | 39.2ms |

## Test plan

- [x] `make format` / `make lint` clean across `langgraph`,
`checkpoint`, `checkpoint-postgres`
- [x] `tests/test_channels.py` — 37 passing including step-based and
eager-snapshot tests
- [x] `tests/test_delta_channel_migration.py` — all passing
- [x] Full suite: 1387 passing, 6 pre-existing failures unrelated to
this branch

---------

Co-authored-by: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>
2026-04-28 16:42:48 -04:00
Sydney Runkle 0ae81f3cff format, lint, restructure 2026-04-24 07:46:57 -04:00
Sydney Runkle afec98f369 internal for now 2026-04-24 07:29:52 -04:00
Sydney RunkleandClaude Opus 4.7 f25d1935ef fix(postgres): handle missing checkpoint_id in _get_channel_writes_history; update test signatures
Two fixes exposed by running the postgres test suite against a local
postgres instance:

1. `PostgresSaver._get_channel_writes_history` /
   `AsyncPostgresSaver._aget_channel_writes_history` required
   `checkpoint_id` in the passed config, raising `KeyError` when called
   with just `thread_id` (e.g. `graph.aget_state({"thread_id": "..."})`).
   Now resolves to the latest checkpoint via `get_tuple`/`aget_tuple`
   when the id is missing.

2. `test_get_checkpoint_no_channel_values` (sync + async) monkeypatched
   `_load_checkpoint_tuple` with the old `(value, cur)` signature. Method
   now takes `(value)` only since delta reconstruction moved out of the
   tuple-load path — updated both tests.

Local postgres (`brew install pgvector postgresql@16`, running on port
5441) now exercises all 40 non-vector postgres tests green.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
2026-04-23 14:47:22 -04:00
Sydney RunkleandClaude Opus 4.7 3a7ed5b454 refactor(delta-channel): honest data model, private experimental API
Restructure DeltaChannel reconstruction so the hydration path matches
pregel's storage axes (blobs + writes) without leaking internal DTOs
into the public checkpoint contract.

Key changes:

* Deleted `DeltaChannelWrites` dataclass and `SEED_UNSET` sentinel.
  Reconstruction data no longer flows through `Checkpoint.channel_values`
  as a wrapped DTO — that field now carries a value or `DELTA_SENTINEL`,
  never a reconstruction shape.
* Added private `_ChannelWritesHistory(seed: Any, writes: list[PendingWrite])`
  NamedTuple as the return type for the new storage-level query.
* Added private, experimental `_get_channel_writes_history` /
  `_aget_channel_writes_history` on `BaseCheckpointSaver` — reference
  impl via `get_tuple` + `parent_config` walk, overridden on
  `InMemorySaver` / `PostgresSaver` / `AsyncPostgresSaver` for perf.
  Fixes a latent migration bug in the base fallback (now inspects
  ancestor `channel_values` for pre-delta seed).
* `DeltaChannel.from_checkpoint(seed)` simplified to two cases
  (sentinel/MISSING → empty, else → seed). New `replay_writes` method
  folds `list[PendingWrite]` through the reducer.
* Delta hydration consolidated inside `channels_from_checkpoint` via
  optional `saver` + `config` kwargs (+ async mirror
  `achannels_from_checkpoint`). All six pregel call sites updated.
  `get_tuple` no longer patches `channel_values` — removed
  `_resolve_delta_channels` (memory) and per-tuple reconstruction from
  `_load_checkpoint_tuple` (postgres sync + async).
* Hydration short-circuits on the target's own blob: if
  `channel_values[k]` is a real value (pre-migration tip, `update_state`
  result), use it directly. Only walks ancestors when the target holds
  sentinel or is missing. Fixes a correctness bug where migration-tip
  and `update_state` values would be lost.
* New test_delta_channel_migration.py: 10 scenarios covering
  BinaryOperatorAggregate → DeltaChannel migration (basic + async,
  time-travel, fork, `update_state`, tip-of-pre-migration, base-saver
  fallback parity, cross-thread isolation).

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
2026-04-23 14:15:02 -04:00
Sydney Runkle 31ef0e942a refactor(delta-channel): drop snapshot_every and saver Overwrite terminator
snapshot_every was a knob for bounding reconstruction cost on deep threads.
Benchmarks (notes/add_messages_replay_problem.md + scratch work on
sr/add-messages-replay-bench) showed the add_messages fast-path
(optimize/add-messages-fast-path) closes the quadratic replay cost for
threads under ~1000 turns, where the crossover to snapshots makes sense.
For deeper threads we'll ship a first-class compaction primitive instead.

Removals:

* DeltaChannel: snapshot_every ctor param, _writes_since_snapshot counter,
  should_snapshot() / snapshot_write() methods, counter threading through
  _apply_write / update / from_checkpoint / copy.
* Pregel loop: post-checkpoint snapshot-injection block and
  SNAPSHOT_TASK_ID import + constant.
* Checkpoint base: _overwrite_types() helper and the ancestor-walk
  short-circuit on user-emitted Overwrite in sync + async
  get_channel_writes.
* InMemory + Postgres savers: same walk-terminator shortcut. The
  pre-delta blob terminator (seed-from-ancestor-blob) stays — it's
  required for migration correctness, not a snapshot optimization.
* Tests for all of the above.

Preserved:

* Channel-level Overwrite semantics in DeltaChannel / BinOpAggregate:
  Overwrite still resets the value at reducer level; same-super-step
  dedup and InvalidUpdateError on multiple Overwrites still enforced.
* Pre-delta migration seeding.
2026-04-23 09:54:12 -04:00
Sydney Runkle d120f127ca refactor(delta-channel): plain SELECT WHERE replaces recursive CTE
The recursive CTE was bottlenecked by a JSON-expression join
(`bl.version = checkpoint->'channel_versions'->>bl.channel`) that the
planner could not index, producing an O(ancestors x blobs) nested-loop.
At depth 1000 it ran ~275 ms and removed ~2M filter rows; the recursion
itself was 2.4 ms.

Switch to three plain indexed SELECTs per delta channel
(checkpoints, checkpoint_writes, checkpoint_blobs); a pure helper on
BasePostgresSaver walks the parent chain and assembles
DeltaChannelWrites. Sync (__init__.py) and async (aio.py) each own
their three-roundtrip I/O wrappers.

Bench numbers (notes/delta_channel_query_bench.md): 3x at depth 50,
15x at depth 200, ~100x at depth 1000. Plain over-fetches sibling rows
when the thread branches but still wins at every realistic depth on
both local and remote postgres.

Multi-channel coalescing dropped — reconstruction is per-channel now.
Same shape as InMemorySaver. Can come back as a SQL-level
optimization later if needed.
2026-04-23 08:21:00 -04:00
Sydney Runkle 9e330c96dc contextvar 2026-04-23 07:42:51 -04:00
Sydney RunkleandClaude Opus 4.7 cd8fad5905 fix(delta-channel): target-exclusion, pre-delta seed, one-query postgres walk
Four fixes from an independent review of the reconstruction pipeline, plus
a structural cleanup:

1. Ancestor walk excludes the target checkpoint itself (matches pregel:
   writes stored under checkpoint_id=T are pending for the NEXT step and
   applied separately via apply_writes). Memory saver previously included
   them, diverging from Postgres and causing pending writes to be folded
   into the reconstructed snapshot — visible via get_state during
   interrupts and time-travel into a non-leaf checkpoint.

2. Pre-delta blob terminator. When the walk hits an ancestor whose blob
   for the channel is a real value (not DELTA_SENTINEL), bind that blob
   as DeltaChannelWrites.seed and stop. Without this, threads migrated
   from pre-delta storage would replay ancestor writes to the root
   forever AND lose any value that lived only in the old blob
   (e.g. from update_state). Per-ancestor, the blob is checked BEFORE
   its writes — a pre-delta blob subsumes writes at the same checkpoint,
   so including them would double-count.

3. Base-fallback get_channel_writes follows parent_checkpoint_id instead
   of list(before=...). The previous form returned every tuple with
   id<target, including sibling branches on forked threads.

4. seed replaces the Overwrite-wrapping hack for pre-delta values.
   DeltaChannelWrites(writes, seed=SEED_UNSET) makes the saver's
   reconstruction terminator semantically explicit; drops the lazy
   _make_overwrite import dance. User-emitted Overwrite still reset the
   chain via _apply_write as before.

Postgres: recursive CTE enumerates on-path ancestors and joins once
against checkpoint_writes and once against checkpoint_blobs for every
delta channel in the get_tuple — one roundtrip instead of the previous
3 queries × N channels.

Tests added:
- Pre-delta blob seeding (seed binding, no double-counting of ancestor
  writes at the terminator, pending-at-target excluded).
- Root checkpoint returns empty writes.
- Seed-based from_checkpoint replay (three scenarios: with writes,
  seed-only, seed=None distinct from SEED_UNSET).

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
2026-04-22 22:09:08 -04:00
Sydney Runkle acc7eda8c5 optimizations i sure hope 2026-04-22 20:47:51 -04:00
Sydney Runkle 5b7fdf5655 eh 2026-04-22 18:41:06 -04:00
Sydney RunkleandClaude Opus 4.7 4cad68f767 fix(delta-channel): unwrap NotRequired[X] for dict/set reducers
Annotated[NotRequired[dict[...]], DeltaChannel(reducer)] (the shape used
by deepagents' filesystem middleware) fell through type inference to
`list`, so the first operator call blew up with
"'list' object is not a mapping". `_is_field_channel` now unwraps a
parameterized Required[X]/NotRequired[X] before stripping extras, which
lets dict/set/mapping outer types reach the abc normalization block.

Also type-annotates the `new` locals in DeltaChannel.copy() and
from_checkpoint() so mypy can infer them through the abstract return
type.

Adds tests covering: dict Overwrite in update and in writes replay,
snapshot_write with a dict reducer, dict backwards-compat checkpoints,
NotRequired type inference, and a filesystem-shaped end-to-end graph.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
2026-04-22 17:58:36 -04:00
Sydney Runkle 9d8c0be068 arbitrary 2026-04-22 17:54:12 -04:00
Sydney Runkle 51154be4ab lint 2026-04-22 17:47:53 -04:00
Sydney Runkle 2e7edb2b60 lint 2026-04-22 17:30:22 -04:00
Sydney Runkle 325cb42f19 lint and snapshot every 2026-04-22 17:06:53 -04:00
Sydney Runkle b9fad696ec cleanup 2026-04-22 16:45:32 -04:00
Sydney Runkle 96760e6267 chore(delta-channel): add PR description 2026-04-22 14:35:48 -04:00
Sydney RunkleandClaude Sonnet 4.6 9342ae215a chore(delta-channel): remove snapshot_every — simpler design, better storage savings
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-04-22 14:12:57 -04:00
Sydney Runkle ee5b3c2639 refactor(postgres): replace recursive CTE with two-query ancestor walk for DeltaChannel
Instead of a recursive SQL CTE, collect the ancestor checkpoint ID chain in
Python by fetching all (checkpoint_id, parent_checkpoint_id) for the thread
in one query, then fetch writes with a plain WHERE checkpoint_id = ANY(...).

Simpler, avoids recursive query planner overhead, and uses well-indexed lookups.
2026-04-22 14:03:37 -04:00
Sydney Runkle d7c3616620 feat(delta-channel): store sentinel in blobs, reconstruct from checkpoint_writes
DeltaChannel.checkpoint() now returns a zero-byte DeltaChannelSentinel
instead of duplicating delta data in checkpoint_blobs. Reconstruction
walks the parent checkpoint chain via checkpoint_writes (which already
holds per-step writes) and replays them through the operator.

In-memory benchmark (100 turns, ~20K tokens):
  storage: 10.2 MB → 40.5 KB (251x reduction)
  read:    0.6ms → 7.9ms (reconstruction cost, amortized by storage savings)

InMemorySaver and PostgresSaver override get_channel_writes() with
efficient implementations (Python dict walk and recursive CTE respectively).
The base class fallback uses self.list() with a thread-local recursion guard.
2026-04-22 14:03:37 -04:00
Sydney Runkle 06e302bbff refactor(delta-channel): infer typ from Annotated outer type instead of constructor arg
Remove the positional `typ` parameter from `DeltaChannel.__init__`. The type is
now injected automatically from the `Annotated` outer type in `_is_field_channel`
(matching how `BinaryOperatorAggregate` receives its type). `copy()` and
`from_checkpoint()` propagate `self.typ` explicitly. Test helpers updated to
use `_get_channel` with the proper `Annotated` path.
2026-04-22 14:03:37 -04:00
Sydney Runkle 10da2326a9 chore(delta-channel): remove supports_delta_channels flag
Rely on the runtime raise in DeltaChannel.from_checkpoint() instead of
a compile-time boolean flag. Savers that assemble DeltaChainValue inside
_load_blobs work transparently; savers that don't will pass through a raw
DeltaValue and hit a clear ValueError on first reload.

Removes: BaseCheckpointSaver.supports_delta_channels, the attribute on
InMemorySaver / PostgresSaver / AsyncPostgresSaver, the compile-time
UserWarning in StateGraph.compile(), and the associated test.
2026-04-22 14:03:37 -04:00
Sydney RunkleandClaude Sonnet 4.6 84f00c51eb chore: remove docs/ from PR
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-04-22 14:03:37 -04:00
ccurmeandSydney Runkle 3cab106bdf fix(prebuilt): handle injected NotRequired keys (#7392)
Resolves https://github.com/langchain-ai/langchain/issues/35585

This would previously raise KeyError:
```python
from typing import Annotated

from langchain_core.tools import tool
from langchain.agents import create_agent
from typing_extensions import NotRequired
from langgraph.prebuilt import InjectedState
from langchain.agents import AgentState


class CustomAgentState(AgentState):
    city: NotRequired[str]


@tool
def get_weather(city: Annotated[str | None, InjectedState("city")] = None) -> str:
    """Get weather for a given city."""
    if city is None:
        city = "Boston"
    return f"It's always sunny in {city}!"


agent = create_agent(
    model="claude-sonnet-4-6",
    tools=[get_weather],
    system_prompt="You are a helpful assistant",
    state_schema=CustomAgentState,
)

input_message = {
    "role": "user",
    "content": "What's the weather?",
}

result = agent.invoke({"messages": [input_message]})
for m in result["messages"]:
    m.pretty_print()
```

---------

Co-authored-by: Sydney Runkle <sydneymarierunkle@gmail.com>
2026-04-22 14:03:37 -04:00
Sydney Runkle 77bb349309 lint 2026-04-22 14:03:37 -04:00
Sydney Runkle 1d8364a749 latest 2026-04-22 14:03:37 -04:00
Sydney Runkle 06e97b0fd2 fix(delta-channel): support non-list reducers (dict) and fix MISSING handling
Use typ() instead of [] throughout DeltaChannel so reducers over dict
(and other non-list types) work correctly. fromCheckpoint(MISSING) now
leaves value as typ() from __init__ instead of overwriting with MISSING.
copy() uses value.copy() to handle dicts. update() initialises base from
typ() when value is MISSING. Add four tests covering the deepagents-style
dict-merge / file-deletion reducer pattern.
2026-04-22 14:03:37 -04:00
Sydney RunkleandClaude Sonnet 4.6 e256a31d00 chore(delta): rename _steps_since_rehydrate → _steps_since_snapshot; add audit tests
- Rename `_steps_since_rehydrate` → `_steps_since_snapshot` in DeltaChannel
  for clarity (counts steps since the last snapshot, not since rehydration)
- Pre-seed cycle-detection `visited` set with current checkpoint ID in both
  sync and async `_assemble_delta_channels` to prevent self-referential chains
- Add 4 new unit tests:
  - `test_delta_channel_snapshot_every_emits_plain_list`: verifies counter
    semantics and snapshot/delta transitions
  - `test_delta_channel_snapshot_every_end_to_end`: graph-level smoke test
  - `test_delta_channel_assembly_fast_path_returns_delta_value`: exercises
    chain traversal via get_channel_blob returning DeltaValue then plain list
  - `test_delta_channel_assembly_broken_chain_logs_warning`: partial chain
    when get_tuple returns None

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-04-22 14:03:37 -04:00
Sydney RunkleandClaude Sonnet 4.6 04b3ae7cd0 chore: rename serde type tag "diff" → "delta" for DeltaValue
Consistent with channel/type naming (DeltaChannel, DeltaValue).

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-04-22 14:03:37 -04:00
Sydney RunkleandClaude Sonnet 4.6 ec52520389 chore: apply format/lint fixes across checkpoint, checkpoint-postgres, prebuilt
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-04-22 14:03:37 -04:00
Sydney RunkleandClaude Sonnet 4.6 82fea763c5 fix: register DeltaValue in SAFE_MSGPACK_TYPES; rename _is_diff_delta; cross-saver benchmark
- Add DeltaValue to SAFE_MSGPACK_TYPES so SQLite and other msgpack-based
  savers don't emit "Deserializing unregistered type" warnings.
- Rename _is_diff_delta → _is_delta_value (leftover from DiffChannel rename).
- Parametrize benchmark by checkpointer: runs InMemory (fast-path) and
  SQLite (get_tuple fallback) in the same table, sharing the _run_turns helper.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-04-22 14:03:37 -04:00
Sydney RunkleandClaude Sonnet 4.6 6cfcad18f3 fix(delta-channel): fix chain assembly and get_state paths
- Fix InMemorySaver.get_channel_blob: use correct storage[thread_id][ns]
  nesting and deserialize the checkpoint before extracting channel_versions.
- Pass checkpoint_id to after_checkpoint() in channels_from_checkpoint so
  DeltaChannel seeds _last_checkpoint_id correctly on load; without this
  every turn broke the chain at its boundary.
- Wire _assemble_delta_channels into _prepare_state_snapshot and
  _aprepare_state_snapshot (get_state / get_state_history paths) and into
  perform_superstep / aperform_superstep (update_state paths) — previously
  only the loop __enter__ path did assembly.
- Fix test_get_channel_blob to use the correct storage structure.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-04-22 14:03:37 -04:00
Sydney Runkle db6b9ec995 test(channels): replace unsupported-saver raise test with fallback assembly test 2026-04-22 14:03:37 -04:00
Sydney Runkle 55fdc7aec6 feat(postgres): remove _load_diff_chains; add get_channel_blob / aget_channel_blob 2026-04-22 14:03:37 -04:00
Sydney Runkle 9969fb9737 feat(memory): implement get_channel_blob; remove diff handling from _load_blobs 2026-04-22 14:03:37 -04:00
Sydney Runkle 40981fdac7 feat(pregel): wire DeltaChannel assembly into loop; pass checkpoint_id to after_checkpoint 2026-04-22 14:03:37 -04:00
Sydney Runkle 5d1b3c4190 feat(pregel): add _assemble_delta_channels helpers for universal DeltaChannel support 2026-04-22 14:03:37 -04:00
Sydney RunkleandClaude Sonnet 4.6 fa83d64eff feat(channels): DeltaChannel tracks checkpoint_id; emits prev_checkpoint_id in DeltaValue
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-04-22 14:03:37 -04:00
Sydney RunkleandClaude Sonnet 4.6 4608af9615 feat(serde): diff type encodes prev_checkpoint_id; loads_typed returns DeltaValue
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-04-22 14:03:37 -04:00
Sydney Runkle 9beda5d3fb docs(checkpoint): expand aget_channel_blob docstring for parity 2026-04-22 14:03:37 -04:00
Sydney Runkle 9a6d7e08fb chore: add .worktrees/ to .gitignore 2026-04-22 14:03:37 -04:00
Sydney Runkle bbeb2759ba feat(checkpoint): DeltaValue uses prev_checkpoint_id; add get_channel_blob stubs 2026-04-22 14:03:37 -04:00
Sydney Runkle b799b95138 chore: rename DiffChannel/DiffDelta/DiffChainValue to Delta* across libs
Renames the diff-channel types to DeltaChannel, DeltaValue, and DeltaChainValue
for consistency with the settled naming convention.
2026-04-22 14:03:37 -04:00
Sydney Runkle ed9711fd33 more tests 2026-04-22 14:03:37 -04:00
Sydney RunkleandClaude Sonnet 4.6 ca883fe4eb chore: format/lint fixes for rehydrate_every benchmark
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-04-22 14:03:37 -04:00
Sydney RunkleandClaude Sonnet 4.6 4b9e4d25ca feat(channels): add rehydrate_every to DiffChannel for bounded chain traversal
Periodic full-snapshot checkpoints cap chain depth, trading a small
amount of extra storage for bounded reconstruction time.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-04-22 14:03:37 -04:00
Sydney RunkleandClaude Sonnet 4.6 ec8fd85ea2 test(channels): add DiffChannel vs BinaryOperatorAggregate storage/time benchmark
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-04-22 14:03:36 -04:00
Sydney RunkleandClaude Sonnet 4.6 ebd98f2e27 chore: format and lint fixes for DiffChannel implementation
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-04-22 14:03:36 -04:00
Sydney RunkleandClaude Sonnet 4.6 318fee9fc6 fix(checkpoint/postgres): pass cursor to avoid deadlock in diff chain traversal
Fixes a critical deadlock that occurs when _load_diff_chains calls self._cursor()
from within _load_blobs while the outer _load_checkpoint_tuple already holds
self._cursor(). On bare (non-pool) connections, the threading.Lock is not
reentrant, causing a deadlock.

Solution: Pass the cursor as a parameter to _load_diff_chains and _load_blobs
instead of acquiring a new cursor within those methods. Updated _load_checkpoint_tuple
to acquire a cursor once at the top level and pass it through the call chain.

Changes:
- Updated _load_blobs signature to accept optional cur parameter
- Updated _load_diff_chains signature (base and implementations) to accept optional cur parameter
- Modified _load_checkpoint_tuple in PostgresSaver to acquire cursor and pass it
- Modified _load_checkpoint_tuple_async to acquire cursor only when diff_payloads exist
- Removed nested self._cursor() calls in _load_diff_chains and _load_diff_chains_async

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-04-22 14:03:36 -04:00
Sydney RunkleandClaude Sonnet 4.6 e37299af87 feat(checkpoint/postgres): diff chain reconstruction in async saver
Add `_load_diff_chains_async` to `AsyncPostgresSaver` and override
`_load_checkpoint_tuple` to inline blob-parsing and diff-chain
resolution via async point-lookup traversal, mirroring the sync
`PostgresSaver._load_diff_chains` implementation. Add integration test
`test_diff_channel_chain_reconstruction` that skips gracefully when
`langgraph` core is not installed.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-04-22 14:03:36 -04:00
Sydney RunkleandClaude Sonnet 4.6 65d6ab2609 feat(checkpoint/postgres): diff chain reconstruction in _load_blobs (sync)
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-04-22 14:03:36 -04:00
Sydney RunkleandClaude Sonnet 4.6 888a308814 test(pregel): strengthen DiffChannel time-travel and reply assertions
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-04-22 14:03:36 -04:00
Sydney RunkleandClaude Sonnet 4.6 c7086ed7e2 feat(pregel): call after_checkpoint hook when loading and saving channels
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-04-22 14:03:36 -04:00
Sydney RunkleandClaude Sonnet 4.6 e645c2a085 fix(checkpoint/memory): warn on broken diff chain, guard against cycles
- Add logger.warning when a mid-chain blob is missing (fixes silent truncation bug)
- Add cycle guard to prevent infinite loops on corrupt blob stores
- Fix type annotation on diff_channels from dict[str, Any] to dict[str, str]

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-04-22 14:03:36 -04:00
Sydney RunkleandClaude Sonnet 4.6 e52b7b2d54 feat(checkpoint/memory): chain-traverse diff blobs in _load_blobs
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-04-22 14:03:36 -04:00
Sydney Runkle 6581fdd0a5 fix(channels): align DiffChannel.is_available with BinaryOperatorAggregate 2026-04-22 14:03:36 -04:00
Sydney RunkleandClaude Sonnet 4.6 fe2bc286fc feat(channels): implement DiffChannel for incremental checkpoint storage
Adds DiffChannel, a new channel type that stores only per-step write
deltas in checkpoints and reconstructs the full list by replaying the
chain through the operator at load time.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-04-22 14:03:36 -04:00
Sydney Runkle c345d337bb feat(channels): add no-op after_checkpoint hook to BaseChannel 2026-04-22 14:03:36 -04:00
Sydney RunkleandClaude Sonnet 4.6 d0af83b746 fix(checkpoint/serde): use lazy isinstance check for DiffDelta
Replace duck-typing check with lazy import inside _is_diff_delta helper
function to avoid module-level circular dependency while using proper
isinstance semantics.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-04-22 14:03:36 -04:00
Sydney RunkleandClaude Sonnet 4.5 eabf926a4f feat(checkpoint/serde): serialize DiffDelta as 'diff' type tag
Add serde support for DiffDelta by implementing dump/load for the "diff" type tag.
This allows the checkpoint system to efficiently store delta objects by serializing
them as msgpack-encoded dicts with {"d": delta, "p": prev_version} structure.

The implementation uses runtime type checking to avoid circular imports and
leverages the existing msgpack ext hooks for proper deserialization of complex
types like LangChain messages.

Co-Authored-By: Claude Sonnet 4.5 <noreply@anthropic.com>
2026-04-22 14:03:36 -04:00
Sydney RunkleandClaude Sonnet 4.6 6852e0478e feat(checkpoint): add DiffDelta and DiffChainValue protocol types
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-04-22 14:03:36 -04:00
Sydney RunkleandClaude Sonnet 4.6 68b2f7bfac docs: add DiffChannel implementation plan
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-04-22 14:03:36 -04:00
Sydney RunkleandClaude Sonnet 4.6 95e0fe060c docs: add DiffChannel incremental checkpoint storage design spec
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-04-22 14:03:36 -04:00
51 changed files with 556 additions and 3147 deletions
@@ -19,7 +19,6 @@ from langgraph.checkpoint.base import (
get_serializable_checkpoint_metadata,
)
from langgraph.checkpoint.serde.base import SerializerProtocol
from langgraph.checkpoint.serde.types import _DeltaSnapshot
from psycopg import Capabilities, Connection, Cursor, Pipeline
from psycopg.rows import DictRow, dict_row
from psycopg.types.json import Jsonb
@@ -27,11 +26,9 @@ from psycopg_pool import ConnectionPool
from langgraph.checkpoint.postgres import _internal
from langgraph.checkpoint.postgres.base import (
SELECT_DELTA_STAGE1_SQL,
SELECT_DELTA_STAGE2_SQL,
SELECT_DELTA_COMBINED_SQL,
BasePostgresSaver,
_DeltaStage1Row,
_DeltaStage2Row,
_DeltaCombinedRow,
)
from langgraph.checkpoint.postgres.shallow import ShallowPostgresSaver
@@ -311,12 +308,7 @@ class PostgresSaver(BasePostgresSaver):
# others are stored in blobs table
blob_values = {}
for k, v in checkpoint["channel_values"].items():
if v is DELTA_SENTINEL:
copy["channel_values"].pop(k)
elif isinstance(v, _DeltaSnapshot):
blob_values[k] = copy["channel_values"].pop(k)
copy["channel_values"][k] = True
elif v is None or isinstance(v, (str, int, float, bool)):
if v is None or isinstance(v, (str, int, float, bool)):
pass
else:
blob_values[k] = copy["channel_values"].pop(k)
@@ -449,49 +441,41 @@ class PostgresSaver(BasePostgresSaver):
) -> _ChannelWritesHistory:
"""Fast-path override of `BaseCheckpointSaver._get_channel_writes_history`.
Two-stage query: stage 1 scans checkpoint metadata to walk the parent
chain and locate the nearest snapshot; stage 2 fetches only the
chain-limited writes and single seed blob.
One combined UNION ALL query (`SELECT_DELTA_COMBINED_SQL`) fetches rows
from `checkpoints`, `checkpoint_writes`, and `checkpoint_blobs` in a
single roundtrip; the ancestor walk runs in Python.
"""
thread_id = config["configurable"]["thread_id"]
checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
checkpoint_id = get_checkpoint_id(config)
if checkpoint_id is None:
# Caller didn't specify a target — resolve to the latest
# checkpoint on the thread. `get_tuple` without `checkpoint_id`
# returns the newest; its config carries the resolved id.
target = self.get_tuple(config)
if target is None:
return _ChannelWritesHistory(seed=DELTA_SENTINEL, writes=[])
checkpoint_id = target.config["configurable"]["checkpoint_id"]
with self._cursor() as cur:
cur.execute(
SELECT_DELTA_STAGE1_SQL,
(channel, channel, thread_id, checkpoint_ns),
)
stage1_rows = cur.fetchall()
chain_cids, seed_version = self._walk_stage1(
cast("list[_DeltaStage1Row]", stage1_rows), checkpoint_id
)
seed_versions = [seed_version] if seed_version else []
with self._cursor() as cur:
cur.execute(
SELECT_DELTA_STAGE2_SQL,
SELECT_DELTA_COMBINED_SQL,
(
channel,
thread_id,
checkpoint_ns,
thread_id,
checkpoint_ns,
channel,
chain_cids,
thread_id,
checkpoint_ns,
channel,
seed_versions,
),
)
stage2_rows = cur.fetchall()
rows = cur.fetchall()
return self._build_delta_channel_writes_history(
channel=channel,
chain_cids=chain_cids,
seed_version=seed_version,
stage2_rows=cast("list[_DeltaStage2Row]", stage2_rows),
target_id=checkpoint_id,
rows=cast("list[_DeltaCombinedRow]", rows),
)
def _load_checkpoint_tuple(self, value: DictRow) -> CheckpointTuple:
@@ -19,7 +19,6 @@ from langgraph.checkpoint.base import (
get_serializable_checkpoint_metadata,
)
from langgraph.checkpoint.serde.base import SerializerProtocol
from langgraph.checkpoint.serde.types import _DeltaSnapshot
from psycopg import AsyncConnection, AsyncCursor, AsyncPipeline, Capabilities
from psycopg.rows import DictRow, dict_row
from psycopg.types.json import Jsonb
@@ -27,11 +26,9 @@ from psycopg_pool import AsyncConnectionPool
from langgraph.checkpoint.postgres import _ainternal
from langgraph.checkpoint.postgres.base import (
SELECT_DELTA_STAGE1_SQL,
SELECT_DELTA_STAGE2_SQL,
SELECT_DELTA_COMBINED_SQL,
BasePostgresSaver,
_DeltaStage1Row,
_DeltaStage2Row,
_DeltaCombinedRow,
)
from langgraph.checkpoint.postgres.shallow import AsyncShallowPostgresSaver
@@ -270,12 +267,7 @@ class AsyncPostgresSaver(BasePostgresSaver):
# others are stored in blobs table
blob_values = {}
for k, v in checkpoint["channel_values"].items():
if v is DELTA_SENTINEL:
copy["channel_values"].pop(k)
elif isinstance(v, _DeltaSnapshot):
blob_values[k] = copy["channel_values"].pop(k)
copy["channel_values"][k] = True
elif v is None or isinstance(v, (str, int, float, bool)):
if v is None or isinstance(v, (str, int, float, bool)):
pass
else:
blob_values[k] = copy["channel_values"].pop(k)
@@ -410,9 +402,10 @@ class AsyncPostgresSaver(BasePostgresSaver):
) -> _ChannelWritesHistory:
"""Fast-path override of `BaseCheckpointSaver._aget_channel_writes_history`.
Two-stage query: stage 1 scans checkpoint metadata to walk the parent
chain and locate the nearest snapshot; stage 2 fetches only the
chain-limited writes and single seed blob.
One combined UNION ALL query (`SELECT_DELTA_COMBINED_SQL`) fetches rows
from `checkpoints`, `checkpoint_writes`, and `checkpoint_blobs` in a
single roundtrip; rows are assembled by the shared pure helper on
`BasePostgresSaver`.
"""
thread_id = config["configurable"]["thread_id"]
checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
@@ -422,37 +415,26 @@ class AsyncPostgresSaver(BasePostgresSaver):
if target is None:
return _ChannelWritesHistory(seed=DELTA_SENTINEL, writes=[])
checkpoint_id = target.config["configurable"]["checkpoint_id"]
async with self._cursor() as cur:
await cur.execute(
SELECT_DELTA_STAGE1_SQL,
(channel, channel, thread_id, checkpoint_ns),
)
stage1_rows = await cur.fetchall()
chain_cids, seed_version = self._walk_stage1(
cast("list[_DeltaStage1Row]", stage1_rows), checkpoint_id
)
seed_versions = [seed_version] if seed_version else []
async with self._cursor() as cur:
await cur.execute(
SELECT_DELTA_STAGE2_SQL,
SELECT_DELTA_COMBINED_SQL,
(
channel,
thread_id,
checkpoint_ns,
thread_id,
checkpoint_ns,
channel,
chain_cids,
thread_id,
checkpoint_ns,
channel,
seed_versions,
),
)
stage2_rows = await cur.fetchall()
rows = await cur.fetchall()
return self._build_delta_channel_writes_history(
channel=channel,
chain_cids=chain_cids,
seed_version=seed_version,
stage2_rows=cast("list[_DeltaStage2Row]", stage2_rows),
target_id=checkpoint_id,
rows=cast("list[_DeltaCombinedRow]", rows),
)
async def _load_checkpoint_tuple(self, value: DictRow) -> CheckpointTuple:
@@ -156,62 +156,62 @@ INSERT_CHECKPOINT_WRITES_SQL = """
"""
class _DeltaStage2Row(TypedDict, total=False):
"""One row from `SELECT_DELTA_STAGE2_SQL` (a UNION ALL of writes and blobs)."""
class _DeltaCombinedRow(TypedDict, total=False):
"""One row from `SELECT_DELTA_COMBINED_SQL` (a UNION ALL of three tables).
_kind: str # "w" or "b"
checkpoint_id: str | None # "w" rows only
Every row carries `_kind` ("p" / "w" / "b") plus whichever columns are
relevant for that kind; irrelevant columns are NULL and typed as `None`.
"""
_kind: str # always present: "p", "w", or "b"
# checkpoint row ("p")
checkpoint_id: str | None
parent_checkpoint_id: str | None
ver: str | None
# write / blob rows ("w", "b")
type: str | None
blob: bytes | None
task_id: str | None # "w" rows only
idx: int | None # "w" rows only
version: str | None # "b" rows only
# write row only ("w")
task_id: str | None
idx: int | None
# blob row only ("b")
version: str | None
# Two-stage DeltaChannel reconstruction. Stage 1 scans checkpoint
# metadata (no blob bytes) to walk the parent chain and locate the
# nearest snapshot marker. Stage 2 fetches only the chain-limited
# writes and the single seed snapshot blob.
# DeltaChannel reconstruction: one UNION ALL query fetches checkpoints,
# writes, and blobs for `channel` in one roundtrip; the ancestor walk runs
# in Python in `_build_delta_channel_writes_history`.
#
# Parameter order:
# stage1: (channel, channel, thread_id, checkpoint_ns)
# stage2: (thread_id, checkpoint_ns, channel, chain_cids[],
# thread_id, checkpoint_ns, channel, seed_versions[])
SELECT_DELTA_STAGE1_SQL = """
SELECT checkpoint_id,
# Parameter order: (channel, thread_id, checkpoint_ns,
# thread_id, checkpoint_ns, channel,
# thread_id, checkpoint_ns, channel)
SELECT_DELTA_COMBINED_SQL = """
SELECT 'p'::text AS _kind,
checkpoint_id,
parent_checkpoint_id,
checkpoint -> 'channel_versions' ->> %s AS ver,
(checkpoint -> 'channel_values' -> %s) IS NOT NULL AS has_snapshot
NULL::text AS type,
NULL::bytea AS blob,
NULL::text AS task_id,
NULL::int AS idx,
NULL::text AS version
FROM checkpoints
WHERE thread_id = %s AND checkpoint_ns = %s
"""
SELECT_DELTA_STAGE2_SQL = """
SELECT 'w'::text AS _kind,
checkpoint_id,
type, blob, task_id, idx, NULL::text AS version
UNION ALL
SELECT 'w',
checkpoint_id, NULL, NULL,
type, blob, task_id, idx, NULL
FROM checkpoint_writes
WHERE thread_id = %s AND checkpoint_ns = %s AND channel = %s
AND checkpoint_id = ANY(%s)
UNION ALL
SELECT 'b', NULL,
SELECT 'b',
NULL, NULL, NULL,
type, blob, NULL, NULL, version
FROM checkpoint_blobs
WHERE thread_id = %s AND checkpoint_ns = %s AND channel = %s
AND version = ANY(%s)
"""
class _DeltaStage1Row(TypedDict):
"""One row from `SELECT_DELTA_STAGE1_SQL`."""
checkpoint_id: str
parent_checkpoint_id: str | None
ver: str | None
has_snapshot: bool
class BasePostgresSaver(BaseCheckpointSaver[str]):
SELECT_SQL = SELECT_SQL
SELECT_PENDING_SENDS_SQL = SELECT_PENDING_SENDS_SQL
@@ -254,59 +254,38 @@ class BasePostgresSaver(BaseCheckpointSaver[str]):
if t.decode() != "empty"
}
@staticmethod
def _walk_stage1(
stage1_rows: Sequence[_DeltaStage1Row],
target_id: str,
) -> tuple[list[str], str | None]:
"""Walk the parent chain from stage 1 metadata rows.
Returns (chain_cids, seed_version):
chain_cids: ancestor checkpoint IDs from target's parent down to
the seed (or root), in newest-first order.
seed_version: the channel blob version at the nearest ancestor
with has_snapshot=True, or None if pure delta.
"""
parent_of: dict[str, str | None] = {}
ver_of: dict[str, str | None] = {}
snapshot_of: dict[str, bool] = {}
for r in stage1_rows:
cid = r["checkpoint_id"]
parent_of[cid] = r["parent_checkpoint_id"]
ver_of[cid] = r["ver"]
snapshot_of[cid] = r["has_snapshot"]
chain_cids: list[str] = []
seed_version: str | None = None
cur_cid: str | None = parent_of.get(target_id)
while cur_cid is not None:
chain_cids.append(cur_cid)
if snapshot_of.get(cur_cid, False):
seed_version = ver_of.get(cur_cid)
break
cur_cid = parent_of.get(cur_cid)
return chain_cids, seed_version
def _build_delta_channel_writes_history(
self,
*,
channel: str,
chain_cids: list[str],
seed_version: str | None,
stage2_rows: Sequence[_DeltaStage2Row],
target_id: str,
rows: Sequence[_DeltaCombinedRow],
) -> _ChannelWritesHistory:
"""Reconstruct delta channel history from two-stage query results.
"""Reconstruct one delta channel's history from the combined UNION ALL rows.
chain_cids are in newest-first order (target's parent first).
stage2_rows contain only writes for chain_cids and the single
seed blob at seed_version.
Pure data transform shared by sync (`PostgresSaver`) and async
(`AsyncPostgresSaver`); both paths run `SELECT_DELTA_COMBINED_SQL`
and feed the tagged rows here.
Walk is newest → oldest from the target's parent. A non-sentinel
blob in `checkpoint_blobs` (a pre-delta snapshot) terminates the
walk and is returned as the seed so replay starts from it.
Writes stored at `target_id` itself are pending writes for the next
step and are excluded — the walk begins at the target's parent.
"""
parent_of: dict[str, str | None] = {}
ver_of: dict[str, str | None] = {}
writes_by_cid: dict[str, list[tuple[str, bytes, str, int]]] = {}
seed_blob: tuple[str, bytes] | None = None
blob_by_ver: dict[str, tuple[str, bytes]] = {}
for r in stage2_rows:
for r in rows:
kind = r["_kind"]
if kind == "w":
if kind == "p":
cid = cast(str, r["checkpoint_id"])
parent_of[cid] = r["parent_checkpoint_id"]
ver_of[cid] = r["ver"]
elif kind == "w":
cid = cast(str, r["checkpoint_id"])
writes_by_cid.setdefault(cid, []).append(
cast(
@@ -315,26 +294,42 @@ class BasePostgresSaver(BaseCheckpointSaver[str]):
)
)
else: # kind == "b"
seed_blob = cast("tuple[str, bytes]", (r["type"], r["blob"]))
blob_by_ver[cast(str, r["version"])] = cast(
"tuple[str, bytes]", (r["type"], r["blob"])
)
# newest write first per ancestor (task_id DESC, idx DESC)
for ws in writes_by_cid.values():
ws.sort(key=lambda w: (w[2], w[3]), reverse=True)
if not chain_cids:
ancestors: list[str] = []
cur_cid: str | None = parent_of.get(target_id)
while cur_cid is not None:
ancestors.append(cur_cid)
cur_cid = parent_of.get(cur_cid)
if not ancestors:
return _ChannelWritesHistory(seed=DELTA_SENTINEL, writes=[])
collected: list[PendingWrite] = []
for cid in chain_cids:
collected: list[PendingWrite] = [] # newest first; reversed at the end
for cid in ancestors:
# Collect writes first — they encode the transition FROM this
# ancestor's state to its child's and must be included even if
# this ancestor is also the seed checkpoint.
for type_tag, write_blob, task_id, _idx in writes_by_cid.get(cid, []):
val = self.serde.loads_typed((type_tag, write_blob))
collected.append((task_id, channel, val))
# Then check seed terminator.
ver = ver_of.get(cid)
if ver is not None:
seed_blob = blob_by_ver.get(ver)
if seed_blob is not None and seed_blob[0] != "empty":
blob_value = self.serde.loads_typed(seed_blob)
if blob_value is not DELTA_SENTINEL:
collected.reverse()
return _ChannelWritesHistory(seed=blob_value, writes=collected)
seed: Any = DELTA_SENTINEL
if seed_blob is not None and seed_blob[0] != "empty":
seed = self.serde.loads_typed(seed_blob)
collected.reverse()
return _ChannelWritesHistory(seed=seed, writes=collected)
collected.reverse() # oldest → newest
return _ChannelWritesHistory(seed=DELTA_SENTINEL, writes=collected)
def _dump_blobs(
self,
+2 -2
View File
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
[project]
name = "langgraph-checkpoint-postgres"
version = "3.1.0a3"
version = "3.0.5"
description = "Library with a Postgres implementation of LangGraph checkpoint saver."
authors = []
requires-python = ">=3.10"
@@ -12,7 +12,7 @@ readme = "README.md"
license = "MIT"
license-files = ['LICENSE']
dependencies = [
"langgraph-checkpoint>=4.1.0a3,<5.0.0",
"langgraph-checkpoint>=4.0.3,<5.0.0",
"orjson>=3.11.5",
"psycopg>=3.2.0",
"psycopg-pool>=3.2.0",
+2 -2
View File
@@ -259,7 +259,7 @@ wheels = [
[[package]]
name = "langgraph-checkpoint"
version = "4.1.0a3"
version = "4.0.3"
source = { editable = "../checkpoint" }
dependencies = [
{ name = "langchain-core" },
@@ -307,7 +307,7 @@ test = [
[[package]]
name = "langgraph-checkpoint-postgres"
version = "3.1.0a3"
version = "3.0.5"
source = { editable = "." }
dependencies = [
{ name = "langgraph-checkpoint" },
+1 -1
View File
@@ -268,7 +268,7 @@ wheels = [
[[package]]
name = "langgraph-checkpoint"
version = "4.1.0a3"
version = "4.0.3"
source = { editable = "../checkpoint" }
dependencies = [
{ name = "langchain-core" },
@@ -452,9 +452,7 @@ class InMemorySaver(
values: dict[str, Any] = c.pop("channel_values") # type: ignore[misc]
for k, v in new_versions.items():
self.blobs[(thread_id, checkpoint_ns, k, v)] = (
self.serde.dumps_typed(values[k])
if k in values and values[k] is not DELTA_SENTINEL
else ("empty", b"")
self.serde.dumps_typed(values[k]) if k in values else ("empty", b"")
)
self.storage[thread_id][checkpoint_ns].update(
{
@@ -34,7 +34,9 @@ from langgraph.checkpoint.serde import _msgpack as _lg_msgpack
from langgraph.checkpoint.serde.base import SerializerProtocol
from langgraph.checkpoint.serde.event_hooks import emit_serde_event
from langgraph.checkpoint.serde.types import (
DELTA_SENTINEL,
SendProtocol,
_DeltaSentinel,
_DeltaSnapshot,
)
from langgraph.store.base import Item
@@ -320,11 +322,14 @@ EXT_PYDANTIC_V1 = 4
EXT_PYDANTIC_V2 = 5
EXT_NUMPY_ARRAY = 6
EXT_DELTA_SNAPSHOT = 7
EXT_DELTA_SENTINEL = 8
def _msgpack_default(obj: Any) -> str | ormsgpack.Ext:
if isinstance(obj, _DeltaSnapshot):
return ormsgpack.Ext(EXT_DELTA_SNAPSHOT, _msgpack_enc(obj.value))
elif isinstance(obj, _DeltaSentinel):
return ormsgpack.Ext(EXT_DELTA_SENTINEL, b"")
elif hasattr(obj, "model_dump") and callable(obj.model_dump): # pydantic v2
return ormsgpack.Ext(
EXT_PYDANTIC_V2,
@@ -651,7 +656,9 @@ def _create_msgpack_ext_hook(
return False
def ext_hook(code: int, data: bytes) -> Any:
if code == EXT_DELTA_SNAPSHOT:
if code == EXT_DELTA_SENTINEL:
return DELTA_SENTINEL
elif code == EXT_DELTA_SNAPSHOT:
return _DeltaSnapshot(
ormsgpack.unpackb(
data, ext_hook=ext_hook, option=ormsgpack.OPT_NON_STR_KEYS
@@ -17,10 +17,12 @@ TASKS = "__pregel_tasks"
class _DeltaSentinel:
"""In-memory marker for a DeltaChannel field with no snapshot.
"""Singleton marker stored (as zero bytes) in checkpoint_blobs for a
DeltaChannel field. The actual per-step writes live in checkpoint_writes
and are replayed through the reducer at load time.
Never serialized to storage — checkpointers strip it before writing.
Compare with `is DELTA_SENTINEL`; always the same module-level instance.
Compare with `is DELTA_SENTINEL` — `loads_typed` always returns the same
module-level instance.
"""
__slots__ = ()
+1 -1
View File
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
[project]
name = "langgraph-checkpoint"
version = "4.1.0a3"
version = "4.0.3"
description = "Library with base interfaces for LangGraph checkpoint savers."
authors = []
requires-python = ">=3.10"
+12
View File
@@ -1048,3 +1048,15 @@ def test_msgpack_nested_pydantic_serializes_as_dict(
# No blocking should occur - inner is serialized as dict, not ext
assert "blocked" not in caplog.text.lower()
assert result == obj
def test_delta_sentinel_serde_round_trip() -> None:
from langgraph.checkpoint.base import DELTA_SENTINEL
from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
serde = JsonPlusSerializer()
type_tag, blob = serde.dumps_typed(DELTA_SENTINEL)
assert type_tag == "msgpack"
assert blob # non-empty ext envelope
loaded = serde.loads_typed((type_tag, blob))
assert loaded is DELTA_SENTINEL
+16 -7
View File
@@ -322,17 +322,26 @@ def test_memory_saver_with_allowlist_proxy_isolated() -> None:
class TestInMemorySaverDeltaChannel:
def test_load_blobs_omits_delta_channel(self) -> None:
"""_load_blobs omits delta channels (stored as 'empty'); reconstruction deferred."""
def test_load_blobs_returns_sentinel_for_delta_channel(self) -> None:
"""_load_blobs returns DELTA_SENTINEL for delta channels (reconstruction deferred)."""
saver = InMemorySaver()
serde = JsonPlusSerializer()
thread_id, ns, channel = "t1", "", "messages"
v1 = "00000000000000000000000000000001.0000000000000000"
saver.blobs[(thread_id, ns, channel, v1)] = ("empty", b"")
saver.blobs[(thread_id, ns, channel, v1)] = serde.dumps_typed(DELTA_SENTINEL)
cp1 = empty_checkpoint()
cp1["id"] = "cp1"
cp1["channel_versions"][channel] = v1
saver.storage[thread_id][ns] = {
"cp1": (serde.dumps_typed(cp1), serde.dumps_typed({}), None),
}
result = saver._load_blobs(thread_id, ns, {channel: v1})
assert channel not in result
assert channel in result
assert result[channel] is DELTA_SENTINEL
def test_get_channel_writes_collects_ancestor_writes_only(self) -> None:
"""_get_channel_writes_history collects ancestor writes oldest→newest,
@@ -573,9 +582,9 @@ class TestPreDeltaBlobTerminator:
# Pre-delta: cp1 stored a real blob for the channel.
saver.blobs[(thread_id, ns, channel, v1)] = serde.dumps_typed(["A"])
# Delta-era: cp2 and cp3 store "empty"; real writes in checkpoint_writes.
saver.blobs[(thread_id, ns, channel, v2)] = ("empty", b"")
saver.blobs[(thread_id, ns, channel, v3)] = ("empty", b"")
# Delta-era: cp2 and cp3 store sentinels; real writes in checkpoint_writes.
saver.blobs[(thread_id, ns, channel, v2)] = serde.dumps_typed(DELTA_SENTINEL)
saver.blobs[(thread_id, ns, channel, v3)] = serde.dumps_typed(DELTA_SENTINEL)
cp1 = empty_checkpoint()
cp1["id"] = "cp1"
+1 -1
View File
@@ -286,7 +286,7 @@ wheels = [
[[package]]
name = "langgraph-checkpoint"
version = "4.1.0a3"
version = "4.0.3"
source = { editable = "." }
dependencies = [
{ name = "langchain-core" },
@@ -12,9 +12,6 @@ RESUME = sys.intern("__resume__")
# for values passed to resume a node after an interrupt
ERROR = sys.intern("__error__")
# for errors raised by nodes
ERROR_SOURCE_NODE = sys.intern("__error_source_node__")
# failed source node name for node-level error handlers
# value format in pending writes: `(task_id, ERROR_SOURCE_NODE, node_name: str)`
NO_WRITES = sys.intern("__no_writes__")
# marker to signal node didn't write anything
TASKS = sys.intern("__pregel_tasks")
@@ -74,10 +71,6 @@ CONFIG_KEY_RESUME_MAP = sys.intern("__pregel_resume_map")
CONFIG_KEY_STREAM_MESSAGES_V2 = sys.intern("__pregel_stream_messages_v2")
# when True, attach StreamMessagesHandlerV2 so content-block (v2) events
# flow through stream_mode="messages"; set by StreamingHandler only.
CONFIG_KEY_NODE_ERROR = sys.intern("__pregel_node_error")
# holds a `NodeError` (failed source node + exception) for the current
# node-level error handler invocation, injected when handler signature
# requests `error: NodeError`
# --- Other constants ---
PUSH = sys.intern("__pregel_push")
@@ -105,7 +98,6 @@ RESERVED = {
INTERRUPT,
RESUME,
ERROR,
ERROR_SOURCE_NODE,
NO_WRITES,
# reserved config.configurable keys
CONFIG_KEY_SEND,
@@ -51,11 +51,9 @@ from langgraph._internal._config import (
)
from langgraph._internal._constants import (
CONF,
CONFIG_KEY_NODE_ERROR,
CONFIG_KEY_RUNTIME,
)
from langgraph._internal._typing import MISSING
from langgraph.errors import NodeError
from langgraph.types import StreamWriter
try:
@@ -196,15 +194,6 @@ KWARGS_CONFIG_KEYS: tuple[tuple[str, tuple[Any, ...], str, Any], ...] = (
"N/A",
inspect.Parameter.empty,
),
(
"error",
(NodeError, "NodeError"),
# we never hit this block, we read directly from configurable
"N/A",
# default to None so non-handler nodes that happen to type a parameter
# `error: NodeError` don't blow up; handlers always receive a NodeError.
None,
),
)
"""List of kwargs that can be passed to functions, and their corresponding
config keys, default values and type annotations.
@@ -378,8 +367,6 @@ class RunnableCallable(Runnable):
kw_value: Any = MISSING
if kw == "config":
kw_value = config
elif kw == "error":
kw_value = config.get(CONF, {}).get(CONFIG_KEY_NODE_ERROR, MISSING)
elif runtime:
if kw == "runtime":
kw_value = runtime
@@ -452,8 +439,6 @@ class RunnableCallable(Runnable):
kw_value: Any = MISSING
if kw == "config":
kw_value = config
elif kw == "error":
kw_value = config.get(CONF, {}).get(CONFIG_KEY_NODE_ERROR, MISSING)
elif runtime:
if kw == "runtime":
kw_value = runtime
+8 -43
View File
@@ -1,7 +1,6 @@
from __future__ import annotations
from collections.abc import Sequence
from dataclasses import dataclass
from enum import Enum
from typing import Any, Literal
from warnings import warn
@@ -16,12 +15,10 @@ from langgraph.warnings import LangGraphDeprecatedSinceV10
__all__ = (
"EmptyChannelError",
"ErrorCode",
"GraphDrained",
"GraphRecursionError",
"InvalidUpdateError",
"GraphBubbleUp",
"GraphInterrupt",
"NodeError",
"NodeInterrupt",
"NodeTimeoutError",
"ParentCommand",
@@ -46,23 +43,6 @@ def create_error_message(*, message: str, error_code: ErrorCode) -> str:
)
class GraphBubbleUp(Exception):
pass
class GraphDrained(GraphBubbleUp):
"""Raised when a graph run exits early due to a drain request.
This indicates the graph stopped cooperatively at a superstep boundary
because `RunControl.request_drain()` was called (e.g., in response to
SIGTERM). The checkpoint is saved and the run can be resumed later.
"""
def __init__(self, reason: str = "shutdown") -> None:
self.reason = reason
super().__init__(f"Graph drained: {reason}")
class GraphRecursionError(RecursionError):
"""Raised when the graph has exhausted the maximum number of steps.
@@ -98,6 +78,10 @@ class InvalidUpdateError(Exception):
pass
class GraphBubbleUp(Exception):
pass
class GraphInterrupt(GraphBubbleUp):
"""Raised when a subgraph is interrupted, suppressed by the root graph.
Never raised directly, or surfaced to the user."""
@@ -144,31 +128,12 @@ class TaskNotFound(Exception):
pass
@dataclass(frozen=True, slots=True)
class NodeError:
"""Failure context passed to a node-level error handler.
Inject by adding a parameter typed `NodeError` to a handler registered via
`StateGraph.add_node(..., error_handler=...)`:
```python
def handler(state: State, error: NodeError) -> Command:
return Command(update={"status": f"recovered from {error.node}: {error.error}"})
```
"""
node: str
"""Name of the node whose execution failed."""
error: BaseException
"""Exception raised by the failed node."""
class NodeTimeoutError(Exception):
class NodeTimeoutError(TimeoutError):
"""Raised when a node invocation exceeds one of its configured timeouts.
Does **not** inherit from the built-in `TimeoutError` (a subclass of
`OSError`) so that the default `RetryPolicy` treats it as retryable.
Subclasses the built-in `TimeoutError`, so existing `except TimeoutError`
handlers keep working. If the node has a `retry_policy` whose `retry_on`
permits `TimeoutError`, the attempt will be retried.
Both `idle_timeout` and `run_timeout` reflect the configured policy at the
time of the failure (each is `None` if not configured). `kind` and
-2
View File
@@ -88,8 +88,6 @@ class StateNodeSpec(Generic[NodeInputT, ContextT]):
input_schema: type[NodeInputT]
retry_policy: RetryPolicy | Sequence[RetryPolicy] | None
cache_policy: CachePolicy | None
is_error_handler: bool = False
error_handler_node: str | None = None
ends: tuple[str, ...] | dict[str, str] | None = EMPTY_SEQ
defer: bool = False
timeout: TimeoutPolicy | None = None
+2 -37
View File
@@ -303,7 +303,6 @@ class StateGraph(Generic[StateT, ContextT, InputT, OutputT]):
input_schema: None = None,
retry_policy: RetryPolicy | Sequence[RetryPolicy] | None = None,
cache_policy: CachePolicy | None = None,
error_handler: StateNode[Any, ContextT] | None = None,
destinations: dict[str, str] | tuple[str, ...] | None = None,
timeout: float | timedelta | TimeoutPolicy | None = None,
**kwargs: Unpack[DeprecatedKwargs],
@@ -372,7 +371,6 @@ class StateGraph(Generic[StateT, ContextT, InputT, OutputT]):
input_schema: type[NodeInputT],
retry_policy: RetryPolicy | Sequence[RetryPolicy] | None = None,
cache_policy: CachePolicy | None = None,
error_handler: StateNode[Any, ContextT] | None = None,
destinations: dict[str, str] | tuple[str, ...] | None = None,
timeout: float | timedelta | TimeoutPolicy | None = None,
**kwargs: Unpack[DeprecatedKwargs],
@@ -446,7 +444,6 @@ class StateGraph(Generic[StateT, ContextT, InputT, OutputT]):
input_schema: None = None,
retry_policy: RetryPolicy | Sequence[RetryPolicy] | None = None,
cache_policy: CachePolicy | None = None,
error_handler: StateNode[Any, ContextT] | None = None,
destinations: dict[str, str] | tuple[str, ...] | None = None,
timeout: float | timedelta | TimeoutPolicy | None = None,
**kwargs: Unpack[DeprecatedKwargs],
@@ -515,7 +512,6 @@ class StateGraph(Generic[StateT, ContextT, InputT, OutputT]):
input_schema: type[NodeInputT],
retry_policy: RetryPolicy | Sequence[RetryPolicy] | None = None,
cache_policy: CachePolicy | None = None,
error_handler: StateNode[Any, ContextT] | None = None,
destinations: dict[str, str] | tuple[str, ...] | None = None,
timeout: float | timedelta | TimeoutPolicy | None = None,
**kwargs: Unpack[DeprecatedKwargs],
@@ -591,7 +587,6 @@ class StateGraph(Generic[StateT, ContextT, InputT, OutputT]):
input_schema: type[NodeInputT] | None = None,
retry_policy: RetryPolicy | Sequence[RetryPolicy] | None = None,
cache_policy: CachePolicy | None = None,
error_handler: StateNode[Any, ContextT] | None = None,
destinations: dict[str, str] | tuple[str, ...] | None = None,
timeout: float | timedelta | TimeoutPolicy | None = None,
**kwargs: Unpack[DeprecatedKwargs],
@@ -612,7 +607,6 @@ class StateGraph(Generic[StateT, ContextT, InputT, OutputT]):
If a sequence is provided, the first matching policy will be applied.
cache_policy: The cache policy for the node.
error_handler: Optional node-level error handler callable for this node.
destinations: Destinations that indicate where a node can route to.
Useful for edgeless graphs with nodes that return `Command` objects.
@@ -772,25 +766,6 @@ class StateGraph(Generic[StateT, ContextT, InputT, OutputT]):
if destinations is not None:
ends = destinations
resolved_input_schema: type[Any] = (
input_schema or inferred_input_schema or self.state_schema
)
handler_node_name: str | None = None
if error_handler is not None:
handler_node_name = f"__error_handler__{node}"
if handler_node_name in self.nodes:
raise ValueError(
f"Auto-generated error handler node `{handler_node_name}` already exists."
)
self.nodes[handler_node_name] = StateNodeSpec[Any, ContextT](
coerce_to_runnable(error_handler, name=handler_node_name, trace=False), # type: ignore[arg-type]
metadata=None,
input_schema=resolved_input_schema,
retry_policy=None,
cache_policy=None,
is_error_handler=True,
)
if input_schema is not None:
self.nodes[node] = StateNodeSpec[NodeInputT, ContextT](
coerce_to_runnable(action, name=node, trace=False), # type: ignore[arg-type]
@@ -798,7 +773,6 @@ class StateGraph(Generic[StateT, ContextT, InputT, OutputT]):
input_schema=input_schema,
retry_policy=retry_policy,
cache_policy=cache_policy,
error_handler_node=handler_node_name,
ends=ends,
defer=defer,
timeout=timeout,
@@ -810,7 +784,6 @@ class StateGraph(Generic[StateT, ContextT, InputT, OutputT]):
input_schema=inferred_input_schema,
retry_policy=retry_policy,
cache_policy=cache_policy,
error_handler_node=handler_node_name,
ends=ends,
defer=defer,
timeout=timeout,
@@ -822,7 +795,6 @@ class StateGraph(Generic[StateT, ContextT, InputT, OutputT]):
input_schema=self.state_schema,
retry_policy=retry_policy,
cache_policy=cache_policy,
error_handler_node=handler_node_name,
ends=ends,
defer=defer,
timeout=timeout,
@@ -1080,6 +1052,7 @@ class StateGraph(Generic[StateT, ContextT, InputT, OutputT]):
for node in interrupt:
if node not in self.nodes:
raise ValueError(f"Interrupt node `{node}` not found")
self.compiled = True
return self
@@ -1128,7 +1101,7 @@ class StateGraph(Generic[StateT, ContextT, InputT, OutputT]):
name: The name to use for the compiled graph.
transformers: Optional sequence of `StreamTransformer` classes or
configured factories. Classes and factories are instantiated
per run whenever `stream_events(version="v3")` / `astream_events(version="v3")` is called and are
per run whenever `stream_v2` / `astream_v2` is called and are
propagated to subgraph scopes. Custom factories should follow
the standard `StreamTransformer` constructor shape by
accepting `scope` as their first argument. Appended after the
@@ -1193,11 +1166,6 @@ class StateGraph(Generic[StateT, ContextT, InputT, OutputT]):
key for key, val in self.channels.items() if not is_managed_value(val)
]
)
node_error_handler_map = {
node_name: spec.error_handler_node
for node_name, spec in self.nodes.items()
if spec.error_handler_node is not None
}
compiled = CompiledStateGraph[StateT, ContextT, InputT, OutputT](
builder=self,
@@ -1220,7 +1188,6 @@ class StateGraph(Generic[StateT, ContextT, InputT, OutputT]):
debug=debug,
store=store,
cache=cache,
node_error_handler_map=node_error_handler_map,
name=name or "LangGraph",
stream_transformers=transformers,
)
@@ -1395,8 +1362,6 @@ class CompiledStateGraph(
metadata=node.metadata,
retry_policy=node.retry_policy,
cache_policy=node.cache_policy,
is_error_handler=node.is_error_handler,
error_handler_node=node.error_handler_node,
bound=node.runnable, # type: ignore[arg-type]
timeout=node.timeout,
)
+1 -189
View File
@@ -39,7 +39,6 @@ from langgraph._internal._constants import (
CONFIG_KEY_CHECKPOINT_MAP,
CONFIG_KEY_CHECKPOINT_NS,
CONFIG_KEY_CHECKPOINTER,
CONFIG_KEY_NODE_ERROR,
CONFIG_KEY_READ,
CONFIG_KEY_RESUME_MAP,
CONFIG_KEY_RUNTIME,
@@ -48,7 +47,6 @@ from langgraph._internal._constants import (
CONFIG_KEY_TASK_ID,
CONFIG_KEY_THREAD_ID,
ERROR,
ERROR_SOURCE_NODE,
INTERRUPT,
NO_WRITES,
NS_END,
@@ -68,7 +66,6 @@ from langgraph.channels.base import BaseChannel
from langgraph.channels.topic import Topic
from langgraph.channels.untracked_value import UntrackedValue
from langgraph.constants import TAG_HIDDEN
from langgraph.errors import NodeError
from langgraph.managed.base import ManagedValueMapping
from langgraph.pregel._call import get_runnable_for_task, identifier
from langgraph.pregel._io import read_channels
@@ -295,15 +292,7 @@ def apply_writes(
pending_writes_by_channel: dict[str, list[Any]] = defaultdict(list)
for task in tasks:
for chan, val in task.writes:
if chan in (
NO_WRITES,
PUSH,
RESUME,
INTERRUPT,
RETURN,
ERROR,
ERROR_SOURCE_NODE,
):
if chan in (NO_WRITES, PUSH, RESUME, INTERRUPT, RETURN, ERROR):
pass
elif chan in channels:
pending_writes_by_channel[chan].append(val)
@@ -761,42 +750,6 @@ def prepare_single_task(
return PregelTask(task_id, name, task_path[:3])
def _coerce_pending_error(value: Any) -> BaseException:
if isinstance(value, BaseException):
return value
return Exception(str(value))
def _read_errors_from_pending_writes(
pending_writes: list[PendingWrite],
) -> list[BaseException]:
errors: list[BaseException] = []
for _, channel, value in pending_writes:
if channel == ERROR:
errors.append(_coerce_pending_error(value))
return errors
def _read_error_for_task_id_from_pending_writes(
pending_writes: list[PendingWrite], task_id: str
) -> BaseException | None:
for pending_task_id, channel, value in reversed(pending_writes):
if pending_task_id == task_id and channel == ERROR:
return _coerce_pending_error(value)
return None
def _read_error_source_node_from_pending_writes(
pending_writes: list[PendingWrite], task_id: str
) -> str | None:
for pending_task_id, channel, value in reversed(pending_writes):
if pending_task_id == task_id and channel == ERROR_SOURCE_NODE:
if isinstance(value, str):
return value
return str(value)
return None
def prepare_push_task_functional(
task_path: tuple[str, tuple, int, str, Call],
# (PUSH, parent task path, idx of PUSH write, id of parent task, Call)
@@ -1107,147 +1060,6 @@ def prepare_push_task_send(
return PregelTask(task_id, packet.node, translated_task_path)
def prepare_node_error_handler_task(
failed_task: PregelExecutableTask,
*,
handler_node_name: str,
failed_error: BaseException,
checkpoint: Checkpoint,
pending_writes: list[PendingWrite],
processes: Mapping[str, PregelNode],
channels: Mapping[str, BaseChannel],
managed: ManagedValueMapping,
config: RunnableConfig,
step: int,
stop: int,
store: BaseStore | None = None,
checkpointer: BaseCheckpointSaver | None = None,
manager: None | ParentRunManager | AsyncParentRunManager = None,
cache_policy: CachePolicy | None = None,
retry_policy: Sequence[RetryPolicy] = (),
) -> PregelExecutableTask | None:
"""Prepare an immediate node-level error handler task for a failed task."""
if handler_node_name not in processes:
return None
proc = processes[handler_node_name]
proc_node = proc.node
if proc_node is None:
return None
checkpoint_id_bytes = binascii.unhexlify(checkpoint["id"].replace("-", ""))
task_id_func = _xxhash_str if checkpoint["v"] > 1 else _uuid5_str
configurable = config.get(CONF, {})
parent_ns = configurable.get(CONFIG_KEY_CHECKPOINT_NS, "")
checkpoint_ns = (
f"{parent_ns}{NS_SEP}{handler_node_name}" if parent_ns else handler_node_name
)
task_id = task_id_func(
checkpoint_id_bytes,
checkpoint_ns,
str(step),
handler_node_name,
PUSH,
"node_error_handler",
failed_task.id,
)
task_checkpoint_ns = f"{checkpoint_ns}:{task_id}"
translated_task_path = (*failed_task.path[:3], "node_error_handler", False)
metadata = {
"langgraph_step": step,
"langgraph_node": handler_node_name,
"langgraph_triggers": PUSH_TRIGGER,
"langgraph_path": translated_task_path,
"langgraph_checkpoint_ns": task_checkpoint_ns,
}
if proc.metadata:
metadata.update(proc.metadata)
writes: deque[tuple[str, Any]] = deque()
effective_retry_policy = proc.retry_policy or retry_policy
effective_cache_policy = proc.cache_policy or cache_policy
if effective_cache_policy:
args_key = effective_cache_policy.key_func(failed_task.input)
cache_key = CacheKey(
(
CACHE_NS_WRITES,
(identifier(proc) or "__dynamic__"),
handler_node_name,
),
xxh3_128_hexdigest(
args_key.encode() if isinstance(args_key, str) else args_key
),
effective_cache_policy.ttl,
)
else:
cache_key = None
scratchpad = _scratchpad(
config[CONF].get(CONFIG_KEY_SCRATCHPAD),
pending_writes,
task_id,
xxh3_128_hexdigest(task_checkpoint_ns.encode()),
config[CONF].get(CONFIG_KEY_RESUME_MAP),
step,
stop,
)
runtime = cast(Runtime, configurable.get(CONFIG_KEY_RUNTIME, DEFAULT_RUNTIME))
runtime = runtime.override(
store=store, previous=checkpoint["channel_values"].get(PREVIOUS, None)
)
additional_config: RunnableConfig = {
"metadata": metadata,
"tags": proc.tags,
}
return PregelExecutableTask(
handler_node_name,
failed_task.input,
proc_node,
writes,
patch_config(
merge_configs(config, additional_config),
run_name=handler_node_name,
callbacks=manager.get_child(f"graph:step:{step}") if manager else None,
configurable={
CONFIG_KEY_TASK_ID: task_id,
CONFIG_KEY_SEND: writes.extend,
CONFIG_KEY_READ: partial(
local_read,
scratchpad,
channels,
managed,
PregelTaskWrites(
translated_task_path,
handler_node_name,
writes,
PUSH_TRIGGER,
),
),
CONFIG_KEY_CHECKPOINTER: (
checkpointer or configurable.get(CONFIG_KEY_CHECKPOINTER)
),
CONFIG_KEY_CHECKPOINT_MAP: {
**configurable.get(CONFIG_KEY_CHECKPOINT_MAP, {}),
parent_ns: checkpoint["id"],
},
CONFIG_KEY_CHECKPOINT_ID: None,
CONFIG_KEY_CHECKPOINT_NS: task_checkpoint_ns,
CONFIG_KEY_SCRATCHPAD: scratchpad,
CONFIG_KEY_RUNTIME: runtime,
CONFIG_KEY_NODE_ERROR: NodeError(
node=failed_task.name, error=failed_error
),
},
),
PUSH_TRIGGER,
effective_retry_policy,
cache_key,
task_id,
translated_task_path,
writers=proc.flat_writers,
subgraphs=proc.subgraphs,
)
def checkpoint_null_version(
checkpoint: Checkpoint,
) -> V | None:
+3 -108
View File
@@ -45,13 +45,11 @@ from langgraph._internal._constants import (
CONFIG_KEY_REPLAY_STATE,
CONFIG_KEY_RESUME_MAP,
CONFIG_KEY_RESUMING,
CONFIG_KEY_RUNTIME,
CONFIG_KEY_SCRATCHPAD,
CONFIG_KEY_STREAM,
CONFIG_KEY_TASK_ID,
CONFIG_KEY_THREAD_ID,
ERROR,
ERROR_SOURCE_NODE,
INPUT,
INTERRUPT,
NS_END,
@@ -89,7 +87,6 @@ from langgraph.pregel._algo import (
checkpoint_null_version,
increment,
prepare_next_tasks,
prepare_node_error_handler_task,
prepare_single_task,
sanitize_untracked_values_in_send,
should_interrupt,
@@ -122,7 +119,6 @@ from langgraph.pregel.debug import (
map_debug_tasks,
)
from langgraph.pregel.protocol import StreamChunk, StreamProtocol
from langgraph.runtime import RunControl, Runtime
from langgraph.types import (
All,
CachePolicy,
@@ -210,12 +206,10 @@ class PregelLoop:
"input",
"pending",
"done",
"draining",
"interrupt_before",
"interrupt_after",
"out_of_steps",
]
control: RunControl | None
tasks: dict[str, PregelExecutableTask]
output: None | dict[str, Any] | Any = None
updated_channels: set[str] | None = None
@@ -323,8 +317,6 @@ class PregelLoop:
else ()
)
self.prev_checkpoint_config = None
runtime = self.config[CONF].get(CONFIG_KEY_RUNTIME)
self.control = runtime.control if isinstance(runtime, Runtime) else None
def _push_graph_lifecycle_event(
self,
@@ -332,16 +324,11 @@ class PregelLoop:
*,
interrupts: tuple[Interrupt, ...] = (),
) -> None:
# drain status never reaches lifecycle events: tick() returns False
# before pushing, and interrupts are raised through GraphInterrupt
if self.status == "draining":
raise RuntimeError("Draining status cannot emit lifecycle events")
status = self.status
if kind == "resume":
self._graph_lifecycle_events.append(
GraphResumeEvent(
run_id=None,
status=status,
status=self.status,
checkpoint_id=self.checkpoint["id"],
checkpoint_ns=self.checkpoint_ns,
)
@@ -350,7 +337,7 @@ class PregelLoop:
self._graph_lifecycle_events.append(
GraphInterruptEvent(
run_id=None,
status=status,
status=self.status,
checkpoint_id=self.checkpoint["id"],
checkpoint_ns=self.checkpoint_ns,
interrupts=interrupts,
@@ -524,16 +511,6 @@ class PregelLoop:
# return the new task, to be started if not run before
return pushed
def schedule_error_handler(
self, failed_task: PregelExecutableTask, error: BaseException
) -> PregelExecutableTask | None:
raise NotImplementedError
async def aschedule_error_handler(
self, failed_task: PregelExecutableTask, error: BaseException
) -> PregelExecutableTask | None:
raise NotImplementedError
def tick(self) -> bool:
"""Execute a single iteration of the Pregel loop.
@@ -592,10 +569,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 not self.is_replaying and self.checkpoint_pending_writes:
self._match_writes(self.tasks)
@@ -662,7 +635,7 @@ class PregelLoop:
def _match_writes(self, tasks: Mapping[str, PregelExecutableTask]) -> None:
for tid, k, v in self.checkpoint_pending_writes:
if k in (ERROR, ERROR_SOURCE_NODE, INTERRUPT, RESUME):
if k in (ERROR, INTERRUPT, RESUME):
continue
if task := tasks.get(tid):
task.writes.append((k, v))
@@ -1239,45 +1212,6 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager):
self.output_writes(task.id, task.writes, cached=True)
return pushed
def schedule_error_handler(
self, failed_task: PregelExecutableTask, error: BaseException
) -> PregelExecutableTask | None:
handler_node = self.nodes[failed_task.name].error_handler_node
if not handler_node:
return None
writes = list(failed_task.writes)
writes.append((ERROR_SOURCE_NODE, failed_task.name))
self.put_writes(
failed_task.id,
writes,
)
handler_task = prepare_node_error_handler_task(
failed_task,
handler_node_name=handler_node,
failed_error=error,
checkpoint=self.checkpoint,
pending_writes=self.checkpoint_pending_writes,
processes=self.nodes,
channels=self.channels,
managed=self.managed,
config=failed_task.config,
step=self.step,
stop=self.stop,
store=self.store,
checkpointer=self.checkpointer,
manager=self.manager,
retry_policy=self.retry_policy,
cache_policy=self.cache_policy,
)
if handler_task is None:
return None
self.tasks[handler_task.id] = handler_task
if not self.is_replaying:
self._match_writes({handler_task.id: handler_task})
for task in self.match_cached_writes():
self.output_writes(task.id, task.writes, cached=True)
return handler_task
def put_writes(self, task_id: str, writes: WritesT) -> None:
"""Put writes for a task, to be read by the next tick."""
super().put_writes(task_id, writes)
@@ -1485,45 +1419,6 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager):
self.output_writes(task.id, task.writes, cached=True)
return pushed
async def aschedule_error_handler(
self, failed_task: PregelExecutableTask, error: BaseException
) -> PregelExecutableTask | None:
handler_node = self.nodes[failed_task.name].error_handler_node
if not handler_node:
return None
writes = list(failed_task.writes)
writes.append((ERROR_SOURCE_NODE, failed_task.name))
self.put_writes(
failed_task.id,
writes,
)
handler_task = prepare_node_error_handler_task(
failed_task,
handler_node_name=handler_node,
failed_error=error,
checkpoint=self.checkpoint,
pending_writes=self.checkpoint_pending_writes,
processes=self.nodes,
channels=self.channels,
managed=self.managed,
config=failed_task.config,
step=self.step,
stop=self.stop,
store=self.store,
checkpointer=self.checkpointer,
manager=self.manager,
retry_policy=self.retry_policy,
cache_policy=self.cache_policy,
)
if handler_task is None:
return None
self.tasks[handler_task.id] = handler_task
if not self.is_replaying:
self._match_writes({handler_task.id: handler_task})
for task in await self.amatch_cached_writes():
self.output_writes(task.id, task.writes, cached=True)
return handler_task
def put_writes(self, task_id: str, writes: WritesT) -> None:
"""Put writes for a task, to be read by the next tick."""
super().put_writes(task_id, writes)
+1 -1
View File
@@ -349,7 +349,7 @@ class StreamMessagesHandlerV2(StreamMessagesHandler, _V2StreamingCallbackHandler
tags: list[str] | None = None,
**kwargs: Any,
) -> Any:
"""Forward a protocol event from `stream_events(version="v3")` as a messages stream part.
"""Forward a protocol event from `stream_v2` as a messages stream part.
Fires once per `MessagesData` event (`message-start`, per-block
`content-block-*`, `message-finish`). The transformer layer
-10
View File
@@ -138,12 +138,6 @@ class PregelNode:
metadata: Mapping[str, Any] | None
"""Metadata to attach to the node for tracing."""
is_error_handler: bool
"""Whether this node is registered as an error handler node."""
error_handler_node: str | None
"""Optional handler node name for failures from this node."""
subgraphs: Sequence[PregelProtocol]
"""Subgraphs used by the node."""
@@ -159,8 +153,6 @@ class PregelNode:
bound: Runnable[Any, Any] | None = None,
retry_policy: RetryPolicy | Sequence[RetryPolicy] | None = None,
cache_policy: CachePolicy | None = None,
is_error_handler: bool = False,
error_handler_node: str | None = None,
subgraphs: Sequence[PregelProtocol] | None = None,
timeout: float | timedelta | TimeoutPolicy | None = None,
) -> None:
@@ -177,8 +169,6 @@ class PregelNode:
self.timeout = coerce_timeout_policy(timeout)
self.tags = tags
self.metadata = metadata
self.is_error_handler = is_error_handler
self.error_handler_node = error_handler_node
if subgraphs is not None:
self.subgraphs = subgraphs
elif self.bound is not DEFAULT_BOUND:
+18 -183
View File
@@ -10,10 +10,8 @@ from collections.abc import (
AsyncIterator,
Awaitable,
Callable,
Collection,
Iterable,
Iterator,
Mapping,
Sequence,
)
from functools import partial
@@ -74,10 +72,6 @@ SKIP_RERAISE_SET: weakref.WeakSet[concurrent.futures.Future | asyncio.Future] =
class FuturesDict(Generic[F, E], dict[F, PregelExecutableTask | None]):
event: E
callback: weakref.ref[Callable[[PregelExecutableTask, BaseException | None], None]]
# Stop condition is injected by PregelRunner instead of hard-coded here.
# This lets the runner treat graph-error-handled exceptions as non-fatal
# so `on_done` does not trigger an early stop for those futures.
should_stop: Callable[[set[F]], bool]
counter: int
done: set[F]
lock: threading.Lock
@@ -88,7 +82,6 @@ class FuturesDict(Generic[F, E], dict[F, PregelExecutableTask | None]):
callback: weakref.ref[
Callable[[PregelExecutableTask, BaseException | None], None]
],
should_stop: Callable[[set[F]], bool],
future_type: type[F],
# used for generic typing, newer py supports FutureDict[...](...)
) -> None:
@@ -96,7 +89,6 @@ class FuturesDict(Generic[F, E], dict[F, PregelExecutableTask | None]):
self.lock = threading.Lock()
self.event = event
self.callback = callback
self.should_stop = should_stop
self.counter = 0
self.done: set[F] = set()
@@ -117,7 +109,6 @@ class FuturesDict(Generic[F, E], dict[F, PregelExecutableTask | None]):
task: PregelExecutableTask,
fut: F,
) -> None:
# Called automatically by future.add_done_callback registered in __setitem__.
try:
if cb := self.callback():
cb(task, _exception(fut))
@@ -125,9 +116,7 @@ class FuturesDict(Generic[F, E], dict[F, PregelExecutableTask | None]):
with self.lock:
self.done.add(fut)
self.counter -= 1
# Wake waiter when all tracked futures are done, or when runner-level
# stop condition is met (for example, a non-handled fatal exception).
if self.counter == 0 or self.should_stop(self.done):
if self.counter == 0 or _should_stop_others(self.done):
self.event.set()
@@ -143,34 +132,11 @@ class PregelRunner:
put_writes: weakref.ref[Callable[[str, Sequence[tuple[str, Any]]], None]],
use_astream: bool = False,
node_finished: Callable[[str], None] | None = None,
node_error_handler_map: Mapping[str, str] | None = None,
schedule_error_handler: Callable[
[PregelExecutableTask, BaseException], PregelExecutableTask | None
]
| None = None,
aschedule_error_handler: Callable[
[PregelExecutableTask, BaseException],
Awaitable[PregelExecutableTask | None],
]
| None = None,
) -> None:
self.submit = submit
self.put_writes = put_writes
self.use_astream = use_astream
self.node_finished = node_finished
self.node_error_handler_map = dict(node_error_handler_map or {})
self.error_handler_nodes = set(self.node_error_handler_map.values())
self.schedule_error_handler = schedule_error_handler
self.aschedule_error_handler = aschedule_error_handler
# Exception object ids that are already routed to graph-level error handler.
# These ids are consulted by stop/panic checks to avoid re-raising handled
# exceptions via the normal fatal path in the same run.
self._handled_exception_ids: set[int] = set()
def _should_route_to_error_handler(self, task: PregelExecutableTask) -> bool:
if task.name in self.error_handler_nodes:
return False
return task.name in self.node_error_handler_map
def tick(
self,
@@ -189,9 +155,6 @@ class PregelRunner:
futures = FuturesDict(
callback=weakref.WeakMethod(self.commit),
event=threading.Event(),
should_stop=partial(
_should_stop_others, handled_exception_ids=self._handled_exception_ids
),
future_type=concurrent.futures.Future,
)
# give control back to the caller
@@ -201,7 +164,6 @@ class PregelRunner:
return
elif len(tasks) == 1 and timeout is None and get_waiter is None:
t = tasks[0]
scheduled_error_handler = False
try:
run_with_retry(
t,
@@ -220,23 +182,12 @@ class PregelRunner:
self.commit(t, None)
except Exception as exc:
self.commit(t, exc)
if (
not isinstance(exc, GraphBubbleUp)
and self._should_route_to_error_handler(t)
and self.schedule_error_handler is not None
):
self._handled_exception_ids.add(id(exc))
if handler_task := self.schedule_error_handler(t, exc):
tasks = (handler_task,)
scheduled_error_handler = True
# Continue to the regular scheduling path for handler execution.
if reraise and futures:
if id(exc) not in self._handled_exception_ids:
# will be re-raised after futures are done
fut: concurrent.futures.Future = concurrent.futures.Future()
fut.set_exception(exc)
futures.done.add(fut)
elif reraise and id(exc) not in self._handled_exception_ids:
# will be re-raised after futures are done
fut: concurrent.futures.Future = concurrent.futures.Future()
fut.set_exception(exc)
futures.done.add(fut)
elif reraise:
if tb := exc.__traceback__:
while tb.tb_next is not None and any(
tb.tb_frame.f_code.co_filename.endswith(name)
@@ -245,12 +196,10 @@ class PregelRunner:
tb = tb.tb_next
exc.__traceback__ = tb
raise
if not futures and not scheduled_error_handler:
# maybe `t` scheduled another task
if not futures: # maybe `t` scheduled another task
return
else:
if not scheduled_error_handler:
tasks = () # don't reschedule this task
tasks = () # don't reschedule this task
# add waiter task if requested
if get_waiter is not None:
futures[get_waiter()] = None
@@ -277,7 +226,6 @@ class PregelRunner:
# each task is independent from all other concurrent tasks
# yield updates/debug output as each task finishes
end_time = timeout + time.monotonic() if timeout else None
handled_futures: set[concurrent.futures.Future[Any]] = set()
while len(futures) > (1 if get_waiter is not None else 0):
done, inflight = concurrent.futures.wait(
futures,
@@ -286,49 +234,17 @@ class PregelRunner:
)
if not done:
break # timed out
done_for_stop: set[concurrent.futures.Future[Any]] = set()
for fut in done:
task = futures.pop(fut)
if task is None:
# waiter task finished, schedule another
if inflight and get_waiter is not None:
futures[get_waiter()] = None
elif (
(task_exc := _exception(fut))
and self._should_route_to_error_handler(task)
and not isinstance(task_exc, GraphBubbleUp)
):
self._handled_exception_ids.add(id(task_exc))
SKIP_RERAISE_SET.add(fut)
handled_futures.add(fut)
if self.schedule_error_handler is not None:
if handler_task := self.schedule_error_handler(task, task_exc):
handler_fut = self.submit()( # type: ignore[misc]
run_with_retry,
handler_task,
retry_policy,
configurable={
CONFIG_KEY_CALL: partial(
_call,
weakref.ref(handler_task),
retry_policy=retry_policy,
futures=weakref.ref(futures),
schedule_task=schedule_task,
submit=self.submit,
),
},
__reraise_on_exit__=reraise,
)
futures[handler_fut] = handler_task
else:
done_for_stop.add(fut)
else:
# remove references to loop vars
del fut, task
# maybe stop other tasks
if _should_stop_others(
done_for_stop, handled_exception_ids=self._handled_exception_ids
):
if _should_stop_others(done):
break
# give control back to the caller
yield
@@ -343,8 +259,6 @@ class PregelRunner:
_panic_or_proceed(
futures.done.union(f for f, t in futures.items() if t is not None),
panic=reraise,
handled_exception_ids=self._handled_exception_ids,
handled_futures=handled_futures,
)
except Exception as exc:
if tb := exc.__traceback__:
@@ -378,9 +292,6 @@ class PregelRunner:
futures = FuturesDict(
callback=weakref.WeakMethod(self.commit),
event=asyncio.Event(),
should_stop=partial(
_should_stop_others, handled_exception_ids=self._handled_exception_ids
),
future_type=asyncio.Future,
)
# give control back to the caller
@@ -390,7 +301,6 @@ class PregelRunner:
return
elif len(tasks) == 1 and get_waiter is None and timeout is None:
t = tasks[0]
scheduled_error_handler = False
try:
await arun_with_retry(
t,
@@ -412,22 +322,12 @@ class PregelRunner:
self.commit(t, None)
except Exception as exc:
self.commit(t, exc)
if (
not isinstance(exc, GraphBubbleUp)
and self._should_route_to_error_handler(t)
and self.aschedule_error_handler is not None
):
self._handled_exception_ids.add(id(exc))
if handler_task := await self.aschedule_error_handler(t, exc):
tasks = (handler_task,)
scheduled_error_handler = True
if reraise and futures:
if id(exc) not in self._handled_exception_ids:
# will be re-raised after futures are done
fut: asyncio.Future = loop.create_future()
fut.set_exception(exc)
futures.done.add(fut)
elif reraise and id(exc) not in self._handled_exception_ids:
# will be re-raised after futures are done
fut: asyncio.Future = loop.create_future()
fut.set_exception(exc)
futures.done.add(fut)
elif reraise:
if tb := exc.__traceback__:
while tb.tb_next is not None and any(
tb.tb_frame.f_code.co_filename.endswith(name)
@@ -436,12 +336,10 @@ class PregelRunner:
tb = tb.tb_next
exc.__traceback__ = tb
raise
if not futures and not scheduled_error_handler:
# maybe `t` scheduled another task
if not futures: # maybe `t` scheduled another task
return
else:
if not scheduled_error_handler:
tasks = () # don't reschedule this task
tasks = () # don't reschedule this task
# add waiter task if requested
if get_waiter is not None:
futures[get_waiter()] = None
@@ -476,7 +374,6 @@ class PregelRunner:
# each task is independent from all other concurrent tasks
# yield updates/debug output as each task finishes
end_time = timeout + loop.time() if timeout else None
handled_futures: set[asyncio.Future[Any]] = set()
while len(futures) > (1 if get_waiter is not None else 0):
done, inflight = await asyncio.wait(
futures,
@@ -485,59 +382,17 @@ class PregelRunner:
)
if not done:
break # timed out
done_for_stop: set[asyncio.Future[Any]] = set()
for fut in done:
task = futures.pop(fut)
if task is None:
# waiter task finished, schedule another
if inflight and get_waiter is not None:
futures[get_waiter()] = None
elif (
(task_exc := _exception(fut))
and self._should_route_to_error_handler(task)
and not isinstance(task_exc, GraphBubbleUp)
):
self._handled_exception_ids.add(id(task_exc))
SKIP_RERAISE_SET.add(fut)
handled_futures.add(fut)
if self.aschedule_error_handler is not None:
if handler_task := await self.aschedule_error_handler(
task, task_exc
):
handler_fut = cast(
asyncio.Future,
self.submit()( # type: ignore[misc]
arun_with_retry,
handler_task,
retry_policy,
stream=self.use_astream,
configurable={
CONFIG_KEY_CALL: partial(
_acall,
weakref.ref(handler_task),
retry_policy=retry_policy,
stream=self.use_astream,
futures=weakref.ref(futures),
schedule_task=schedule_task,
submit=self.submit,
loop=loop,
),
},
__name__=handler_task.name,
__cancel_on_exit__=True,
__reraise_on_exit__=reraise,
),
)
futures[handler_fut] = handler_task
else:
done_for_stop.add(fut)
else:
# remove references to loop vars
del fut, task
# maybe stop other tasks
if _should_stop_others(
done_for_stop, handled_exception_ids=self._handled_exception_ids
):
if _should_stop_others(done):
break
# give control back to the caller
yield
@@ -557,8 +412,6 @@ class PregelRunner:
futures.done.union(f for f, t in futures.items() if t is not None),
timeout_exc_cls=asyncio.TimeoutError,
panic=reraise,
handled_exception_ids=self._handled_exception_ids,
handled_futures=handled_futures,
)
except Exception as exc:
if tb := exc.__traceback__:
@@ -594,11 +447,6 @@ class PregelRunner:
else:
# save error to checkpointer
task.writes.append((ERROR, exception))
if self._should_route_to_error_handler(task) and not isinstance(
exception, GraphBubbleUp
):
# Mark early in commit path; loop-side routing may happen later.
self._handled_exception_ids.add(id(exception))
self.put_writes()(task.id, task.writes) # type: ignore[misc]
else:
if self.node_finished and (
@@ -614,8 +462,6 @@ class PregelRunner:
def _should_stop_others(
done: set[F],
*,
handled_exception_ids: set[int] | None = None,
) -> bool:
"""Check if any task failed, if so, cancel all other tasks.
GraphInterrupts are not considered failures."""
@@ -623,11 +469,7 @@ def _should_stop_others(
if fut.cancelled():
continue
elif exc := fut.exception():
if (
id(exc) not in (handled_exception_ids or set())
and not isinstance(exc, GraphBubbleUp)
and fut not in SKIP_RERAISE_SET
):
if not isinstance(exc, GraphBubbleUp) and fut not in SKIP_RERAISE_SET:
return True
return False
@@ -651,9 +493,6 @@ def _panic_or_proceed(
*,
timeout_exc_cls: type[Exception] = TimeoutError,
panic: bool = True,
handled_exception_ids: set[int] | None = None,
handled_futures: Collection[concurrent.futures.Future[Any] | asyncio.Future[Any]]
| None = None,
) -> None:
"""Cancel remaining tasks if any failed, re-raise exception if panic is True."""
done: set[concurrent.futures.Future[Any] | asyncio.Future[Any]] = set()
@@ -670,10 +509,6 @@ def _panic_or_proceed(
# if any task failed
fut = done.pop()
if exc := _exception(fut):
if fut in (handled_futures or set()):
continue
if id(exc) in (handled_exception_ids or set()):
continue
# cancel all pending tasks
while inflight:
inflight.pop().cancel()
+60 -217
View File
@@ -30,7 +30,6 @@ from typing import (
)
from uuid import UUID, uuid5
from langchain_core._api import beta
from langchain_core.globals import get_debug
from langchain_core.runnables import (
RunnableSequence,
@@ -42,7 +41,6 @@ from langchain_core.runnables.config import (
get_callback_manager_for_config,
)
from langchain_core.runnables.graph import Graph
from langchain_core.runnables.schema import StreamEvent
from langgraph.cache.base import BaseCache
from langgraph.checkpoint.base import (
BaseCheckpointSaver,
@@ -113,7 +111,6 @@ from langgraph.config import get_config
from langgraph.constants import END
from langgraph.errors import (
ErrorCode,
GraphDrained,
GraphRecursionError,
InvalidUpdateError,
create_error_message,
@@ -159,7 +156,6 @@ from langgraph.pregel.protocol import PregelProtocol, StreamChunk, StreamProtoco
from langgraph.runtime import (
DEFAULT_RUNTIME,
BaseUser,
RunControl,
Runtime,
ServerInfo,
)
@@ -378,7 +374,7 @@ def _collect_stream_modes(mux: Any) -> list[StreamMode]:
"""Return the union of `required_stream_modes` across registered transformers.
Transformers declare the stream modes they need to function, and
`stream_events(version="v3")` asks the graph for exactly that union no hardcoded
`stream_v2` asks the graph for exactly that union no hardcoded
default set. If zero transformers declare a given mode, the graph
does not stream events for it.
"""
@@ -408,14 +404,14 @@ def _normalize_stream_transformer_factories(
for spec in specs or ():
if isinstance(spec, StreamTransformer):
raise TypeError(
"stream_events(version='v3') transformers must be scope-aware callables, "
"stream_v2 transformers must be scope-aware callables, "
f"got pre-built instance {type(spec).__name__}. Pass the "
"transformer class or a factory like "
"`lambda scope: MyTransformer(scope, ...)`."
)
if not callable(spec):
raise TypeError(
"stream_events(version='v3') transformers must be scope-aware callables, "
"stream_v2 transformers must be scope-aware callables, "
f"got {type(spec).__name__}."
)
@@ -732,7 +728,6 @@ class Pregel(
name: str = "LangGraph"
trigger_to_nodes: Mapping[str, Sequence[str]]
node_error_handler_map: Mapping[str, str]
def __init__(
self,
@@ -757,7 +752,6 @@ class Pregel(
context_schema: type[ContextT] | None = None,
config: RunnableConfig | None = None,
trigger_to_nodes: Mapping[str, Sequence[str]] | None = None,
node_error_handler_map: Mapping[str, str] | None = None,
name: str = "LangGraph",
stream_transformers: Sequence[Callable[[tuple[str, ...]], Any]] | None = None,
**deprecated_kwargs: Unpack[DeprecatedKwargs],
@@ -805,7 +799,6 @@ class Pregel(
self.context_schema = context_schema
self.config = config
self.trigger_to_nodes = trigger_to_nodes or {}
self.node_error_handler_map = node_error_handler_map or {}
self.name = name
self.stream_transformers: tuple[Callable[[tuple[str, ...]], Any], ...] = tuple(
stream_transformers or ()
@@ -2577,7 +2570,6 @@ class Pregel(
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
durability: Durability | None = None,
control: RunControl | None = None,
subgraphs: bool = False,
debug: bool | None = None,
version: Literal["v2"],
@@ -2597,7 +2589,6 @@ class Pregel(
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
durability: Durability | None = None,
control: RunControl | None = None,
subgraphs: bool = False,
debug: bool | None = None,
version: Literal["v1"] = ...,
@@ -2616,7 +2607,6 @@ class Pregel(
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
durability: Durability | None = None,
control: RunControl | None = None,
subgraphs: bool = False,
debug: bool | None = None,
version: Literal["v1", "v2"] = "v1",
@@ -2661,7 +2651,6 @@ class Pregel(
- `"sync"`: Changes are persisted synchronously before the next step starts.
- `"async"`: Changes are persisted asynchronously while the next step executes.
- `"exit"`: Changes are persisted only when the graph exits.
control: Optional run control used to request cooperative drain.
subgraphs: Whether to stream events from inside subgraphs, defaults to `False`.
If `True`, the events will be emitted as tuples `(namespace, data)`,
@@ -2826,7 +2815,6 @@ class Pregel(
previous=None,
execution_info=None,
server_info=server_info,
control=control or parent_runtime.control or RunControl(),
)
runtime = parent_runtime.merge(runtime)
config[CONF][CONFIG_KEY_RUNTIME] = runtime
@@ -2876,8 +2864,6 @@ class Pregel(
),
put_writes=weakref.WeakMethod(loop.put_writes),
node_finished=config[CONF].get(CONFIG_KEY_NODE_FINISHED),
node_error_handler_map=self.node_error_handler_map,
schedule_error_handler=loop.schedule_error_handler,
)
# enable subgraph streaming
if subgraphs:
@@ -2959,10 +2945,6 @@ class Pregel(
error_code=ErrorCode.GRAPH_RECURSION_LIMIT,
)
raise GraphRecursionError(msg)
elif loop.status == "draining":
if loop.control is None:
raise RuntimeError("Draining status requires run control")
raise GraphDrained(loop.control.drain_reason or "shutdown")
# set final channel values as run output
run_manager.on_chain_end(loop.output)
except BaseException as e:
@@ -2983,7 +2965,6 @@ class Pregel(
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
durability: Durability | None = None,
control: RunControl | None = None,
subgraphs: bool = False,
debug: bool | None = None,
version: Literal["v2"],
@@ -3003,7 +2984,6 @@ class Pregel(
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
durability: Durability | None = None,
control: RunControl | None = None,
subgraphs: bool = False,
debug: bool | None = None,
version: Literal["v1"] = ...,
@@ -3022,7 +3002,6 @@ class Pregel(
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
durability: Durability | None = None,
control: RunControl | None = None,
subgraphs: bool = False,
debug: bool | None = None,
version: Literal["v1", "v2"] = "v1",
@@ -3067,7 +3046,6 @@ class Pregel(
- `"sync"`: Changes are persisted synchronously before the next step starts.
- `"async"`: Changes are persisted asynchronously while the next step executes.
- `"exit"`: Changes are persisted only when the graph exits.
control: Optional run control used to request cooperative drain.
subgraphs: Whether to stream events from inside subgraphs, defaults to `False`.
If `True`, the events will be emitted as tuples `(namespace, data)`,
@@ -3267,7 +3245,6 @@ class Pregel(
previous=None,
execution_info=None,
server_info=server_info,
control=control or parent_runtime.control or RunControl(),
)
runtime = parent_runtime.merge(runtime)
config[CONF][CONFIG_KEY_RUNTIME] = runtime
@@ -3329,8 +3306,6 @@ class Pregel(
put_writes=weakref.WeakMethod(loop.put_writes),
use_astream=do_stream,
node_finished=config[CONF].get(CONFIG_KEY_NODE_FINISHED),
node_error_handler_map=self.node_error_handler_map,
aschedule_error_handler=loop.aschedule_error_handler,
)
# enable subgraph streaming
if subgraphs:
@@ -3438,10 +3413,6 @@ class Pregel(
error_code=ErrorCode.GRAPH_RECURSION_LIMIT,
)
raise GraphRecursionError(msg)
elif loop.status == "draining":
if loop.control is None:
raise RuntimeError("Draining status requires run control")
raise GraphDrained(loop.control.drain_reason or "shutdown")
# set final channel values as run output
await run_manager.on_chain_end(loop.output)
except BaseException as e:
@@ -3449,22 +3420,49 @@ class Pregel(
await asyncio.shield(run_manager.on_chain_error(e))
raise
@beta(message="The v3 streaming protocol on Pregel is experimental.")
def _pregel_stream_v3(
def stream_v2(
self,
input: InputT | Command | None,
config: RunnableConfig | None = None,
*,
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
control: RunControl | None = None,
transformers: Sequence[Callable[[tuple[str, ...]], Any]] | None = None,
) -> Any:
"""Internal v3 sync streaming implementation. Public entry: stream_events(version='v3').
"""Start a sync v2 streaming run driven by transformer projections.
!!! warning
Builds a `StreamMux` from the built-in transformers, this
graph's compile-time `stream_transformers`, and any additional
`transformers=` supplied at the call site. Returns a
`GraphRunStream` that the caller drives by iterating any
projection no background thread.
The v3 streaming protocol is experimental and may change.
`run.output`, `run.interrupted` and `run.interrupts` work
regardless of which transformers are registered.
Note:
Nesting v1 `stream(stream_mode="messages")` inside a node
of a `stream_v2` run is not fully supported. The outer v2
messages handler reroutes `BaseChatModel.invoke` through
the v2 event protocol, so the inner v1 handler does not see
`on_llm_new_token` chunks. The inner stream still yields a
finalized message via `on_llm_end`. Use `stream_v2` for
the inner graph as well, or call
`chat_model.stream(...)` explicitly, to get token-level
streaming.
Args:
input: Graph input.
config: Optional runnable config forwarded to the graph.
interrupt_before: Nodes to interrupt before, if any.
interrupt_after: Nodes to interrupt after, if any.
transformers: Extra transformer classes or configured factories
appended after compile-time `stream_transformers`. Factories
are called as `factory(scope)` so they can propagate to
subgraph scopes.
Returns:
A `GraphRunStream` the caller iterates to drive the run.
"""
parent_ns = _resolve_parent_ns(self.config, config)
compiled_factories = _normalize_stream_transformer_factories(
@@ -3492,27 +3490,43 @@ class Pregel(
version="v2",
interrupt_before=interrupt_before,
interrupt_after=interrupt_after,
control=control,
)
)
return GraphRunStream(graph_iter, mux)
@beta(message="The v3 streaming protocol on Pregel is experimental.")
async def _apregel_stream_v3(
async def astream_v2(
self,
input: InputT | Command | None,
config: RunnableConfig | None = None,
*,
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
control: RunControl | None = None,
transformers: Sequence[Callable[[tuple[str, ...]], Any]] | None = None,
) -> Any:
"""Internal v3 async streaming implementation. Public entry: astream_events(version='v3').
"""Async counterpart to `stream_v2`.
!!! warning
Returns an `AsyncGraphRunStream` whose projections can be awaited
concurrently; each subscribed cursor drives the pump when its
buffer is empty.
The v3 streaming protocol is experimental and may change.
Note:
Same nesting limitation as `stream_v2`: nesting v1
`astream(stream_mode="messages")` inside a node of an
`astream_v2` run drops `on_llm_new_token` chunks because
the outer v2 handler reroutes `BaseChatModel.invoke`
through the v2 event protocol. Use `astream_v2` for the
inner graph as well, or call `chat_model.astream(...)`
explicitly, to get token-level streaming.
Args:
input: Graph input.
config: Optional runnable config forwarded to the graph.
interrupt_before: Nodes to interrupt before, if any.
interrupt_after: Nodes to interrupt after, if any.
transformers: Extra transformer classes or configured factories
appended after compile-time `stream_transformers`. Factories
are called as `factory(scope)` so they can propagate to
subgraph scopes.
"""
parent_ns = _resolve_parent_ns(self.config, config)
compiled_factories = _normalize_stream_transformer_factories(
@@ -3539,166 +3553,9 @@ class Pregel(
version="v2",
interrupt_before=interrupt_before,
interrupt_after=interrupt_after,
control=control,
).__aiter__()
return AsyncGraphRunStream(graph_aiter, mux)
@overload
def stream_events(
self,
input: InputT | Command | None,
config: RunnableConfig | None = None,
*,
version: Literal["v1", "v2"] = "v2",
**kwargs: Any,
) -> Iterator[StreamEvent]: ...
@overload
def stream_events(
self,
input: InputT | Command | None,
config: RunnableConfig | None = None,
*,
version: Literal["v3"],
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
control: RunControl | None = None,
transformers: Sequence[Callable[[tuple[str, ...]], Any]] | None = None,
) -> Any: ...
def stream_events(
self,
input: InputT | Command | None,
config: RunnableConfig | None = None,
*,
version: Literal["v1", "v2", "v3"] = "v2",
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
control: RunControl | None = None,
transformers: Sequence[Callable[[tuple[str, ...]], Any]] | None = None,
**kwargs: Any,
) -> Any:
"""Stream events from this graph.
For `version="v1"` / `"v2"`, yields `StreamEvent` dicts (see
`Runnable.stream_events`). For `version="v3"`, returns a
`GraphRunStream` whose typed projections the caller drives by
iterating no background thread.
!!! warning
The `version="v3"` API is experimental and may change.
Builds a `StreamMux` from the built-in transformers, this
graph's compile-time `stream_transformers`, and any additional
`transformers=` supplied at the call site. `run.output`,
`run.interrupted`, and `run.interrupts` work regardless of
which transformers are registered.
Note:
Nesting v1 `stream(stream_mode="messages")` inside a node
of a `stream_events(version="v3")` run is not fully
supported. The outer v3 messages handler reroutes
`BaseChatModel.invoke` through the v2 event protocol, so
the inner v1 handler does not see `on_llm_new_token`
chunks. The inner stream still yields a finalized message
via `on_llm_end`. Use `stream_events(version="v3")` for the
inner graph as well, or call `chat_model.stream(...)`
explicitly, to get token-level streaming.
Args:
input: Graph input.
config: Optional runnable config.
version: Streaming-event schema version. `"v3"` selects the
content-block-centric streaming protocol.
interrupt_before: Nodes to interrupt before, if any. Only
used for `version="v3"`.
interrupt_after: Nodes to interrupt after, if any. Only
used for `version="v3"`.
control: Optional run control used to request cooperative
drain. Only used for `version="v3"`.
transformers: Extra transformer classes or configured
factories appended after compile-time
`stream_transformers`. Factories are called as
`factory(scope)` so they can propagate to subgraph
scopes. Only used for `version="v3"`.
**kwargs: Forwarded to the v1/v2 path.
Returns:
For `version="v3"`, a `GraphRunStream` the caller iterates
to drive the run. Otherwise an `Iterator[StreamEvent]`.
"""
if version == "v3":
return self._pregel_stream_v3(
input,
config,
interrupt_before=interrupt_before,
interrupt_after=interrupt_after,
control=control,
transformers=transformers,
)
return super().stream_events(input, config, version=version, **kwargs)
@overload
def astream_events(
self,
input: InputT | Command | None,
config: RunnableConfig | None = None,
*,
version: Literal["v1", "v2"] = "v2",
**kwargs: Any,
) -> AsyncIterator[StreamEvent]: ...
@overload
def astream_events(
self,
input: InputT | Command | None,
config: RunnableConfig | None = None,
*,
version: Literal["v3"],
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
control: RunControl | None = None,
transformers: Sequence[Callable[[tuple[str, ...]], Any]] | None = None,
) -> Awaitable[Any]: ...
def astream_events(
self,
input: InputT | Command | None,
config: RunnableConfig | None = None,
*,
version: Literal["v1", "v2", "v3"] = "v2",
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
control: RunControl | None = None,
transformers: Sequence[Callable[[tuple[str, ...]], Any]] | None = None,
**kwargs: Any,
) -> AsyncIterator[StreamEvent] | Awaitable[Any]:
"""Async variant of `stream_events`.
For `version="v3"`, returns an `AsyncGraphRunStream` whose
projections can be awaited concurrently; each subscribed cursor
drives the pump when its buffer is empty. The same nesting
limitation as the sync path applies see `stream_events` for
details.
!!! warning
The `version="v3"` API is experimental and may change.
See `stream_events` for full argument and return documentation.
"""
if version == "v3":
return self._apregel_stream_v3(
input,
config,
interrupt_before=interrupt_before,
interrupt_after=interrupt_after,
control=control,
transformers=transformers,
)
return super().astream_events(input, config, version=version, **kwargs)
@overload
def invoke(
self,
@@ -3712,7 +3569,6 @@ class Pregel(
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
durability: Durability | None = None,
control: RunControl | None = None,
version: Literal["v2"],
**kwargs: Any,
) -> GraphOutput[OutputT]: ...
@@ -3730,7 +3586,6 @@ class Pregel(
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
durability: Durability | None = None,
control: RunControl | None = None,
version: Literal["v2"],
**kwargs: Any,
) -> list[StreamPart[StateT, OutputT]]: ...
@@ -3748,7 +3603,6 @@ class Pregel(
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
durability: Durability | None = None,
control: RunControl | None = None,
version: Literal["v1"] = ...,
**kwargs: Any,
) -> dict[str, Any] | Any: ...
@@ -3765,7 +3619,6 @@ class Pregel(
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
durability: Durability | None = None,
control: RunControl | None = None,
version: Literal["v1", "v2"] = "v1",
**kwargs: Any,
) -> dict[str, Any] | Any:
@@ -3790,7 +3643,6 @@ class Pregel(
- `"sync"`: Changes are persisted synchronously before the next step starts.
- `"async"`: Changes are persisted asynchronously while the next step executes.
- `"exit"`: Changes are persisted only when the graph exits.
control: Optional run control used to request cooperative drain.
version: The streaming format version. `"v1"` (default) returns the
traditional format, `"v2"` returns `StreamPart` typed dicts when
`stream_mode` is not `"values"`.
@@ -3818,7 +3670,6 @@ class Pregel(
interrupt_before=interrupt_before,
interrupt_after=interrupt_after,
durability=durability,
control=control,
version=version,
**kwargs,
):
@@ -3842,7 +3693,6 @@ class Pregel(
interrupt_before=interrupt_before,
interrupt_after=interrupt_after,
durability=durability,
control=control,
**kwargs,
):
if stream_mode == "values":
@@ -3889,7 +3739,6 @@ class Pregel(
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
durability: Durability | None = None,
control: RunControl | None = None,
version: Literal["v2"],
**kwargs: Any,
) -> GraphOutput[OutputT]: ...
@@ -3907,7 +3756,6 @@ class Pregel(
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
durability: Durability | None = None,
control: RunControl | None = None,
version: Literal["v2"],
**kwargs: Any,
) -> list[StreamPart[StateT, OutputT]]: ...
@@ -3925,7 +3773,6 @@ class Pregel(
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
durability: Durability | None = None,
control: RunControl | None = None,
version: Literal["v1"] = ...,
**kwargs: Any,
) -> dict[str, Any] | Any: ...
@@ -3942,7 +3789,6 @@ class Pregel(
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
durability: Durability | None = None,
control: RunControl | None = None,
version: Literal["v1", "v2"] = "v1",
**kwargs: Any,
) -> dict[str, Any] | Any:
@@ -3967,7 +3813,6 @@ class Pregel(
- `"sync"`: Changes are persisted synchronously before the next step starts.
- `"async"`: Changes are persisted asynchronously while the next step executes.
- `"exit"`: Changes are persisted only when the graph exits.
control: Optional run control used to request cooperative drain.
version: The streaming format version. `"v1"` (default) returns the
traditional format, `"v2"` returns `StreamPart` typed dicts when
`stream_mode` is not `"values"`.
@@ -3995,7 +3840,6 @@ class Pregel(
interrupt_before=interrupt_before,
interrupt_after=interrupt_after,
durability=durability,
control=control,
version=version,
**kwargs,
):
@@ -4019,7 +3863,6 @@ class Pregel(
interrupt_before=interrupt_before,
interrupt_after=interrupt_after,
durability=durability,
control=control,
**kwargs,
):
if stream_mode == "values":
@@ -4193,7 +4036,7 @@ def _resolve_parent_ns(
) -> tuple[str, ...]:
"""Return the checkpoint namespace the caller is running under.
`stream_events(version="v3")` uses this to scope its native projections
`stream_v2` uses this to scope its native projections
(`ValuesTransformer`, `MessagesTransformer`) to events emitted at
the run's own level. A root call resolves to `()`; a call made
from inside a node carries the outer graph's task namespace so the
+2 -49
View File
@@ -16,7 +16,6 @@ from langgraph.typing import ContextT
__all__ = (
"BaseUser",
"ExecutionInfo",
"RunControl",
"Runtime",
"ServerInfo",
"get_runtime",
@@ -76,34 +75,6 @@ class ServerInfo:
"""
class RunControl:
"""Run-scoped control surface for cooperative draining.
Intended for a single graph run. Create a fresh `RunControl` per run;
reusing a control after `request_drain()` leaves it drained.
Safe to call from any thread: the drain request is represented by a
single attribute write, so no lock is needed for this signal.
If more mutable state is added here, add synchronization.
"""
__slots__ = ("_drain_reason",)
def __init__(self) -> None:
self._drain_reason: str | None = None
def request_drain(self, reason: str = "shutdown") -> None:
self._drain_reason = reason
@property
def drain_requested(self) -> bool:
return self._drain_reason is not None
@property
def drain_reason(self) -> str | None:
return self._drain_reason
def _no_op_stream_writer(_: Any) -> None: ...
@@ -118,7 +89,6 @@ class _RuntimeOverrides(TypedDict, Generic[ContextT], total=False):
previous: Any
execution_info: ExecutionInfo
server_info: ServerInfo | None
control: RunControl | None
@dataclass(**_DC_KWARGS)
@@ -197,7 +167,7 @@ class Runtime(Generic[ContextT]):
context: ContextT = field(default=None) # type: ignore[assignment]
"""Static context for the graph run, like `user_id`, `db_conn`, etc.
Can also be thought of as 'run dependencies'."""
store: BaseStore | None = field(default=None)
@@ -218,7 +188,7 @@ class Runtime(Generic[ContextT]):
previous: Any = field(default=None)
"""The previous return value for the given thread.
Only available with the functional API when a checkpointer is provided.
"""
@@ -230,13 +200,6 @@ class Runtime(Generic[ContextT]):
server_info: ServerInfo | None = field(default=None)
"""Metadata injected by LangGraph Server. None when running open-source LangGraph without LangSmith deployments."""
control: RunControl | None = field(default=None)
"""Run-scoped control plane for cooperative draining.
Populated automatically during graph runs. None outside an active
graph runtime.
"""
def merge(self, other: Runtime[ContextT]) -> Runtime[ContextT]:
"""Merge two runtimes together.
@@ -254,7 +217,6 @@ class Runtime(Generic[ContextT]):
previous=self.previous if other.previous is None else other.previous,
execution_info=other.execution_info or self.execution_info,
server_info=other.server_info or self.server_info,
control=other.control or self.control,
)
def override(
@@ -273,14 +235,6 @@ class Runtime(Generic[ContextT]):
execution_info=self.execution_info.patch(**overrides),
)
@property
def drain_requested(self) -> bool:
return self.control.drain_requested if self.control is not None else False
@property
def drain_reason(self) -> str | None:
return self.control.drain_reason if self.control is not None else None
DEFAULT_RUNTIME = Runtime(
context=None,
@@ -289,7 +243,6 @@ DEFAULT_RUNTIME = Runtime(
heartbeat=_no_op_heartbeat,
previous=None,
execution_info=None,
control=None,
)
+2 -2
View File
@@ -1,7 +1,7 @@
"""Streaming infrastructure for LangGraph.
Compile a graph with `transformers=[...]` and call `graph.stream_events(version="v3")` /
`graph.astream_events(version="v3")` to drive a transformer pipeline that projects the
Compile a graph with `transformers=[...]` and call `graph.stream_v2()` /
`graph.astream_v2()` to drive a transformer pipeline that projects the
graph's raw events into ergonomic per-channel streams.
"""
-7
View File
@@ -91,11 +91,9 @@ class StreamMux:
self._assign_seq = _assign_seq
self._events: StreamChannel[ProtocolEvent] = StreamChannel()
self._events._bind(is_async=is_async)
self._events._bind_mux(self)
self._transformers: list[StreamTransformer] = []
self._channels: list[StreamChannel[Any]] = []
self._seq = 0
self._push_seq = 0
self.extensions: dict[str, Any] = {}
self.native_keys: set[str] = set()
@@ -126,10 +124,6 @@ class StreamMux:
"""Return the transformer that contributed `key` to the projection."""
return self._transformer_by_key.get(key)
def _next_push_seq(self) -> int:
self._push_seq += 1
return self._push_seq
# ------------------------------------------------------------------
# Pump wiring + mini-mux nesting
# ------------------------------------------------------------------
@@ -455,7 +449,6 @@ class StreamMux:
for value in projection.values():
if isinstance(value, StreamChannel):
value._bind(is_async=self.is_async)
value._bind_mux(self)
self._channels.append(value)
if value.name is not None:
method = value.name if native else f"custom:{value.name}"
+1 -1
View File
@@ -88,7 +88,7 @@ class StreamTransformer(ABC):
required_stream_modes: Stream modes the graph must emit for
this transformer to have anything to process. Computed as
the union across all registered transformers to determine
which modes a `stream_events(version="v3")` run requests from the graph.
which modes a `stream_v2` run requests from the graph.
Empty tuple means the transformer consumes only synthetic
events (or is purely passive).
"""
+27 -88
View File
@@ -5,8 +5,6 @@ from collections.abc import AsyncIterator, Awaitable, Callable, Iterator, Mappin
from types import MappingProxyType, TracebackType
from typing import TYPE_CHECKING, Any
from langchain_core._api import beta
from langgraph.stream._convert import convert_to_protocol_event
from langgraph.stream._mux import StreamMux
from langgraph.stream._types import ProtocolEvent
@@ -27,7 +25,6 @@ async def _adrive_until_done(pump: Callable[[], Awaitable[bool]]) -> None:
pass
@beta(message="The v3 streaming protocol on Pregel is experimental.")
class GraphRunStream:
"""Sync run stream with caller-driven pumping.
@@ -41,11 +38,6 @@ class GraphRunStream:
All transformer projections live in `extensions`. Native transformer
projections (those with `_native = True`) are also set as direct
attributes on this instance (e.g. `run.values`, `run.messages`).
!!! warning
Returned by `Pregel.stream_events(version="v3")`, which is
experimental and may change.
"""
def __init__(
@@ -193,25 +185,28 @@ class GraphRunStream:
return iter(self._mux._events)
def interleave(self, *names: str) -> Iterator[tuple[str, Any]]:
"""Iterate multiple projections in arrival order, yielding ``(name, item)``.
"""Iterate multiple projections round-robin, yielding ``(name, item)``.
Items are ordered by a monotonic push stamp assigned when each
transformer pushes into its `StreamChannel`. This gives strict
arrival ordering across projections, unlike round-robin.
Each turn advances one projection's cursor; when a cursor's buffer
is empty, pulling from it drives the pump once, which fans out to
every subscribed projection log. Projections whose items aren't
consumed on this turn sit in their own buffers only until the next
turn reaches them, bounding memory by the skew between projection
rates rather than letting any single log grow to the full run
length.
Projections are exhausted independently; a projection that finishes
early drops out of the rotation while others continue. The overall
iterator ends once all named projections are done.
Args:
*names: Projection keys to interleave. Must match keys in
``extensions``.
Yields:
``(name, item)`` tuples in arrival order across the named
``(name, item)`` tuples in round-robin order across the named
projections.
Each named channel is locked for the duration of iteration and
released when the generator completes, is closed, or raises.
Channels cannot be subscribed concurrently use `.tee(n)` if
you need fan-out.
Raises:
KeyError: If a name doesn't match a registered projection.
@@ -224,73 +219,22 @@ class GraphRunStream:
print("val:", item)
```
"""
from langgraph.stream.stream_channel import StreamChannel
channels: dict[str, StreamChannel[Any]] = {}
try:
for name in names:
ch = self.extensions[name]
if not isinstance(ch, StreamChannel):
raise TypeError(
f"interleave() requires StreamChannel projections, "
f"got {type(ch).__name__} for {name!r}"
)
if ch._is_async is None:
raise TypeError(
f"StreamChannel {name!r} has not been bound yet. "
"Register the transformer with a StreamMux first."
)
if ch._is_async:
raise TypeError(
f"StreamChannel {name!r} is bound to async mode — "
"sync interleave() cannot consume async channels."
)
if ch._subscribed:
raise RuntimeError(
f"StreamChannel {name!r} already has a subscriber; "
"use .tee(n) for fan-out."
)
ch._subscribed = True
channels[name] = ch
done: set[str] = set()
while len(done) < len(channels):
best: tuple[int, str] | None = None
for name, ch in channels.items():
if name in done:
continue
if ch._closed and not ch._items:
if ch._error is not None:
raise ch._error
done.add(name)
continue
if ch._items:
stamp = ch._items[0][0]
if best is None or stamp < best[0]:
best = (stamp, name)
if best is not None:
_stamp, item = channels[best[1]]._items.popleft()
yield (best[1], item)
else:
pump = self._mux._pump_fn
if pump is None or not pump():
before = len(done)
for name, ch in channels.items():
if name not in done and not ch._items:
if ch._closed:
if ch._error is not None:
raise ch._error
done.add(name)
if len(done) == before:
break
finally:
for ch in channels.values():
ch._subscribed = False
cursors: dict[str, Iterator[Any]] = {
name: iter(self.extensions[name]) for name in names
}
done: set[str] = set()
while len(done) < len(cursors):
for name, cursor in cursors.items():
if name in done:
continue
try:
item = next(cursor)
except StopIteration:
done.add(name)
continue
yield (name, item)
@beta(message="The v3 streaming protocol on Pregel is experimental.")
class AsyncGraphRunStream:
"""Async run stream with caller-driven pumping.
@@ -312,11 +256,6 @@ class AsyncGraphRunStream:
async for msg in run.messages:
...
```
!!! warning
Awaited from `Pregel.astream_events(version="v3")`, which is
experimental and may change.
"""
def __init__(
@@ -3,10 +3,7 @@ from __future__ import annotations
import asyncio
from collections import deque
from collections.abc import AsyncIterator, Awaitable, Callable, Iterator
from typing import TYPE_CHECKING, Generic, TypeVar
if TYPE_CHECKING:
from langgraph.stream._mux import StreamMux
from typing import Generic, TypeVar
T = TypeVar("T")
@@ -67,7 +64,7 @@ class StreamChannel(Generic[T]):
if maxlen is not None and maxlen <= 0:
raise ValueError("StreamChannel maxlen must be a positive int or None")
self.name = name
self._items: deque[tuple[int, T]] = deque()
self._items: deque[T] = deque()
self._maxlen: int | None = maxlen
self._closed = False
self._error: BaseException | None = None
@@ -80,15 +77,11 @@ class StreamChannel(Generic[T]):
self._arequest_more: Callable[[], Awaitable[bool]] | None = None
self._wire_fn: Callable[[T], None] | None = None
self._mux: StreamMux | None = None
# ------------------------------------------------------------------
# Binding
# ------------------------------------------------------------------
def _bind_mux(self, mux: StreamMux) -> None:
self._mux = mux
def _bind(self, *, is_async: bool) -> None:
"""Bind this channel to sync or async mode.
@@ -124,18 +117,13 @@ class StreamChannel(Generic[T]):
registered, but auto-forwarding always fires so wired events
reach the main event log regardless of subscription state.
Items are stored as `(stamp, item)` tuples where stamp is a
monotonic counter from the owning mux. Stamps are stripped by
the default cursors; raw stamped tuples are visible on `_items`.
Raises:
RuntimeError: If the channel is closed (and subscribed).
"""
if self._subscribed:
if self._closed:
raise RuntimeError("Cannot push to a closed StreamChannel")
stamp = self._mux._next_push_seq() if self._mux is not None else 0
self._items.append((stamp, item))
self._items.append(item)
if self._wire_fn is not None:
self._wire_fn(item)
@@ -182,8 +170,7 @@ class StreamChannel(Generic[T]):
def _sync_cursor(self) -> Iterator[T]:
while True:
if self._items:
_stamp, item = self._items.popleft()
yield item
yield self._items.popleft()
elif self._closed:
if self._error is not None:
raise self._error
@@ -225,8 +212,7 @@ class StreamChannel(Generic[T]):
async def _async_cursor(self) -> AsyncIterator[T]:
while True:
if self._items:
_stamp, item = self._items.popleft()
yield item
yield self._items.popleft()
elif self._closed:
if self._error is not None:
raise self._error
+13 -13
View File
@@ -12,7 +12,7 @@ from langchain_core.messages import AIMessageChunk, BaseMessage
from langchain_protocol.protocol import MessagesData
from typing_extensions import NotRequired, TypedDict
from langgraph.errors import GraphDrained, GraphInterrupt
from langgraph.errors import GraphInterrupt
from langgraph.stream._types import ProtocolEvent, StreamTransformer
from langgraph.stream.run_stream import AsyncSubgraphRunStream, SubgraphRunStream
from langgraph.stream.stream_channel import StreamChannel
@@ -38,8 +38,8 @@ class ValuesTransformer(StreamTransformer):
Only values events at the run's own level are captured; snapshots
from deeper subgraphs are left in the main event log but excluded
from the projection. "Own level" is defined by `scope`, which
`stream_events(version="v3")` / `astream_events(version="v3")` populate from the caller's
checkpoint namespace so that a nested `stream_events(version="v3")` call still
`stream_v2` / `astream_v2` populate from the caller's
checkpoint namespace so that a nested `stream_v2` call still
sees its own root snapshots.
"""
@@ -165,7 +165,7 @@ class MessagesTransformer(StreamTransformer):
metadata)` from `StreamMessagesHandler`):
1. Protocol event (dict with `"event"` key) emitted by
`stream_events(version="v3")` / `astream_events(version="v3")` via the `on_stream_event`
`stream_v2()` / `astream_v2()` via the `on_stream_event`
callback. Routed to an existing `ChatModelStream` by
`metadata["run_id"]`. A `message-start` event creates a new
stream; `message-finish` closes it.
@@ -177,15 +177,15 @@ class MessagesTransformer(StreamTransformer):
V1 `AIMessageChunk` tuples (from `on_llm_new_token`) are not
streamed into this projection: chat models that want to populate
`run.messages` with content-block streaming must use
`stream_events(version="v3")` / `astream_events(version="v3")`. Models called via the legacy
`stream_v2()` / `astream_v2()`. Models called via the legacy
`stream()` method still surface their final `AIMessage` via
`on_chain_end` when a node returns it as state.
Only events at the run's own level are projected; tokens from
deeper subgraphs are left in the main event log but excluded from
`.messages`. "Own level" is defined by `scope`, which
`stream_events(version="v3")` / `astream_events(version="v3")` populate from the caller's checkpoint
namespace so that a `stream_events(version="v3")` call inside a node still sees its
`stream_v2` / `astream_v2` populate from the caller's checkpoint
namespace so that a `stream_v2` call inside a node still sees its
own root chat model streams on `.messages`. Consumers that need
subgraph tokens should iterate the raw event stream or register a
custom transformer.
@@ -281,7 +281,7 @@ class MessagesTransformer(StreamTransformer):
):
self._route_whole_message(payload, node=node)
# Legacy AIMessageChunk tuples (from on_llm_new_token) are ignored;
# v1 streaming callers must switch to stream_events(version="v3") to populate this
# v1 streaming callers must switch to stream_v2() to populate this
# projection.
return True
@@ -327,7 +327,7 @@ class MessagesTransformer(StreamTransformer):
self._by_run.clear()
SubgraphStatus = Literal["started", "completed", "failed", "interrupted", "drained"]
SubgraphStatus = Literal["started", "completed", "failed", "interrupted"]
def _parse_ns_segment(segment: str) -> tuple[str, str | None]:
@@ -472,8 +472,10 @@ class _TasksLifecycleBase(StreamTransformer):
self._open.clear()
def fail(self, err: BaseException) -> None:
"""Emit terminal status for any tracked namespace still open."""
status, error_str = _status_from_exception(err)
"""Emit `failed` / `interrupted` for any tracked namespace still open."""
is_interrupt = isinstance(err, GraphInterrupt)
status: SubgraphStatus = "interrupted" if is_interrupt else "failed"
error_str = None if is_interrupt else str(err)
for ns in list(self._open):
self._on_terminal(ns, status, error_str)
self._open.clear()
@@ -481,8 +483,6 @@ class _TasksLifecycleBase(StreamTransformer):
def _status_from_exception(err: BaseException) -> tuple[SubgraphStatus, str | None]:
"""Map a run exception to a subgraph terminal status and error string."""
if isinstance(err, GraphDrained):
return "drained", None
if isinstance(err, GraphInterrupt):
return "interrupted", None
return "failed", str(err)
+3 -3
View File
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
[project]
name = "langgraph"
version = "1.2.0a3"
version = "1.1.10"
description = "Building stateful, multi-actor applications with LLMs"
authors = []
requires-python = ">=3.10"
@@ -24,8 +24,8 @@ classifiers = [
'Programming Language :: Python :: 3.13',
]
dependencies = [
"langchain-core>=1.4.0a2,<2",
"langgraph-checkpoint>=4.1.0a3,<5.0.0",
"langchain-core>=1.3.2,<2",
"langgraph-checkpoint>=4.0.3,<5.0.0",
"langgraph-sdk>=0.3.0,<0.4.0",
"langgraph-prebuilt>=1.0.12,<1.1.0",
"xxhash>=3.5.0",
+3 -2
View File
@@ -371,7 +371,8 @@ def test_delta_channel_inmemory_saver_assembles_writes() -> None:
saved = saver.get_tuple(config)
assert saved is not None
assert "messages" not in saved.checkpoint["channel_values"]
assert "messages" in saved.checkpoint["channel_values"]
assert saved.checkpoint["channel_values"]["messages"] is DELTA_SENTINEL
state = graph.get_state(config)
assert len(state.values["messages"]) == 4 # 2 human + 2 AI
@@ -561,7 +562,7 @@ def test_delta_channel_dict_reducer_end_to_end_filesystem() -> None:
saved = saver.get_tuple(config)
assert saved is not None
assert "files" not in saved.checkpoint["channel_values"]
assert saved.checkpoint["channel_values"]["files"] is DELTA_SENTINEL
state = graph.get_state(config)
assert state.values["files"] == {
"/doc_1.txt": "content for turn 1",
@@ -1,359 +0,0 @@
"""Tests for arrival-ordered interleave and push stamps."""
from __future__ import annotations
import operator
from typing import Annotated, Any
import pytest
from typing_extensions import TypedDict
from langgraph.constants import END, START
from langgraph.graph import StateGraph
from langgraph.stream import StreamChannel, StreamTransformer
from langgraph.stream._mux import StreamMux
from langgraph.stream._types import ProtocolEvent
from langgraph.stream.run_stream import GraphRunStream
from langgraph.stream.transformers import ValuesTransformer
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
class _TwoChannelTransformer(StreamTransformer):
"""Transformer that exposes two named channels for testing interleave."""
_native = True
def __init__(self, scope: tuple[str, ...] = ()) -> None:
super().__init__(scope)
self._alpha: StreamChannel[str] = StreamChannel("alpha")
self._beta: StreamChannel[str] = StreamChannel("beta")
def init(self) -> dict[str, Any]:
return {"alpha": self._alpha, "beta": self._beta}
def process(self, event: ProtocolEvent) -> bool:
return True
class SimpleState(TypedDict):
value: str
items: Annotated[list[str], operator.add]
def _build_simple_graph():
def node_a(state: SimpleState) -> dict:
return {"value": state["value"] + "A", "items": ["a"]}
def node_b(state: SimpleState) -> dict:
return {"value": state["value"] + "B", "items": ["b"]}
builder = StateGraph(SimpleState)
builder.add_node("node_a", node_a)
builder.add_node("node_b", node_b)
builder.add_edge(START, "node_a")
builder.add_edge("node_a", "node_b")
builder.add_edge("node_b", END)
return builder.compile()
# ---------------------------------------------------------------------------
# Unit tests: push stamps on StreamChannel
# ---------------------------------------------------------------------------
class TestPushStamps:
def test_stamps_are_monotonic_across_channels(self) -> None:
mux = StreamMux(
factories=[ValuesTransformer, _TwoChannelTransformer],
is_async=False,
)
alpha = mux.extensions["alpha"]
beta = mux.extensions["beta"]
alpha._subscribed = True
beta._subscribed = True
alpha.push("a1")
beta.push("b1")
alpha.push("a2")
beta.push("b2")
all_stamped = list(alpha._items) + list(beta._items)
stamps = [s for s, _ in all_stamped]
assert len(set(stamps)) == 4
items_by_arrival = [item for _, item in sorted(all_stamped)]
assert items_by_arrival == ["a1", "b1", "a2", "b2"]
def test_regular_iter_strips_stamps(self) -> None:
mux = StreamMux(
factories=[ValuesTransformer, _TwoChannelTransformer],
is_async=False,
)
alpha = mux.extensions["alpha"]
it = iter(alpha)
alpha.push("a1")
alpha.push("a2")
alpha.close()
items = list(it)
assert items == ["a1", "a2"]
assert all(isinstance(item, str) for item in items)
def test_events_channel_gets_real_stamps(self) -> None:
mux = StreamMux(
factories=[ValuesTransformer, _TwoChannelTransformer],
is_async=False,
)
alpha = mux.extensions["alpha"]
alpha._subscribed = True
alpha.push("a1")
mux._events._subscribed = True
mux._events.push({"method": "test", "data": "x"})
alpha.push("a2")
all_stamps = [s for s, _ in alpha._items] + [s for s, _ in mux._events._items]
assert len(set(all_stamps)) == len(all_stamps), "all stamps should be unique"
assert all(s > 0 for s in all_stamps), "no stamp should be zero"
def test_channel_without_mux_gets_zero_stamp(self) -> None:
ch: StreamChannel[str] = StreamChannel()
ch._bind(is_async=False)
ch._subscribed = True
ch.push("x")
assert list(ch._items) == [(0, "x")]
# ---------------------------------------------------------------------------
# Unit tests: interleave arrival order
# ---------------------------------------------------------------------------
class TestInterleaveArrivalOrder:
def test_arrival_order_not_round_robin(self) -> None:
mux = StreamMux(
factories=[ValuesTransformer, _TwoChannelTransformer],
is_async=False,
)
alpha = mux.extensions["alpha"]
beta = mux.extensions["beta"]
run = GraphRunStream(None, mux, wire_pump=False)
# interleave() subscribes channels directly and reads _items
# for stamp-ordered iteration. We simulate the pump by wiring
# a custom callback that pushes items in a known order.
push_script = [
("alpha", "a1"),
("alpha", "a2"),
("beta", "b1"),
("alpha", "a3"),
("beta", "b2"),
]
push_iter = iter(push_script)
channels = {"alpha": alpha, "beta": beta}
def fake_pump() -> bool:
try:
name, item = next(push_iter)
channels[name].push(item)
return True
except StopIteration:
mux.close()
return False
mux.bind_pump(fake_pump)
result = list(run.interleave("alpha", "beta"))
names = [name for name, _ in result]
items = [item for _, item in result]
assert items == ["a1", "a2", "b1", "a3", "b2"]
assert names == ["alpha", "alpha", "beta", "alpha", "beta"]
def test_single_projection(self) -> None:
mux = StreamMux(
factories=[ValuesTransformer, _TwoChannelTransformer],
is_async=False,
)
alpha = mux.extensions["alpha"]
run = GraphRunStream(None, mux, wire_pump=False)
push_script = [("alpha", "a1"), ("alpha", "a2")]
push_iter = iter(push_script)
def fake_pump() -> bool:
try:
_, item = next(push_iter)
alpha.push(item)
return True
except StopIteration:
mux.close()
return False
mux.bind_pump(fake_pump)
result = list(run.interleave("alpha"))
assert result == [("alpha", "a1"), ("alpha", "a2")]
def test_empty_projection(self) -> None:
mux = StreamMux(
factories=[ValuesTransformer, _TwoChannelTransformer],
is_async=False,
)
alpha = mux.extensions["alpha"]
run = GraphRunStream(None, mux, wire_pump=False)
push_script = [("alpha", "a1"), ("alpha", "a2")]
push_iter = iter(push_script)
channels = {"alpha": alpha}
def fake_pump() -> bool:
try:
name, item = next(push_iter)
channels[name].push(item)
return True
except StopIteration:
mux.close()
return False
mux.bind_pump(fake_pump)
result = list(run.interleave("alpha", "beta"))
assert result == [("alpha", "a1"), ("alpha", "a2")]
def test_unknown_projection_raises(self) -> None:
mux = StreamMux(
factories=[ValuesTransformer, _TwoChannelTransformer],
is_async=False,
)
run = GraphRunStream(None, mux, wire_pump=False)
mux.close()
with pytest.raises((KeyError, AttributeError)):
list(run.interleave("alpha", "does_not_exist"))
def test_all_empty(self) -> None:
mux = StreamMux(
factories=[ValuesTransformer, _TwoChannelTransformer],
is_async=False,
)
run = GraphRunStream(None, mux, wire_pump=False)
def fake_pump() -> bool:
mux.close()
return False
mux.bind_pump(fake_pump)
result = list(run.interleave("alpha", "beta"))
assert result == []
def test_error_propagation(self) -> None:
mux = StreamMux(
factories=[ValuesTransformer, _TwoChannelTransformer],
is_async=False,
)
alpha = mux.extensions["alpha"]
beta = mux.extensions["beta"]
run = GraphRunStream(None, mux, wire_pump=False)
err = RuntimeError("boom")
push_script = [
("alpha", "a1"),
("beta", "b1"),
]
push_iter = iter(push_script)
channels = {"alpha": alpha, "beta": beta}
def fake_pump() -> bool:
try:
name, item = next(push_iter)
channels[name].push(item)
return True
except StopIteration:
alpha.fail(err)
beta.close()
return False
mux.bind_pump(fake_pump)
collected = []
with pytest.raises(RuntimeError, match="boom"):
for pair in run.interleave("alpha", "beta"):
collected.append(pair)
assert ("alpha", "a1") in collected
assert ("beta", "b1") in collected
# ---------------------------------------------------------------------------
# Integration test: interleave with stream_events(version="v3")
# ---------------------------------------------------------------------------
class TestInterleaveIntegration:
def test_interleave_values_and_messages(self) -> None:
run = _build_simple_graph().stream_events(
{"value": "x", "items": []}, version="v3"
)
tagged = list(run.interleave("values", "messages"))
names = [name for name, _ in tagged]
assert set(names).issubset({"values", "messages"})
assert names.count("values") >= 1
def test_interleave_rejects_already_subscribed(self) -> None:
mux = StreamMux(
factories=[ValuesTransformer, _TwoChannelTransformer],
is_async=False,
)
alpha = mux.extensions["alpha"]
run = GraphRunStream(None, mux, wire_pump=False)
# Subscribe alpha via iter first
_ = iter(alpha)
mux.close()
with pytest.raises(RuntimeError, match="already has a subscriber"):
list(run.interleave("alpha"))
def test_interleave_releases_projections_on_completion(self) -> None:
run = _build_simple_graph().stream_events(
{"value": "x", "items": []}, version="v3"
)
list(run.interleave("values", "messages"))
# Subscriptions should be released after the generator completes,
# so the channels can be re-iterated (they'll be empty / closed).
assert run.extensions["values"]._subscribed is False
assert run.extensions["messages"]._subscribed is False
def test_interleave_releases_projections_on_early_break(self) -> None:
run = _build_simple_graph().stream_events(
{"value": "x", "items": []}, version="v3"
)
gen = run.interleave("values", "messages")
next(gen)
gen.close()
assert run.extensions["values"]._subscribed is False
assert run.extensions["messages"]._subscribed is False
def test_interleave_releases_projections_on_validation_failure(self) -> None:
mux = StreamMux(
factories=[ValuesTransformer, _TwoChannelTransformer],
is_async=False,
)
alpha = mux.extensions["alpha"]
# Pre-subscribe alpha so that interleave will fail validation when
# it gets to the second name. The first (already-validated) channel
# should still be released.
run = GraphRunStream(None, mux, wire_pump=False)
mux.close()
alpha._subscribed = True
with pytest.raises(RuntimeError, match="already has a subscriber"):
list(run.interleave("values", "alpha"))
assert mux.extensions["values"]._subscribed is False
-23
View File
@@ -122,29 +122,6 @@ def test_graph_validation() -> None:
graph.invoke({"hello": "there"})
def test_request_drain_allows_inflight_call_scheduling(
sync_checkpointer: BaseCheckpointSaver,
) -> None:
from langgraph.runtime import RunControl
@task
def child(x: int) -> int:
return x + 1
control = RunControl()
@entrypoint(checkpointer=sync_checkpointer)
def graph(x: int) -> int:
control.request_drain()
fut = child(x)
return fut.result()
config = {"configurable": {"thread_id": "drain-call-sync"}}
assert graph.invoke(1, config=config, control=control) == 2
assert control.drain_requested
def test_invalid_checkpointer_type() -> None:
class State(TypedDict):
foo: str
-149
View File
@@ -16,7 +16,6 @@ from typing import (
Literal,
Optional,
)
from unittest.mock import patch
from uuid import UUID
import pytest
@@ -49,7 +48,6 @@ from langgraph.channels.topic import Topic
from langgraph.errors import (
GraphRecursionError,
InvalidUpdateError,
NodeError,
ParentCommand,
)
from langgraph.func import entrypoint, task
@@ -217,30 +215,6 @@ async def test_checkpoint_errors() -> None:
pass
@NEEDS_CONTEXTVARS
async def test_request_drain_allows_inflight_acall_scheduling(
async_checkpointer: BaseCheckpointSaver,
) -> None:
from langgraph.runtime import RunControl
@task
async def child(x: int) -> int:
return x + 1
control = RunControl()
@entrypoint(checkpointer=async_checkpointer)
async def graph(x: int) -> int:
control.request_drain()
fut = child(x)
return await fut
config = {"configurable": {"thread_id": "drain-call-async"}}
assert await graph.ainvoke(1, config=config, control=control) == 2
assert control.drain_requested
async def test_py_async_with_cancel_behavior() -> None:
"""This test confirms that in all versions of Python we support, __aexit__
is not cancelled when the coroutine containing the async with block is cancelled."""
@@ -9765,126 +9739,3 @@ async def test_fork_does_not_apply_pending_writes(
# 1 (input) + 20 (forked node_a) + 100 (node_b) = 121
assert result == {"value": 121}
async def test_graph_error_handler_async_runtime_info() -> None:
class State(TypedDict):
foo: str
attempts = 0
captured: dict[str, object] = {}
async def always_failing_node(state: State) -> State:
nonlocal attempts
attempts += 1
raise ValueError("Always fails async")
async def err_handler_node(state: State, error: NodeError) -> State:
captured["from_node_name"] = error.node
captured["from_node_error"] = error.error
return {"foo": "handled_async"}
graph = (
StateGraph(State)
.add_node(
"always_failing",
always_failing_node,
retry_policy=RetryPolicy(
max_attempts=2,
initial_interval=0.01,
jitter=False,
retry_on=ValueError,
),
error_handler=err_handler_node,
)
.add_edge(START, "always_failing")
.compile()
)
with patch("asyncio.sleep"):
result = await graph.ainvoke({"foo": ""})
assert attempts == 2
assert result["foo"] == "handled_async"
assert captured["from_node_name"] == "always_failing"
assert isinstance(captured["from_node_error"], BaseException)
@NEEDS_CONTEXTVARS
async def test_graph_error_handler_does_not_swallow_interrupt_concurrent() -> None:
"""When a graph error handler is configured and a node calls interrupt()
concurrently with other nodes, the interrupt must still be raised not
silently swallowed."""
class State(TypedDict):
foo: str
async def node_a(state: State) -> State:
val = interrupt("need human input")
return {"foo": f"a_{val}"}
async def node_b(state: State) -> State:
return {}
async def err_handler(state: State) -> State:
return {"foo": "handled"}
checkpointer = InMemorySaver()
graph = (
StateGraph(State)
.add_node("node_a", node_a, error_handler=err_handler)
.add_node("node_b", node_b)
.add_edge(START, "node_a")
.add_edge(START, "node_b")
.compile(checkpointer=checkpointer)
)
config = {"configurable": {"thread_id": "test-interrupt-concurrent-async"}}
await graph.ainvoke({"foo": ""}, config)
state = await graph.aget_state(config)
assert len(state.tasks) > 0
interrupts = [t for t in state.tasks if hasattr(t, "interrupts") and t.interrupts]
assert len(interrupts) > 0, (
"GraphInterrupt was swallowed — interrupt() in node_a "
"should have paused execution"
)
async def test_node_error_handler_handles_subgraph_internal_failure_async() -> None:
class SubState(TypedDict):
foo: str
class ParentState(TypedDict):
foo: str
captured: dict[str, object] = {}
async def sub_fail_node(state: SubState) -> SubState:
raise ValueError("async subgraph boom")
async def parent_handler(state: ParentState, error: NodeError) -> ParentState:
captured["from_node_name"] = error.node
captured["from_node_error"] = error.error
return {"foo": "handled_async_subgraph"}
subgraph = (
StateGraph(SubState)
.add_node("sub_fail_node", sub_fail_node)
.add_edge(START, "sub_fail_node")
.compile()
)
parent_graph = (
StateGraph(ParentState)
.add_node("subgraph_node", subgraph, error_handler=parent_handler)
.add_edge(START, "subgraph_node")
.compile()
)
result = await parent_graph.ainvoke({"foo": ""})
assert result["foo"] == "handled_async_subgraph"
assert captured["from_node_name"] == "subgraph_node"
assert isinstance(captured["from_node_error"], BaseException)
@@ -1,4 +1,4 @@
"""Tests for Pregel.stream_events(version="v3") / astream_events(version="v3") and the transformer pipeline."""
"""Tests for Pregel.stream_v2 / astream_v2 and the transformer pipeline."""
from __future__ import annotations
@@ -125,7 +125,7 @@ def _build_custom_stream_graph():
class _CustomPassthroughTransformer(StreamTransformer):
"""Opts a run into the `custom` stream mode without building a projection.
`stream_events(version="v3")` requests only the modes that registered transformers
`stream_v2` requests only the modes that registered transformers
declare via `required_stream_modes`. Custom events are raw user
emissions from `StreamWriter`, so tests that want them visible on
the main event log register this pass-through transformer.
@@ -390,31 +390,25 @@ class TestStreamChannelNamed:
# ---------------------------------------------------------------------------
# stream_events(version="v3") sync tests
# stream_v2 sync tests
# ---------------------------------------------------------------------------
class TestStreamV2Sync:
def test_values_projection(self) -> None:
run = _build_simple_graph().stream_events(
{"value": "x", "items": []}, version="v3"
)
run = _build_simple_graph().stream_v2({"value": "x", "items": []})
snapshots = list(run.values)
assert len(snapshots) >= 1
last = snapshots[-1]
assert "A" in last["value"] and "B" in last["value"]
def test_output(self) -> None:
run = _build_simple_graph().stream_events(
{"value": "x", "items": []}, version="v3"
)
run = _build_simple_graph().stream_v2({"value": "x", "items": []})
output = run.output
assert output == {"value": "xAB", "items": ["a", "b"]}
def test_raw_event_iteration(self) -> None:
run = _build_simple_graph().stream_events(
{"value": "x", "items": []}, version="v3"
)
run = _build_simple_graph().stream_v2({"value": "x", "items": []})
events = list(run)
assert len(events) > 0
for event in events:
@@ -424,27 +418,22 @@ class TestStreamV2Sync:
assert isinstance(event["params"]["timestamp"], int)
def test_extensions_has_native_keys(self) -> None:
run = _build_simple_graph().stream_events(
{"value": "x", "items": []}, version="v3"
)
run = _build_simple_graph().stream_v2({"value": "x", "items": []})
_ = run.output
assert "values" in run.extensions and "messages" in run.extensions
assert run.values is run.extensions["values"]
assert run.messages is run.extensions["messages"]
def test_extensions_is_read_only(self) -> None:
run = _build_simple_graph().stream_events(
{"value": "x", "items": []}, version="v3"
)
run = _build_simple_graph().stream_v2({"value": "x", "items": []})
with pytest.raises(TypeError):
run.extensions["new_key"] = object() # type: ignore[index]
with pytest.raises(TypeError):
del run.extensions["values"] # type: ignore[attr-defined]
def test_custom_stream_events(self) -> None:
run = _build_custom_stream_graph().stream_events(
run = _build_custom_stream_graph().stream_v2(
{"value": "x", "items": []},
version="v3",
transformers=[_CustomPassthroughTransformer],
)
custom_events = [e for e in run if e["method"] == "custom"]
@@ -455,34 +444,27 @@ class TestStreamV2Sync:
def test_custom_events_suppressed_without_transformer(self) -> None:
"""Without a transformer declaring `"custom"`, no custom events flow.
`stream_events(version="v3")` asks the graph only for the modes that registered
`stream_v2` asks the graph only for the modes that registered
transformers require. Built-ins cover `values` / `messages`;
consumers that want raw custom events surface them by
registering a transformer whose `required_stream_modes`
includes `"custom"`.
"""
run = _build_custom_stream_graph().stream_events(
{"value": "x", "items": []}, version="v3"
)
run = _build_custom_stream_graph().stream_v2({"value": "x", "items": []})
custom_events = [e for e in run if e["method"] == "custom"]
assert custom_events == []
def test_interleave_values_and_messages(self) -> None:
run = _build_simple_graph().stream_events(
{"value": "x", "items": []}, version="v3"
)
run = _build_simple_graph().stream_v2({"value": "x", "items": []})
tagged = list(run.interleave("values", "messages"))
names = [name for name, _ in tagged]
assert set(names).issubset({"values", "messages"})
assert names.count("values") >= 1
# interleave releases its subscription on completion.
assert run.extensions["values"]._subscribed is False
assert run.extensions["messages"]._subscribed is False
with pytest.raises(RuntimeError, match="already has a subscriber"):
list(run.values)
def test_abort_marks_exhausted_and_closes_mux(self) -> None:
run = _build_simple_graph().stream_events(
{"value": "x", "items": []}, version="v3"
)
run = _build_simple_graph().stream_v2({"value": "x", "items": []})
values_iter = iter(run.values)
_ = next(values_iter)
run.abort()
@@ -491,63 +473,48 @@ class TestStreamV2Sync:
run.abort() # idempotent
def test_context_manager_calls_abort_on_exit(self) -> None:
with _build_simple_graph().stream_events(
{"value": "x", "items": []}, version="v3"
) as run:
with _build_simple_graph().stream_v2({"value": "x", "items": []}) as run:
_ = next(iter(run.values))
assert run._exhausted is True
def test_interleave_unknown_projection(self) -> None:
run = _build_simple_graph().stream_events(
{"value": "x", "items": []}, version="v3"
)
run = _build_simple_graph().stream_v2({"value": "x", "items": []})
with pytest.raises(KeyError):
list(run.interleave("values", "does_not_exist"))
class TestStreamV2SyncErrors:
def test_error_propagation_output(self) -> None:
run = _build_error_graph().stream_events(
{"value": "x", "items": []}, version="v3"
)
run = _build_error_graph().stream_v2({"value": "x", "items": []})
with pytest.raises(ValueError, match="boom"):
_ = run.output
def test_error_propagation_values(self) -> None:
run = _build_error_graph().stream_events(
{"value": "x", "items": []}, version="v3"
)
run = _build_error_graph().stream_v2({"value": "x", "items": []})
with pytest.raises(ValueError, match="boom"):
list(run.values)
def test_error_propagation_raw_events(self) -> None:
run = _build_error_graph().stream_events(
{"value": "x", "items": []}, version="v3"
)
run = _build_error_graph().stream_v2({"value": "x", "items": []})
with pytest.raises(ValueError, match="boom"):
list(run)
def test_error_propagation_interrupted(self) -> None:
run = _build_error_graph().stream_events(
{"value": "x", "items": []}, version="v3"
)
run = _build_error_graph().stream_v2({"value": "x", "items": []})
with pytest.raises(ValueError, match="boom"):
_ = run.interrupted
def test_error_propagation_interrupts(self) -> None:
run = _build_error_graph().stream_events(
{"value": "x", "items": []}, version="v3"
)
run = _build_error_graph().stream_v2({"value": "x", "items": []})
with pytest.raises(ValueError, match="boom"):
_ = run.interrupts
class TestStreamV2SyncInterrupt:
def test_interrupted(self) -> None:
run = _build_interrupt_graph().stream_events(
run = _build_interrupt_graph().stream_v2(
{"value": "x", "items": []},
{"configurable": {"thread_id": "t1"}},
version="v3",
)
_ = run.output
assert run.interrupted is True
@@ -555,7 +522,7 @@ class TestStreamV2SyncInterrupt:
# ---------------------------------------------------------------------------
# astream_events(version="v3") async tests
# astream_v2 async tests
# ---------------------------------------------------------------------------
@@ -563,34 +530,26 @@ class TestStreamV2SyncInterrupt:
@NEEDS_CONTEXTVARS
class TestStreamV2Async:
async def test_values_projection(self) -> None:
run = await _build_simple_graph().astream_events(
{"value": "x", "items": []}, version="v3"
)
run = await _build_simple_graph().astream_v2({"value": "x", "items": []})
snapshots = [s async for s in run.values]
assert len(snapshots) >= 1
last = snapshots[-1]
assert "A" in last["value"] and "B" in last["value"]
async def test_output(self) -> None:
run = await _build_simple_graph().astream_events(
{"value": "x", "items": []}, version="v3"
)
run = await _build_simple_graph().astream_v2({"value": "x", "items": []})
output = await run.output()
assert output == {"value": "xAB", "items": ["a", "b"]}
async def test_raw_event_iteration(self) -> None:
run = await _build_simple_graph().astream_events(
{"value": "x", "items": []}, version="v3"
)
run = await _build_simple_graph().astream_v2({"value": "x", "items": []})
events = [e async for e in run]
assert len(events) > 0
for event in events:
assert event["type"] == "event"
async def test_abort_marks_exhausted_and_closes_mux(self) -> None:
run = await _build_simple_graph().astream_events(
{"value": "x", "items": []}, version="v3"
)
run = await _build_simple_graph().astream_v2({"value": "x", "items": []})
values_iter = aiter(run.values)
_ = await anext(values_iter)
await run.abort()
@@ -600,26 +559,21 @@ class TestStreamV2Async:
await run.abort() # idempotent
async def test_context_manager_calls_abort_on_exit(self) -> None:
run = await _build_simple_graph().astream_events(
{"value": "x", "items": []}, version="v3"
)
run = await _build_simple_graph().astream_v2({"value": "x", "items": []})
async with run:
_ = await anext(aiter(run.values))
assert run._exhausted is True
async def test_extensions_has_native_keys(self) -> None:
run = await _build_simple_graph().astream_events(
{"value": "x", "items": []}, version="v3"
)
run = await _build_simple_graph().astream_v2({"value": "x", "items": []})
_ = await run.output()
assert "values" in run.extensions and "messages" in run.extensions
assert run.values is run.extensions["values"]
assert run.messages is run.extensions["messages"]
async def test_custom_stream_events(self) -> None:
run = await _build_custom_stream_graph().astream_events(
run = await _build_custom_stream_graph().astream_v2(
{"value": "x", "items": []},
version="v3",
transformers=[_CustomPassthroughTransformer],
)
events = [e async for e in run]
@@ -633,39 +587,29 @@ class TestStreamV2Async:
@NEEDS_CONTEXTVARS
class TestStreamV2AsyncErrors:
async def test_error_propagation_output(self) -> None:
run = await _build_error_graph().astream_events(
{"value": "x", "items": []}, version="v3"
)
run = await _build_error_graph().astream_v2({"value": "x", "items": []})
with pytest.raises(ValueError, match="boom"):
await run.output()
async def test_error_propagation_values(self) -> None:
run = await _build_error_graph().astream_events(
{"value": "x", "items": []}, version="v3"
)
run = await _build_error_graph().astream_v2({"value": "x", "items": []})
with pytest.raises(ValueError, match="boom"):
async for _ in run.values:
pass
async def test_error_propagation_raw_events(self) -> None:
run = await _build_error_graph().astream_events(
{"value": "x", "items": []}, version="v3"
)
run = await _build_error_graph().astream_v2({"value": "x", "items": []})
with pytest.raises(ValueError, match="boom"):
async for _ in run:
pass
async def test_error_propagation_interrupted(self) -> None:
run = await _build_error_graph().astream_events(
{"value": "x", "items": []}, version="v3"
)
run = await _build_error_graph().astream_v2({"value": "x", "items": []})
with pytest.raises(ValueError, match="boom"):
await run.interrupted()
async def test_error_propagation_interrupts(self) -> None:
run = await _build_error_graph().astream_events(
{"value": "x", "items": []}, version="v3"
)
run = await _build_error_graph().astream_v2({"value": "x", "items": []})
with pytest.raises(ValueError, match="boom"):
await run.interrupts()
@@ -674,10 +618,9 @@ class TestStreamV2AsyncErrors:
@NEEDS_CONTEXTVARS
class TestStreamV2AsyncInterrupt:
async def test_interrupted(self) -> None:
run = await _build_interrupt_graph().astream_events(
run = await _build_interrupt_graph().astream_v2(
{"value": "x", "items": []},
{"configurable": {"thread_id": "t2"}},
version="v3",
)
_ = await run.output()
assert await run.interrupted() is True
@@ -1041,8 +984,8 @@ class TestCustomTransformer:
self._channel.push(self._count)
return True
run = _build_simple_graph().stream_events(
{"value": "x", "items": []}, version="v3", transformers=[CounterTransformer]
run = _build_simple_graph().stream_v2(
{"value": "x", "items": []}, transformers=[CounterTransformer]
)
assert "counter" in run.extensions
counter_iter = iter(run.extensions["counter"])
@@ -1067,15 +1010,15 @@ class TestCustomTransformer:
self._log.push("saw_values")
return True
run = _build_simple_graph().stream_events(
{"value": "x", "items": []}, version="v3", transformers=[FooTransformer]
run = _build_simple_graph().stream_v2(
{"value": "x", "items": []}, transformers=[FooTransformer]
)
foo_iter = iter(run.foo)
_ = run.output
assert "foo" in run.extensions and run.foo is run.extensions["foo"]
assert "saw_values" in list(foo_iter)
def test_stream_events_v3_rejects_transformer_instances(self) -> None:
def test_stream_v2_rejects_transformer_instances(self) -> None:
class InstanceTransformer(StreamTransformer):
def init(self) -> dict[str, Any]:
return {}
@@ -1084,10 +1027,8 @@ class TestCustomTransformer:
return True
with pytest.raises(TypeError, match="pre-built instance"):
_build_simple_graph().stream_events(
{"value": "x", "items": []},
version="v3",
transformers=[InstanceTransformer()],
_build_simple_graph().stream_v2(
{"value": "x", "items": []}, transformers=[InstanceTransformer()]
)
def test_stream_channel_auto_forward(self) -> None:
@@ -1106,8 +1047,8 @@ class TestCustomTransformer:
self._channel.push("emitted")
return True
run = _build_simple_graph().stream_events(
{"value": "x", "items": []}, version="v3", transformers=[EmitterTransformer]
run = _build_simple_graph().stream_v2(
{"value": "x", "items": []}, transformers=[EmitterTransformer]
)
custom_events = [e for e in run if e["method"] == "custom:emitter"]
assert len(custom_events) > 0
@@ -1151,10 +1092,8 @@ class TestCustomTransformer:
return True
with pytest.raises(ValueError, match=r"conflict.*'values'.*ValuesTransformer"):
_build_simple_graph().stream_events(
{"value": "x", "items": []},
version="v3",
transformers=[ConflictTransformer],
_build_simple_graph().stream_v2(
{"value": "x", "items": []}, transformers=[ConflictTransformer]
)
@@ -1236,8 +1175,8 @@ class TestStreamChannelAutoLifecycle:
self._log.push("got_it")
return True
run = _build_simple_graph().stream_events(
{"value": "x", "items": []}, version="v3", transformers=[MinimalTransformer]
run = _build_simple_graph().stream_v2(
{"value": "x", "items": []}, transformers=[MinimalTransformer]
)
minimal_iter = iter(run.extensions["minimal"])
_ = run.output
@@ -1497,8 +1436,8 @@ class TestAsyncTransformerLane:
async def afinalize(self) -> None:
self._log.close()
run = await _build_simple_graph().astream_events(
{"value": "x", "items": []}, version="v3", transformers=[Scorer]
run = await _build_simple_graph().astream_v2(
{"value": "x", "items": []}, transformers=[Scorer]
)
scores_cursor = aiter(run.extensions["scores"])
_ = await run.output()
@@ -1514,9 +1453,7 @@ class TestAsyncTransformerLane:
@NEEDS_CONTEXTVARS
class TestMemoryBounds:
def test_sync_subscribed_buffer_stays_at_most_one_between_yields(self) -> None:
run = _build_simple_graph().stream_events(
{"value": "x", "items": []}, version="v3"
)
run = _build_simple_graph().stream_v2({"value": "x", "items": []})
events_iter = iter(run)
max_buffered = 0
count = 0
@@ -1529,9 +1466,7 @@ class TestMemoryBounds:
)
def test_unsubscribed_projections_never_accumulate(self) -> None:
run = _build_simple_graph().stream_events(
{"value": "x", "items": []}, version="v3"
)
run = _build_simple_graph().stream_v2({"value": "x", "items": []})
list(run)
values_log = run.extensions["values"]
messages_log = run.extensions["messages"]
@@ -1539,25 +1474,19 @@ class TestMemoryBounds:
assert len(messages_log._items) == 0 and not messages_log._subscribed
def test_output_path_does_not_retain_values(self) -> None:
run = _build_simple_graph().stream_events(
{"value": "x", "items": []}, version="v3"
)
run = _build_simple_graph().stream_v2({"value": "x", "items": []})
_ = run.output
values_log = run.extensions["values"]
assert len(values_log._items) == 0 and not values_log._subscribed
def test_drained_subscriber_buffer_returns_to_empty(self) -> None:
run = _build_simple_graph().stream_events(
{"value": "x", "items": []}, version="v3"
)
run = _build_simple_graph().stream_v2({"value": "x", "items": []})
list(run.values)
assert len(run.extensions["values"]._items) == 0
@pytest.mark.anyio
async def test_async_single_consumer_buffer_stays_at_most_one(self) -> None:
run = await _build_simple_graph().astream_events(
{"value": "x", "items": []}, version="v3"
)
run = await _build_simple_graph().astream_v2({"value": "x", "items": []})
max_buffered = 0
count = 0
async for _ in run:
@@ -1568,9 +1497,7 @@ class TestMemoryBounds:
@pytest.mark.anyio
async def test_async_unsubscribed_projections_never_accumulate(self) -> None:
run = await _build_simple_graph().astream_events(
{"value": "x", "items": []}, version="v3"
)
run = await _build_simple_graph().astream_v2({"value": "x", "items": []})
_ = await run.output()
values_log = run.extensions["values"]
messages_log = run.extensions["messages"]
+4 -516
View File
@@ -1,5 +1,4 @@
import asyncio
import operator
import sys
import threading
import time
@@ -16,7 +15,7 @@ from langchain_core.language_models.fake_chat_models import GenericFakeChatModel
from langchain_core.messages import AIMessage, AIMessageChunk, BaseMessage, HumanMessage
from langchain_core.outputs import ChatGeneration, ChatGenerationChunk, ChatResult
from langchain_core.runnables import RunnableLambda, RunnableParallel
from langgraph.checkpoint.memory import InMemorySaver, MemorySaver
from langgraph.checkpoint.memory import MemorySaver
from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
from typing_extensions import TypedDict
@@ -35,7 +34,7 @@ from langgraph._internal._runnable import RunnableCallable
from langgraph._internal._timeout import coerce_timeout_policy
from langgraph.channels.ephemeral_value import EphemeralValue
from langgraph.channels.last_value import LastValue
from langgraph.errors import GraphInterrupt, NodeError, NodeTimeoutError, ParentCommand
from langgraph.errors import GraphInterrupt, NodeTimeoutError, ParentCommand
from langgraph.func import entrypoint, task
from langgraph.graph import END, START, StateGraph, add_messages
from langgraph.pregel import NodeBuilder, Pregel
@@ -210,15 +209,6 @@ def test_should_retry_default_retry_on():
req_error_no_resp.response = None
assert _should_retry_on(policy, req_error_no_resp) is True
# NodeTimeoutError should be retryable by default
assert (
_should_retry_on(
policy,
NodeTimeoutError("node", 1.0, kind="run", run_timeout=0.5),
)
is True
)
# Should retry on other exceptions by default
class CustomException(Exception):
pass
@@ -1465,14 +1455,14 @@ async def test_state_graph_add_node_timeout_composes_with_retry():
async def flaky(state: _TimeoutState) -> _TimeoutState:
attempts.append(len(attempts))
if len(attempts) < 2:
await asyncio.sleep(1.0)
await asyncio.sleep(0.5)
return {"x": state["x"] + 1}
builder = StateGraph(_TimeoutState)
builder.add_node(
"flaky",
flaky,
timeout=TimeoutPolicy(idle_timeout=0.3),
timeout=TimeoutPolicy(idle_timeout=0.1),
retry_policy=RetryPolicy(
max_attempts=3,
initial_interval=0.0,
@@ -1765,505 +1755,3 @@ async def test_arun_with_retry_timeout_observer_treats_bubble_up_as_non_error():
assert finish.status == "success"
assert finish.error_type is None
assert finish.error_message is None
# ---------------------------------------------------------------------------
# Watcher invariant: any timeout that retry/error_handler can recover from
# MUST emit `finish=error` BEFORE the in-process recovery work happens. The
# external watchdog (langgraph-api) relies on this so it only kills a worker
# when no `finish` arrives within the deadline. The tests below pin down the
# three recovery paths so a refactor that moves `_finish_timed_attempt` past
# an `await` (or past the final `raise`) trips CI.
# ---------------------------------------------------------------------------
@pytest.mark.anyio
async def test_arun_with_retry_observer_emits_finish_before_retry_backoff():
"""`finish=error` of attempt N must arrive before the retry backoff sleep."""
timeline: list[tuple[float, Any]] = []
class TimingOutOnceProc:
def __init__(self) -> None:
self.calls = 0
async def ainvoke(self, input, config):
self.calls += 1
if self.calls == 1:
await asyncio.sleep(1.0)
return "ok"
backoff = 0.25
policy = RetryPolicy(
max_attempts=2,
initial_interval=backoff,
backoff_factor=1.0,
max_interval=backoff,
jitter=False,
retry_on=NodeTimeoutError,
)
task = _make_task(
TimingOutOnceProc(),
timeout=_idle_timeout(0.05),
retry_policy=(policy,),
name="backoff_watcher",
)
task.config[CONF][CONFIG_KEY_TIMED_ATTEMPT_OBSERVER] = lambda ev: timeline.append(
(time.monotonic(), ev)
)
assert await arun_with_retry(task, retry_policy=None) == "ok"
starts = [(t, ev) for t, ev in timeline if ev.event == "start"]
finishes = [(t, ev) for t, ev in timeline if ev.event == "finish"]
assert [ev.context.attempt for _, ev in starts] == [1, 2]
assert [ev.status for _, ev in finishes] == ["error", "success"]
first_finish_t = finishes[0][0]
second_start_t = starts[1][0]
# The watcher relies on this gap: `finish=error` for attempt 1 must arrive
# before `arun_with_retry` enters `await asyncio.sleep(backoff)`. We give a
# generous slack to keep this stable on slow CI; the structural invariant
# is "finish lands first", not "the gap equals exactly backoff".
assert second_start_t - first_finish_t >= backoff * 0.5, (
f"finish=error appears to be emitted after retry backoff sleep; "
f"gap was {second_start_t - first_finish_t:.3f}s, expected >= {backoff * 0.5:.3f}s"
)
@pytest.mark.anyio
async def test_state_graph_observer_emits_finish_before_error_handler_start():
"""Original task's `finish=error` must arrive before the error_handler task's `start`."""
class State(TypedDict):
foo: str
async def slow_node(state: State) -> State:
await asyncio.sleep(1.0)
return {"foo": "should-not-happen"}
async def handler_node(state: State, error: NodeError) -> State:
return {"foo": "handled"}
events: list = []
graph = (
StateGraph(State)
.add_node(
"slow",
slow_node,
timeout=TimeoutPolicy(idle_timeout=0.05),
error_handler=handler_node,
)
.add_edge(START, "slow")
.compile()
)
result = await graph.ainvoke(
{"foo": ""},
config={
"configurable": {CONFIG_KEY_TIMED_ATTEMPT_OBSERVER: events.append},
},
)
assert result["foo"] == "handled"
# Filter to events from the failing node only — the handler node has no
# timeout configured here, so it doesn't appear in the observer stream.
slow_events = [ev for ev in events if ev.context.task_name == "slow"]
starts = [ev for ev in slow_events if ev.event == "start"]
finishes = [ev for ev in slow_events if ev.event == "finish"]
assert len(starts) == 1
assert len(finishes) == 1
assert finishes[0].status == "error"
assert finishes[0].error_type == "NodeTimeoutError"
# The slow task's finish-error event must precede every event for any
# follow-up task in the same observer stream.
slow_finish_index = events.index(finishes[0])
for ev in events[slow_finish_index + 1 :]:
assert ev.context.task_name == "slow" or ev.event == "start", (
f"unexpected event {ev.event} for {ev.context.task_name} "
f"after slow's finish=error"
)
@pytest.mark.anyio
async def test_arun_with_retry_observer_emits_finish_before_final_raise_on_exhaustion():
"""When retry exhausts and the timeout propagates, the final `finish=error` must
be emitted before `arun_with_retry` re-raises."""
events: list = []
class AlwaysTimingOutProc:
async def ainvoke(self, input, config):
await asyncio.sleep(1.0)
return "never"
policy = RetryPolicy(
max_attempts=2,
initial_interval=0.0,
jitter=False,
retry_on=NodeTimeoutError,
)
task = _make_task(
AlwaysTimingOutProc(),
timeout=_idle_timeout(0.05),
retry_policy=(policy,),
name="never_succeeds",
)
task.config[CONF][CONFIG_KEY_TIMED_ATTEMPT_OBSERVER] = events.append
with pytest.raises(NodeTimeoutError):
await arun_with_retry(task, retry_policy=None)
starts = [ev for ev in events if ev.event == "start"]
finishes = [ev for ev in events if ev.event == "finish"]
assert [ev.context.attempt for ev in starts] == [1, 2]
assert [ev.context.attempt for ev in finishes] == [1, 2]
assert [ev.status for ev in finishes] == ["error", "error"]
assert all(ev.error_type == "NodeTimeoutError" for ev in finishes)
# Both finish events were observed BEFORE arun_with_retry raised, otherwise
# the `with pytest.raises` block would have exited before `events` got
# populated with the second finish.
@pytest.mark.anyio
async def test_sync_sleep_in_async_node_bypasses_timeout_and_emits_finish_success():
"""Sync `time.sleep` inside an async node blocks the event loop so the
in-process watchdog cannot fire. We document the resulting behavior here:
1. `NodeTimeoutError` is NOT raised, even though the sync sleep exceeds
`idle_timeout`.
2. The node's normal return value flows through.
3. `finish=success` is emitted to the observer.
This is the canonical case where the in-process timeout is defeated and
the only safety net is the external watcher (langgraph-api), which
SIGKILLs the worker when no `finish` arrives within its deadline. The
catch is that with a *short* sync sleep the event loop unblocks before
the watcher's deadline expires, so the watcher legitimately does not
kill meaning the configured `idle_timeout` is silently honored at the
process level only when the block is long enough to outlast the
watcher's grace.
This is the documented "Cooperative cancellation" caveat on
`TimeoutPolicy`. The test pins the behavior so any future change that
starts raising `NodeTimeoutError` for sync-blocked async nodes (or stops
emitting `finish=success`) is caught.
"""
events: list = []
class SyncSleepingProc:
async def ainvoke(self, input, config):
time.sleep(0.1)
return "completed_despite_timeout"
task = _make_task(
SyncSleepingProc(),
timeout=_idle_timeout(0.05),
name="sync_sleeper",
)
task.config[CONF][CONFIG_KEY_TIMED_ATTEMPT_OBSERVER] = events.append
result = await arun_with_retry(task, retry_policy=None)
assert result == "completed_despite_timeout"
starts = [ev for ev in events if ev.event == "start"]
finishes = [ev for ev in events if ev.event == "finish"]
assert len(starts) == 1
assert len(finishes) == 1
assert finishes[0].status == "success"
assert finishes[0].error_type is None
def test_graph_error_handler_runs_after_retry_exhaustion():
class State(TypedDict):
foo: str
attempts = 0
captured: dict[str, object] = {}
def always_failing_node(state: State) -> State:
nonlocal attempts
attempts += 1
raise ValueError("Always fails")
def err_handler_node(state: State, error: NodeError) -> Command:
captured["from_node_name"] = error.node
captured["from_node_error"] = error.error
return Command(update={"foo": "handled"}, goto="after_handler")
def after_handler(state: State) -> State:
return {"foo": f"{state['foo']}_after"}
retry_policy = RetryPolicy(
max_attempts=2,
initial_interval=0.01,
jitter=False,
retry_on=ValueError,
)
graph = (
StateGraph(State)
.add_node(
"always_failing",
always_failing_node,
retry_policy=retry_policy,
error_handler=err_handler_node,
)
.add_node("after_handler", after_handler)
.add_edge(START, "always_failing")
.compile()
)
with patch("time.sleep"):
result = graph.invoke({"foo": ""})
assert attempts == 2
assert result["foo"] == "handled_after"
assert captured["from_node_name"] == "always_failing"
assert isinstance(captured["from_node_error"], BaseException)
def test_graph_error_handler_can_route_with_command():
class State(TypedDict):
foo: str
attempts = 0
def always_failing_node(state: State) -> State:
nonlocal attempts
attempts += 1
raise ValueError("Always fails")
def err_handler_node(state: State) -> Command:
return Command(update={"foo": "handled"}, goto="next_node")
def next_node(state: State) -> State:
return {"foo": f"{state['foo']}_next"}
retry_policy = RetryPolicy(
max_attempts=1,
initial_interval=0.01,
jitter=False,
retry_on=ValueError,
)
graph = (
StateGraph(State)
.add_node(
"always_failing",
always_failing_node,
retry_policy=retry_policy,
error_handler=err_handler_node,
)
.add_node("next_node", next_node)
.add_edge(START, "always_failing")
.compile()
)
result = graph.invoke({"foo": ""})
assert attempts == 1
assert result["foo"] == "handled_next"
def test_graph_error_handler_failure_fails_run():
class State(TypedDict):
foo: str
def always_failing_node(state: State) -> State:
raise ValueError("Always fails")
def err_handler_node(state: State) -> State:
raise RuntimeError("handler failed")
graph = (
StateGraph(State)
.add_node("always_failing", always_failing_node, error_handler=err_handler_node)
.add_edge(START, "always_failing")
.compile()
)
with pytest.raises(RuntimeError, match="handler failed"):
graph.invoke({"foo": ""})
def test_graph_error_handler_handles_subgraph_internal_failure():
class SubState(TypedDict):
foo: str
class ParentState(TypedDict):
foo: str
parent_handler_called = False
captured: dict[str, object] = {}
def sub_fail_node(state: SubState) -> SubState:
raise ValueError("subgraph boom")
def parent_handler(state: ParentState, error: NodeError) -> ParentState:
nonlocal parent_handler_called
parent_handler_called = True
captured["from_node_name"] = error.node
captured["from_node_error"] = error.error
return {"foo": "handled_by_parent"}
subgraph = (
StateGraph(SubState)
.add_node("sub_fail_node", sub_fail_node)
.add_edge(START, "sub_fail_node")
.compile()
)
parent_graph = (
StateGraph(ParentState)
.add_node("subgraph_node", subgraph, error_handler=parent_handler)
.add_edge(START, "subgraph_node")
.compile()
)
result = parent_graph.invoke({"foo": ""})
assert result["foo"] == "handled_by_parent"
assert parent_handler_called is True
assert captured["from_node_name"] == "subgraph_node"
assert isinstance(captured["from_node_error"], BaseException)
def test_graph_error_handler_error_context_survives_checkpoint_resume():
class State(TypedDict):
foo: str
captured: dict[str, object] = {}
def always_failing_node(state: State) -> State:
raise RuntimeError("failed before handler")
def err_handler_node(state: State, error: NodeError) -> State:
captured["from_node_name"] = error.node
captured["from_node_error"] = error.error
return {"foo": "handled_after_resume"}
checkpointer = InMemorySaver()
config = {"configurable": {"thread_id": "graph-error-resume"}}
graph = (
StateGraph(State)
.add_node("always_failing", always_failing_node, error_handler=err_handler_node)
.add_edge(START, "always_failing")
.compile(
checkpointer=checkpointer,
interrupt_before=["__error_handler__always_failing"],
)
)
# First run pauses before handler, after failure context is checkpointed.
graph.invoke({"foo": ""}, config)
# Resume should execute handler and recover serialized error context.
result = graph.invoke(None, config)
assert result["foo"] == "handled_after_resume"
assert captured["from_node_name"] == "always_failing"
assert isinstance(captured["from_node_error"], BaseException)
def test_graph_error_handler_does_not_swallow_interrupt_concurrent():
"""When a graph error handler is configured and a node calls interrupt()
concurrently with other nodes, the interrupt must still be raised not
silently swallowed."""
from langgraph.types import interrupt
class State(TypedDict):
foo: str
def node_a(state: State) -> State:
# This node uses interrupt() which raises GraphInterrupt
val = interrupt("need human input")
return {"foo": f"a_{val}"}
def node_b(state: State) -> State:
return {}
def err_handler(state: State) -> State:
return {"foo": "handled"}
checkpointer = InMemorySaver()
graph = (
StateGraph(State)
.add_node("node_a", node_a, error_handler=err_handler)
.add_node("node_b", node_b)
# Fan-out: both node_a and node_b run concurrently
.add_edge(START, "node_a")
.add_edge(START, "node_b")
.compile(checkpointer=checkpointer)
)
config = {"configurable": {"thread_id": "test-interrupt-concurrent"}}
# First invoke should pause at the interrupt, not silently complete
graph.invoke({"foo": ""}, config)
# The graph should have an interrupt pending
state = graph.get_state(config)
assert len(state.tasks) > 0
# There should be a pending interrupt from node_a
interrupts = [t for t in state.tasks if hasattr(t, "interrupts") and t.interrupts]
assert len(interrupts) > 0, (
"GraphInterrupt was swallowed — interrupt() in node_a "
"should have paused execution"
)
def test_node_error_handlers_route_to_matching_handler():
class State(TypedDict):
route: str
foo: Annotated[list[str], operator.add]
def route_node(state: State) -> State:
return {"foo": []}
def choose_node(state: State) -> str:
return state["route"]
def fail_a(state: State) -> State:
raise ValueError("a failed")
def fail_b(state: State) -> State:
raise RuntimeError("b failed")
def handler_a(state: State, error: NodeError) -> State:
assert error.node == "fail_a"
return {"foo": ["handled_a"]}
def handler_b(state: State, error: NodeError) -> State:
assert error.node == "fail_b"
return {"foo": ["handled_b"]}
graph = (
StateGraph(State)
.add_node("route_node", route_node)
.add_node("fail_a", fail_a, error_handler=handler_a)
.add_node("fail_b", fail_b, error_handler=handler_b)
.add_edge(START, "route_node")
.add_conditional_edges("route_node", choose_node, path_map=["fail_a", "fail_b"])
.compile()
)
result_a = graph.invoke({"route": "fail_a", "foo": []})
result_b = graph.invoke({"route": "fail_b", "foo": []})
assert result_a["foo"] == ["handled_a"]
assert result_b["foo"] == ["handled_b"]
def test_node_without_error_handler_still_fails_run():
class State(TypedDict):
foo: str
def fail_without_handler(state: State) -> State:
raise ValueError("no handler")
graph = (
StateGraph(State)
.add_node("fail_without_handler", fail_without_handler)
.add_edge(START, "fail_without_handler")
.compile()
)
with pytest.raises(ValueError, match="no handler"):
graph.invoke({"foo": ""})
+1 -516
View File
@@ -1,6 +1,3 @@
import asyncio
import threading
import time
from dataclasses import dataclass
from typing import Any
@@ -9,15 +6,8 @@ from langgraph.checkpoint.memory import MemorySaver
from pydantic import BaseModel, ValidationError
from typing_extensions import TypedDict
from langgraph.errors import GraphDrained
from langgraph.graph import END, START, StateGraph
from langgraph.runtime import (
ExecutionInfo,
RunControl,
Runtime,
ServerInfo,
get_runtime,
)
from langgraph.runtime import ExecutionInfo, Runtime, ServerInfo, get_runtime
def test_injected_runtime() -> None:
@@ -89,183 +79,6 @@ def test_merge_runtime() -> None:
assert runtime1.merge(runtime3).context.api_key == "abc" # type: ignore
def test_merge_runtime_preserves_run_control() -> None:
control = RunControl()
runtime1 = Runtime(control=control)
runtime2 = Runtime(context=None)
assert runtime1.merge(runtime2).control is control
def test_run_control_request_drain_stops_future_steps() -> None:
class State(TypedDict, total=False):
first: str
second: str
control = RunControl()
def first_node(state: State) -> dict[str, str]:
control.request_drain()
return {"first": "done"}
def second_node(state: State) -> dict[str, str]:
return {"second": "should-not-run"}
graph = StateGraph(State)
graph.add_node("first", first_node)
graph.add_node("second", second_node)
graph.add_edge(START, "first")
graph.add_edge("first", "second")
graph.add_edge("second", END)
with pytest.raises(GraphDrained, match="shutdown"):
graph.compile().invoke({}, control=control)
@pytest.mark.anyio
async def test_run_control_request_drain_stops_future_steps_async() -> None:
class State(TypedDict, total=False):
first: str
second: str
control = RunControl()
async def first_node(state: State) -> dict[str, str]:
control.request_drain()
return {"first": "done"}
async def second_node(state: State) -> dict[str, str]:
return {"second": "should-not-run"}
graph = StateGraph(State)
graph.add_node("first", first_node)
graph.add_node("second", second_node)
graph.add_edge(START, "first")
graph.add_edge("first", "second")
graph.add_edge("second", END)
with pytest.raises(GraphDrained, match="shutdown"):
await graph.compile().ainvoke({}, control=control)
def test_drain_requested_in_terminal_step_finishes_normally() -> None:
class State(TypedDict, total=False):
value: str
control = RunControl()
def node(state: State) -> dict[str, str]:
control.request_drain()
return {"value": "done"}
graph = StateGraph(State)
graph.add_node("node", node)
graph.add_edge(START, "node")
graph.add_edge("node", END)
assert graph.compile().invoke({}, control=control) == {"value": "done"}
assert control.drain_requested
def test_drain_with_exit_durability_persists_resume_checkpoint() -> None:
class State(TypedDict, total=False):
first: str
second: str
control = RunControl()
def first_node(state: State) -> dict[str, str]:
control.request_drain("sigterm")
return {"first": "done"}
def second_node(state: State) -> dict[str, str]:
return {"second": "done"}
graph = StateGraph(State)
graph.add_node("first", first_node)
graph.add_node("second", second_node)
graph.add_edge(START, "first")
graph.add_edge("first", "second")
graph.add_edge("second", END)
compiled = graph.compile(checkpointer=MemorySaver())
config = {"configurable": {"thread_id": "drain-exit"}}
with pytest.raises(GraphDrained, match="sigterm"):
compiled.invoke({}, config, durability="exit", control=control)
assert compiled.invoke(None, config, durability="exit") == {
"first": "done",
"second": "done",
}
def test_drain_from_subgraph_can_resume_parent() -> None:
class State(TypedDict, total=False):
child_first: str
child_second: str
parent_second: str
control = RunControl()
def child_first(state: State) -> dict[str, str]:
control.request_drain("sigterm")
return {"child_first": "done"}
def child_second(state: State) -> dict[str, str]:
return {"child_second": "done"}
child_builder = StateGraph(State)
child_builder.add_node("child_first", child_first)
child_builder.add_node("child_second", child_second)
child_builder.add_edge(START, "child_first")
child_builder.add_edge("child_first", "child_second")
child_builder.add_edge("child_second", END)
child_graph = child_builder.compile(checkpointer=True)
def parent_second(state: State) -> dict[str, str]:
return {"parent_second": "done"}
parent_builder = StateGraph(State)
parent_builder.add_node("child", child_graph)
parent_builder.add_node("parent_second", parent_second)
parent_builder.add_edge(START, "child")
parent_builder.add_edge("child", "parent_second")
parent_builder.add_edge("parent_second", END)
compiled = parent_builder.compile(checkpointer=MemorySaver())
config = {"configurable": {"thread_id": "drain-subgraph"}}
with pytest.raises(GraphDrained, match="sigterm"):
compiled.invoke({}, config, control=control)
assert compiled.invoke(None, config) == {
"child_first": "done",
"child_second": "done",
"parent_second": "done",
}
@pytest.mark.anyio
async def test_drain_requested_in_terminal_step_finishes_normally_async() -> None:
class State(TypedDict, total=False):
value: str
control = RunControl()
async def node(state: State) -> dict[str, str]:
control.request_drain()
return {"value": "done"}
graph = StateGraph(State)
graph.add_node("node", node)
graph.add_edge(START, "node")
graph.add_edge("node", END)
assert await graph.compile().ainvoke({}, control=control) == {"value": "done"}
assert control.drain_requested
def test_runtime_propogated_to_subgraph() -> None:
@dataclass
class Context:
@@ -579,334 +392,6 @@ def test_context_coercion_pydantic_validation_errors() -> None:
)
def test_external_drain_concurrent_sync() -> None:
"""External thread calls request_drain() while graph is mid-execution."""
class State(TypedDict, total=False):
first: str
second: str
started = threading.Event()
def first_node(state: State) -> dict[str, str]:
started.set()
time.sleep(0.05)
return {"first": "done"}
def second_node(state: State) -> dict[str, str]:
return {"second": "should-not-run"}
graph = StateGraph(State)
graph.add_node("first", first_node)
graph.add_node("second", second_node)
graph.add_edge(START, "first")
graph.add_edge("first", "second")
graph.add_edge("second", END)
control = RunControl()
compiled = graph.compile()
exc_holder: list[BaseException | None] = [None]
def run_graph() -> None:
try:
compiled.invoke({}, control=control)
except GraphDrained as e:
exc_holder[0] = e
t = threading.Thread(target=run_graph)
t.start()
started.wait(timeout=5)
control.request_drain("sigterm")
t.join(timeout=10)
exc = exc_holder[0]
assert isinstance(exc, GraphDrained)
assert exc.reason == "sigterm"
@pytest.mark.anyio
async def test_external_drain_concurrent_async() -> None:
"""External task calls request_drain() while graph is mid-execution."""
class State(TypedDict, total=False):
first: str
second: str
started = asyncio.Event()
async def first_node(state: State) -> dict[str, str]:
started.set()
await asyncio.sleep(0.05)
return {"first": "done"}
async def second_node(state: State) -> dict[str, str]:
return {"second": "should-not-run"}
graph = StateGraph(State)
graph.add_node("first", first_node)
graph.add_node("second", second_node)
graph.add_edge(START, "first")
graph.add_edge("first", "second")
graph.add_edge("second", END)
control = RunControl()
compiled = graph.compile()
async def drain_after_start() -> None:
await started.wait()
control.request_drain("sigterm")
drain_task = asyncio.create_task(drain_after_start())
with pytest.raises(GraphDrained, match="sigterm"):
await compiled.ainvoke({}, control=control)
await drain_task
@pytest.mark.anyio
async def test_drain_then_cancel_after_graceful_timeout() -> None:
"""Simulate: drain requested -> node still running -> graceful timeout -> cancel.
This shows what happens when a long-running node doesn't finish within
the graceful period after drain is requested.
"""
class State(TypedDict, total=False):
first: str
second: str
node_started = asyncio.Event()
node_cancelled = asyncio.Event()
node_finished = asyncio.Event()
async def slow_node(state: State) -> dict[str, str]:
node_started.set()
try:
await asyncio.sleep(30) # very long operation
except asyncio.CancelledError:
node_cancelled.set()
raise
node_finished.set()
return {"first": "done"}
async def second_node(state: State) -> dict[str, str]:
return {"second": "should-not-run"}
graph = StateGraph(State)
graph.add_node("first", slow_node)
graph.add_node("second", second_node)
graph.add_edge(START, "first")
graph.add_edge("first", "second")
graph.add_edge("second", END)
control = RunControl()
compiled = graph.compile()
# Phase 1: start graph
graph_task = asyncio.create_task(compiled.ainvoke({}, control=control))
# Phase 2: wait for node to start, then request drain
await node_started.wait()
control.request_drain("sigterm")
# Phase 3: graceful timeout — node is still running, cancel after 1s
graceful_timeout = 1.0
await asyncio.sleep(graceful_timeout)
assert not node_finished.is_set(), "node should still be running"
assert not node_cancelled.is_set(), "node should not be cancelled yet"
# Phase 4: force cancel
graph_task.cancel()
with pytest.raises(asyncio.CancelledError):
await graph_task
# The node received CancelledError at the await point
assert node_cancelled.is_set(), "node should have received CancelledError"
assert not node_finished.is_set(), "node should NOT have finished normally"
@pytest.mark.anyio
async def test_cancel_ainvoke_with_async_node() -> None:
"""Cancel ainvoke running an async node: CancelledError is delivered
at the await point and the node stops immediately."""
class State(TypedDict, total=False):
first: str
second: str
timeline: list[str] = []
node_started = asyncio.Event()
async def slow_async_node(state: State) -> dict[str, str]:
timeline.append(f"async_node:start thread={threading.current_thread().name}")
node_started.set()
try:
await asyncio.sleep(30)
except asyncio.CancelledError:
timeline.append("async_node:cancelled")
raise
timeline.append("async_node:finished")
return {"first": "done"}
async def second_node(state: State) -> dict[str, str]:
timeline.append("second_node:run")
return {"second": "should-not-run"}
graph = StateGraph(State)
graph.add_node("first", slow_async_node)
graph.add_node("second", second_node)
graph.add_edge(START, "first")
graph.add_edge("first", "second")
graph.add_edge("second", END)
compiled = graph.compile()
graph_task = asyncio.create_task(compiled.ainvoke({}))
await node_started.wait()
timeline.append("test:cancel")
graph_task.cancel()
with pytest.raises(asyncio.CancelledError):
await graph_task
timeline.append("test:done")
# async node runs on the event loop thread (MainThread)
assert any("MainThread" in e for e in timeline if "async_node:start" in e)
# CancelledError was delivered at the await point — node stopped
assert "async_node:cancelled" in timeline
# Node did NOT run to completion
assert "async_node:finished" not in timeline
# Second node never ran
assert "second_node:run" not in timeline
@pytest.mark.anyio
async def test_cancel_ainvoke_with_sync_node() -> None:
"""Cancel ainvoke running a sync node.
Sync nodes in ainvoke run on a separate thread (via run_in_executor),
NOT on the event loop thread. Cancelling the asyncio task disconnects
from the thread future, but the thread keeps running as an orphan and
completes on its own.
Key difference from async nodes:
- async node: CancelledError stops the coroutine at an await point
- sync node: cancel only disconnects asyncio; the thread runs to completion
In shutdown case, we will ignore this because the instance will be destroyed soon.
"""
class State(TypedDict, total=False):
first: str
second: str
timeline: list[str] = []
node_started = threading.Event()
node_finished = threading.Event()
def slow_sync_node(state: State) -> dict[str, str]:
timeline.append(f"sync_node:start thread={threading.current_thread().name}")
node_started.set()
time.sleep(1)
timeline.append("sync_node:after_sleep")
node_finished.set()
return {"first": "done"}
def second_node(state: State) -> dict[str, str]:
timeline.append("second_node:run")
return {"second": "should-not-run"}
graph = StateGraph(State)
graph.add_node("first", slow_sync_node)
graph.add_node("second", second_node)
graph.add_edge(START, "first")
graph.add_edge("first", "second")
graph.add_edge("second", END)
control = RunControl()
compiled = graph.compile()
timeline.append(f"test:main thread={threading.current_thread().name}")
graph_task = asyncio.create_task(compiled.ainvoke({}, control=control))
loop = asyncio.get_event_loop()
await loop.run_in_executor(None, node_started.wait, 5)
timeline.append("test:cancel+drain")
graph_task.cancel()
control.request_drain("sigterm")
with pytest.raises(asyncio.CancelledError):
await graph_task
timeline.append("test:exc=CancelledError")
# Sync node runs on a background thread (asyncio_*), NOT MainThread
sync_start = next(e for e in timeline if "sync_node:start" in e)
assert "MainThread" not in sync_start, (
"sync node should run on a background thread, not the event loop thread"
)
# At this point, the asyncio task is done but the thread is orphaned.
# The sync node has NOT finished yet — cancel only disconnected asyncio.
assert not node_finished.is_set(), (
"sync node should still be running in its background thread"
)
# Wait for the orphaned thread to complete on its own.
await loop.run_in_executor(None, node_finished.wait, 5)
assert node_finished.is_set()
# After the orphaned thread finishes, the full timeline looks like:
# test:main thread=MainThread
# sync_node:start thread=asyncio_N <- background thread
# test:cancel+drain <- cancel + drain fired
# test:exc=CancelledError <- asyncio disconnected
# sync_node:after_sleep <- thread ran to completion anyway
assert "sync_node:after_sleep" in timeline
# Second node never ran
assert "second_node:run" not in timeline
# Verify timeline ordering: cancel happened before node finished
cancel_idx = timeline.index("test:cancel+drain")
sleep_idx = timeline.index("sync_node:after_sleep")
assert cancel_idx < sleep_idx, (
"cancel was issued while the sync node was still sleeping"
)
def test_drain_with_control_parameter_sync() -> None:
"""Control parameter is wired through invoke -> stream."""
class State(TypedDict, total=False):
value: str
ran = False
def node(state: State) -> dict[str, str]:
nonlocal ran
ran = True
return {"value": "done"}
graph = StateGraph(State)
graph.add_node("node", node)
graph.add_edge(START, "node")
graph.add_edge("node", END)
# Pre-drained control stops before executing the first pending task.
control = RunControl()
control.request_drain("pre-drained")
with pytest.raises(GraphDrained, match="pre-drained"):
graph.compile().invoke({}, control=control)
assert not ran
# --- ExecutionInfo unit tests ---
@@ -4,7 +4,7 @@ These transformers capture raw protocol events for their respective stream
modes and expose them as native projections on the run stream (run.custom,
run.updates, run.checkpoints, run.debug, run.tasks). Tests dispatch synthetic
protocol events through a StreamMux to isolate transformer logic; the final
group exercises real graphs through stream_events(version="v3").
group exercises real graphs through stream_v2.
"""
from __future__ import annotations
@@ -77,13 +77,8 @@ def _arm(mux: StreamMux, transformer: Any) -> None:
transformer._log._subscribed = True
def _unstamped(items):
"""Strip push stamps from a StreamChannel's internal buffer."""
return [item for _stamp, item in items]
def _drain(transformer: Any) -> list[Any]:
return _unstamped(transformer._log._items)
return list(transformer._log._items)
# ---------------------------------------------------------------------------
@@ -145,7 +140,7 @@ def test_custom_does_not_suppress_from_main_log() -> None:
mux.push(_custom_event([], "data"))
methods = [evt["method"] for evt in _unstamped(mux._events._items)]
methods = [evt["method"] for evt in mux._events._items]
assert "custom" in methods
@@ -224,7 +219,7 @@ def test_checkpoints_does_not_suppress_from_main_log() -> None:
mux.push(_checkpoints_event([], {"values": {}}))
methods = [evt["method"] for evt in _unstamped(mux._events._items)]
methods = [evt["method"] for evt in mux._events._items]
assert "checkpoints" in methods
@@ -289,7 +284,7 @@ def test_debug_does_not_suppress_from_main_log() -> None:
mux.push(_debug_event([], {"step": 0}))
methods = [evt["method"] for evt in _unstamped(mux._events._items)]
methods = [evt["method"] for evt in mux._events._items]
assert "debug" in methods
@@ -365,7 +360,7 @@ def test_tasks_does_not_suppress_from_main_log() -> None:
mux.push(_tasks_event([], {"id": "t1"}))
methods = [evt["method"] for evt in _unstamped(mux._events._items)]
methods = [evt["method"] for evt in mux._events._items]
assert "tasks" in methods
@@ -434,7 +429,7 @@ def test_updates_does_not_suppress_from_main_log() -> None:
mux.push(_updates_event([], {"n": {}}))
methods = [evt["method"] for evt in _unstamped(mux._events._items)]
methods = [evt["method"] for evt in mux._events._items]
assert "updates" in methods
@@ -474,11 +469,11 @@ def test_unrelated_events_ignored_by_all() -> None:
)
for t in transformers:
assert _unstamped(t._log._items) == []
assert list(t._log._items) == []
# ---------------------------------------------------------------------------
# End-to-end: real graphs through stream_events(version="v3")
# End-to-end: real graphs through stream_v2
# ---------------------------------------------------------------------------
@@ -503,11 +498,11 @@ def _make_simple_graph() -> Any:
return builder.compile()
def test_stream_events_v3_custom_projection_opt_in() -> None:
def test_stream_v2_custom_projection_opt_in() -> None:
"""run.custom surfaces get_stream_writer() payloads when opted in."""
graph = _make_simple_graph()
run = graph.stream_events(
{"value": "hello", "items": []}, version="v3", transformers=[CustomTransformer]
run = graph.stream_v2(
{"value": "hello", "items": []}, transformers=[CustomTransformer]
)
custom_events = list(run.custom)
@@ -515,11 +510,11 @@ def test_stream_events_v3_custom_projection_opt_in() -> None:
assert any(e.get("status") == "working" for e in custom_events)
def test_stream_events_v3_custom_and_values_coexist() -> None:
def test_stream_v2_custom_and_values_coexist() -> None:
"""Both run.custom and run.values work in the same run."""
graph = _make_simple_graph()
run = graph.stream_events(
{"value": "hello", "items": []}, version="v3", transformers=[CustomTransformer]
run = graph.stream_v2(
{"value": "hello", "items": []}, transformers=[CustomTransformer]
)
custom_events = list(run.custom)
@@ -528,12 +523,10 @@ def test_stream_events_v3_custom_and_values_coexist() -> None:
assert len(custom_events) >= 1
def test_stream_events_v3_tasks_projection_opt_in() -> None:
def test_stream_v2_tasks_projection_opt_in() -> None:
"""run.tasks surfaces raw task events when opted in via transformers=."""
graph = _make_simple_graph()
run = graph.stream_events(
{"value": "x", "items": []}, transformers=[TasksTransformer], version="v3"
)
run = graph.stream_v2({"value": "x", "items": []}, transformers=[TasksTransformer])
tasks_events = list(run.tasks)
assert len(tasks_events) >= 1
@@ -541,12 +534,10 @@ def test_stream_events_v3_tasks_projection_opt_in() -> None:
assert "my_node" in names
def test_stream_events_v3_debug_projection_opt_in() -> None:
def test_stream_v2_debug_projection_opt_in() -> None:
"""run.debug surfaces debug events when opted in via transformers=."""
graph = _make_simple_graph()
run = graph.stream_events(
{"value": "x", "items": []}, transformers=[DebugTransformer], version="v3"
)
run = graph.stream_v2({"value": "x", "items": []}, transformers=[DebugTransformer])
debug_events = list(run.debug)
assert len(debug_events) >= 1
@@ -554,11 +545,11 @@ def test_stream_events_v3_debug_projection_opt_in() -> None:
assert types & {"checkpoint", "task", "task_result"}
def test_stream_events_v3_updates_projection_opt_in() -> None:
def test_stream_v2_updates_projection_opt_in() -> None:
"""run.updates surfaces node output dicts when opted in via transformers=."""
graph = _make_simple_graph()
run = graph.stream_events(
{"value": "x", "items": []}, version="v3", transformers=[UpdatesTransformer]
run = graph.stream_v2(
{"value": "x", "items": []}, transformers=[UpdatesTransformer]
)
updates = list(run.updates)
@@ -567,12 +558,11 @@ def test_stream_events_v3_updates_projection_opt_in() -> None:
assert "my_node" in node_names
def test_stream_events_v3_all_transformers_interleaved() -> None:
def test_stream_v2_all_transformers_interleaved() -> None:
"""All five transformers registered together, consumed via interleave."""
graph = _make_simple_graph()
run = graph.stream_events(
run = graph.stream_v2(
{"value": "x", "items": []},
version="v3",
transformers=[
CustomTransformer,
UpdatesTransformer,
@@ -604,7 +594,7 @@ def test_stream_events_v3_all_transformers_interleaved() -> None:
assert run.output["value"] == "x!"
def test_stream_events_v3_all_transformers_with_checkpointer() -> None:
def test_stream_v2_all_transformers_with_checkpointer() -> None:
"""All transformers with a checkpointer — run.checkpoints populated."""
from langgraph.checkpoint.memory import InMemorySaver
@@ -614,9 +604,8 @@ def test_stream_events_v3_all_transformers_with_checkpointer() -> None:
builder.add_edge("my_node", END)
graph = builder.compile(checkpointer=InMemorySaver())
run = graph.stream_events(
run = graph.stream_v2(
{"value": "x", "items": []},
version="v3",
config={"configurable": {"thread_id": "test-all"}},
transformers=[
CustomTransformer,
@@ -643,7 +632,7 @@ def test_stream_events_v3_all_transformers_with_checkpointer() -> None:
assert len(collected["custom"]) >= 1
def test_stream_events_v3_checkpoints_projection_opt_in() -> None:
def test_stream_v2_checkpoints_projection_opt_in() -> None:
"""run.checkpoints surfaces checkpoint data when opted in with a checkpointer."""
from langgraph.checkpoint.memory import InMemorySaver
@@ -653,9 +642,8 @@ def test_stream_events_v3_checkpoints_projection_opt_in() -> None:
builder.add_edge("my_node", END)
graph = builder.compile(checkpointer=InMemorySaver())
run = graph.stream_events(
run = graph.stream_v2(
{"value": "x", "items": []},
version="v3",
config={"configurable": {"thread_id": "test-ckpt-standalone"}},
transformers=[CheckpointsTransformer],
)
@@ -686,7 +674,7 @@ def test_tasks_and_lifecycle_coregistration() -> None:
assert _drain(tasks) == [task_data]
methods = [evt["method"] for evt in _unstamped(mux._events._items)]
methods = [evt["method"] for evt in mux._events._items]
assert "tasks" not in methods
@@ -695,9 +683,8 @@ def test_tasks_and_lifecycle_coregistration_e2e() -> None:
is present and suppressing them from the main log.
"""
graph = _make_simple_graph()
run = graph.stream_events(
run = graph.stream_v2(
{"value": "x", "items": []},
version="v3",
transformers=[TasksTransformer],
)
@@ -6,7 +6,7 @@ on the `lifecycle` channel for both in-process iteration via
events. Most tests dispatch synthetic protocol events through a
`StreamMux` to keep the inference logic isolated; the end-of-file
group exercises the path through real graphs (multi-depth
discovery, nested `stream_events(version="v3")` calls with non-empty `parent_ns`).
discovery, nested `stream_v2` calls with non-empty `parent_ns`).
"""
from __future__ import annotations
@@ -92,16 +92,11 @@ def _arm(mux: StreamMux) -> None:
transformer._channel._subscribed = True
def _unstamped(items):
"""Strip push stamps from a StreamChannel's internal buffer."""
return [item for _stamp, item in items]
def _drain_lifecycle(mux: StreamMux) -> list[LifecyclePayload]:
"""Snapshot the lifecycle channel's buffer."""
transformer = mux.transformer_by_key("lifecycle")
assert isinstance(transformer, LifecycleTransformer)
return _unstamped(transformer._channel._items)
return list(transformer._channel._items)
def _build_lifecycle_mux(*, scope: tuple[str, ...] = ()) -> StreamMux:
@@ -306,7 +301,7 @@ def test_protocol_event_method_is_native() -> None:
mux = _build_lifecycle_mux()
mux.push(_tasks_start(["agent:abc"], task_id="t1", name="tool"))
methods = {evt["method"] for evt in _unstamped(mux._events._items)}
methods = {evt["method"] for evt in mux._events._items}
assert "lifecycle" in methods
assert "custom:lifecycle" not in methods
@@ -317,14 +312,14 @@ def test_tasks_events_suppressed_from_main_log() -> None:
mux.push(_tasks_start(["agent:abc"], task_id="t1", name="tool"))
mux.push(_tasks_result([], task_id="abc", name="agent"))
methods = [evt["method"] for evt in _unstamped(mux._events._items)]
methods = [evt["method"] for evt in mux._events._items]
assert "tasks" not in methods
# Lifecycle events did make it through, though.
assert "lifecycle" in methods
# ---------------------------------------------------------------------------
# End-to-end: real graphs through stream_events(version="v3")
# End-to-end: real graphs through stream_v2
# ---------------------------------------------------------------------------
@@ -358,10 +353,10 @@ def _make_two_level_nested() -> Any:
return outer_b.compile()
def test_stream_events_v3_real_graph_emits_lifecycle_at_each_depth() -> None:
def test_stream_v2_real_graph_emits_lifecycle_at_each_depth() -> None:
"""Outer graph with two nested subgraphs surfaces lifecycle for both."""
graph = _make_two_level_nested()
run = graph.stream_events({"value": "x", "items": []}, version="v3")
run = graph.stream_v2({"value": "x", "items": []})
# Iterating the projection drives the pump and drains synthesized
# lifecycle events at the same time.
@@ -384,17 +379,17 @@ def test_stream_events_v3_real_graph_emits_lifecycle_at_each_depth() -> None:
)
def test_stream_events_v3_with_nested_parent_ns_scopes_lifecycle() -> None:
"""When `stream_events(version="v3")` is called with a non-empty checkpoint_ns in config,
def test_stream_v2_with_nested_parent_ns_scopes_lifecycle() -> None:
"""When `stream_v2` is called with a non-empty checkpoint_ns in config,
`_resolve_parent_ns` returns that namespace and the registered
`LifecycleTransformer` is constructed with `scope=parent_ns`. This
exercises the path that exists today purely for nested-stream_events(version="v3")
exercises the path that exists today purely for nested-stream_v2
callers; the test simulates such a caller by injecting a
checkpoint_ns into the config.
"""
graph = _make_two_level_nested()
config = {CONF: {CONFIG_KEY_CHECKPOINT_NS: "outer:abc"}}
run = graph.stream_events({"value": "x", "items": []}, config=config, version="v3")
run = graph.stream_v2({"value": "x", "items": []}, config=config)
payloads = list(run.lifecycle)
# Every emitted lifecycle namespace must extend the caller's scope —
@@ -1,5 +1,5 @@
"""Tests for MessagesTransformer: protocol event routing, whole-message fallback,
legacy v1 chunk filtering, and end-to-end via stream_events(version="v3") / astream_events(version="v3")."""
legacy v1 chunk filtering, and end-to-end via stream_v2 / astream_v2."""
from __future__ import annotations
@@ -26,11 +26,6 @@ from langgraph.stream.transformers import MessagesTransformer, ValuesTransformer
TS = int(time.time() * 1000)
def _unstamped(items):
"""Strip push stamps from a StreamChannel's internal buffer."""
return [item for _stamp, item in items]
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
@@ -152,7 +147,7 @@ def _lifecycle(
def _simple_graph():
def call_model(state: MessagesState) -> dict[str, Any]:
model = GenericFakeChatModel(messages=iter(["hello world"]))
stream = model.stream_events(state["messages"], version="v3")
stream = model.stream_v2(state["messages"])
return {"messages": stream.output}
return (
@@ -179,7 +174,7 @@ class TestProtocolEventRouting:
)
)
log.close()
(stream,) = _unstamped(log._items)
(stream,) = list(log._items)
assert isinstance(stream, ChatModelStream)
assert stream.message_id == "run-1"
@@ -188,7 +183,7 @@ class TestProtocolEventRouting:
for evt in _lifecycle(text="hello world"):
t.process(_proto_event(evt, run_id="run-1"))
log.close()
(stream,) = _unstamped(log._items)
(stream,) = list(log._items)
assert stream.done
assert stream.output.text == "hello world"
@@ -211,7 +206,7 @@ class TestProtocolEventRouting:
)
)
log.close()
assert _unstamped(log._items) == []
assert list(log._items) == []
def test_concurrent_streams_routed_by_run_id(self) -> None:
t, log = _make_sync_transformer()
@@ -221,7 +216,7 @@ class TestProtocolEventRouting:
t.process(_proto_event(a, run_id="run-a"))
t.process(_proto_event(b, run_id="run-b"))
log.close()
streams = _unstamped(log._items)
streams = list(log._items)
assert len(streams) == 2
by_id = {s.message_id: s for s in streams}
assert by_id["run-a"].output.text == "aaaa"
@@ -232,7 +227,7 @@ class TestProtocolEventRouting:
for evt in _lifecycle(text="abcdef"):
t.process(_proto_event(evt))
log.close()
(stream,) = _unstamped(log._items)
(stream,) = list(log._items)
assert "".join(stream._text_proj._deltas) == "abcdef"
def test_stream_pushed_on_message_start_not_finish(self) -> None:
@@ -255,7 +250,7 @@ class TestProtocolEventRouting:
node="my_llm",
)
)
(stream,) = _unstamped(log._items)
(stream,) = list(log._items)
assert stream.node == "my_llm"
@@ -269,7 +264,7 @@ class TestWholeMessageFallback:
t, log = _make_sync_transformer()
t.process(_whole_msg("the full answer"))
log.close()
(stream,) = _unstamped(log._items)
(stream,) = list(log._items)
assert stream.done
assert stream.output.text == "the full answer"
@@ -277,7 +272,7 @@ class TestWholeMessageFallback:
t, log = _make_sync_transformer()
t.process(_whole_msg("full"))
log.close()
(stream,) = _unstamped(log._items)
(stream,) = list(log._items)
assert [e["event"] for e in stream._events] == [
"message-start",
"content-block-start",
@@ -323,16 +318,16 @@ class TestFiltering:
}
)
log.close()
assert _unstamped(log._items) == []
assert list(log._items) == []
def test_legacy_v1_chunks_ignored(self) -> None:
# v1 AIMessageChunk tuples (from on_llm_new_token) are not streamed
# into this projection; callers must migrate to stream_events(version="v3").
# into this projection; callers must migrate to stream_v2.
t, log = _make_sync_transformer()
t.process(_v1_chunk("hello"))
t.process(_v1_chunk(" world", finish=True))
log.close()
assert _unstamped(log._items) == []
assert list(log._items) == []
# ---------------------------------------------------------------------------
@@ -348,7 +343,7 @@ class TestLifecycle:
{"event": "message-start", "message_id": "run-1"}, run_id="run-1"
)
)
streams = _unstamped(log._items)
streams = list(log._items)
err = RuntimeError("graph died")
t.fail(err)
assert t._by_run == {}
@@ -376,14 +371,14 @@ class TestAsyncMode:
t, log = _make_async_transformer()
for evt in _lifecycle(text="async stream"):
t.process(_proto_event(evt))
assert isinstance(_unstamped(log._items)[0], AsyncChatModelStream)
assert isinstance(list(log._items)[0], AsyncChatModelStream)
@pytest.mark.anyio
async def test_text_projection_yields_deltas(self) -> None:
t, log = _make_async_transformer()
for evt in _lifecycle(text="hello world"):
t.process(_proto_event(evt))
(stream,) = _unstamped(log._items)
(stream,) = list(log._items)
assert isinstance(stream, AsyncChatModelStream)
assert "".join([d async for d in stream.text]) == "hello world"
@@ -392,7 +387,7 @@ class TestAsyncMode:
t, log = _make_async_transformer()
for evt in _lifecycle(text="async"):
t.process(_proto_event(evt))
(stream,) = _unstamped(log._items)
(stream,) = list(log._items)
assert (await stream.output).text == "async"
@@ -424,7 +419,7 @@ class TestWireRequestMore:
for evt in _lifecycle():
messages_t.process(_proto_event(evt))
(stream,) = _unstamped(log._items)
(stream,) = list(log._items)
assert stream._request_more is messages_t._pump_fn
@@ -450,14 +445,14 @@ class TestViaMux:
for evt in _lifecycle(text="mux stream"):
mux.push(_proto_event(evt))
mux.close()
(stream,) = _unstamped(log._items)
(stream,) = list(log._items)
assert stream.output.text == "mux stream"
def test_whole_message_via_mux(self) -> None:
t, mux, log = self._make_mux()
mux.push(_whole_msg("result"))
mux.close()
(stream,) = _unstamped(log._items)
(stream,) = list(log._items)
assert stream.output.text == "result"
@pytest.mark.anyio
@@ -471,24 +466,24 @@ class TestViaMux:
for evt in _lifecycle(text="async mux"):
await mux.apush(_proto_event(evt))
(stream,) = _unstamped(log._items)
(stream,) = list(log._items)
assert (await stream.output).text == "async mux"
await mux.aclose()
# ---------------------------------------------------------------------------
# End-to-end: graph → stream_events(version="v3") → run.messages (node calls stream_events)
# End-to-end: graph → stream_v2 → run.messages (node calls stream_v2)
# ---------------------------------------------------------------------------
class TestEndToEnd:
"""stream_events(version="v3") path: node calls model.stream_events() explicitly."""
"""stream_v2 path: node calls model.stream_v2() explicitly."""
def test_node_calling_stream_v2_populates_messages(self) -> None:
model = GenericFakeChatModel(messages=iter(["hello world"]))
def call_model(state: MessagesState) -> dict[str, Any]:
stream = model.stream_events(state["messages"], version="v3")
stream = model.stream_v2(state["messages"])
return {"messages": stream.output}
graph = (
@@ -499,7 +494,7 @@ class TestEndToEnd:
.compile()
)
run = graph.stream_events({"messages": "hi"}, version="v3")
run = graph.stream_v2({"messages": "hi"})
(stream,) = list(run.messages)
assert isinstance(stream, ChatModelStream)
assert stream.output.text == "hello world"
@@ -509,7 +504,7 @@ class TestEndToEnd:
model = GenericFakeChatModel(messages=iter(["streamed answer"]))
def call_model(state: MessagesState) -> dict[str, Any]:
stream = model.stream_events(state["messages"], version="v3")
stream = model.stream_v2(state["messages"])
return {"messages": stream.output}
graph = (
@@ -520,7 +515,7 @@ class TestEndToEnd:
.compile()
)
run = graph.stream_events({"messages": "go"}, version="v3")
run = graph.stream_v2({"messages": "go"})
(stream,) = list(run.messages)
assert "".join(stream.text) == "streamed answer"
@@ -538,7 +533,7 @@ class TestEndToEnd:
.compile()
)
run = graph.stream_events({"messages": "hi"}, version="v3")
run = graph.stream_v2({"messages": "hi"})
(stream,) = list(run.messages)
assert stream.output.text == "hardcoded"
@@ -547,7 +542,7 @@ class TestEndToEnd:
model = GenericFakeChatModel(messages=iter(["async answer"]))
async def call_model(state: MessagesState) -> dict[str, Any]:
stream = await model.astream_events(state["messages"], version="v3")
stream = await model.astream_v2(state["messages"])
return {"messages": await stream}
graph = (
@@ -558,7 +553,7 @@ class TestEndToEnd:
.compile()
)
run = await graph.astream_events({"messages": "hi"}, version="v3")
run = await graph.astream_v2({"messages": "hi"})
streams = [s async for s in run.messages]
assert len(streams) == 1
assert isinstance(streams[0], AsyncChatModelStream)
@@ -572,7 +567,7 @@ class TestEndToEnd:
model = GenericFakeChatModel(messages=iter(["hello world"]))
async def call_model(state: MessagesState) -> dict[str, Any]:
stream = await model.astream_events(state["messages"], version="v3")
stream = await model.astream_v2(state["messages"])
return {"messages": await stream}
graph = (
@@ -583,7 +578,7 @@ class TestEndToEnd:
.compile()
)
run = await graph.astream_events({"messages": "hi"}, version="v3")
run = await graph.astream_v2({"messages": "hi"})
async def consume() -> list[str]:
collected: list[str] = []
@@ -596,12 +591,12 @@ class TestEndToEnd:
# ---------------------------------------------------------------------------
# End-to-end: graph → stream_events(version="v3") → run.messages (node calls invoke)
# End-to-end: graph → stream_v2 → run.messages (node calls invoke)
# ---------------------------------------------------------------------------
class TestEndToEndV2Invoke:
"""Auto-routing path: stream_events(version="v3") injects CONFIG_KEY_STREAM_MESSAGES_V2,
"""Auto-routing path: stream_v2 injects CONFIG_KEY_STREAM_MESSAGES_V2,
causing BaseChatModel to drive the v2 protocol event generator even for
model.invoke()."""
@@ -620,7 +615,7 @@ class TestEndToEndV2Invoke:
def test_invoke_populates_messages(self) -> None:
run = self._graph(
GenericFakeChatModel(messages=iter(["hello world"]))
).stream_events({"messages": "hi"}, version="v3")
).stream_v2({"messages": "hi"})
(stream,) = list(run.messages)
assert isinstance(stream, ChatModelStream)
assert stream.output.text == "hello world"
@@ -629,7 +624,7 @@ class TestEndToEndV2Invoke:
"""Iterating the stream yields the full v2 lifecycle, not v1 chunks."""
run = self._graph(
GenericFakeChatModel(messages=iter(["streamed answer"]))
).stream_events({"messages": "go"}, version="v3")
).stream_v2({"messages": "go"})
(stream,) = list(run.messages)
events = list(stream)
@@ -650,7 +645,7 @@ class TestEndToEndV2Invoke:
def test_invoke_text_deltas_iterate(self) -> None:
run = self._graph(
GenericFakeChatModel(messages=iter(["delta streaming works"]))
).stream_events({"messages": "hi"}, version="v3")
).stream_v2({"messages": "hi"})
(stream,) = list(run.messages)
assert "".join(stream.text) == "delta streaming works"
@@ -674,7 +669,7 @@ class TestEndToEndV2Invoke:
.compile()
)
streams = list(graph.stream_events({"messages": "hi"}, version="v3").messages)
streams = list(graph.stream_v2({"messages": "hi"}).messages)
assert len(streams) == 2
assert {s.output.text for s in streams} == {"alpha", "beta"}
@@ -698,7 +693,7 @@ class TestEndToEndV2Invoke:
.compile()
)
run = graph.stream_events({"messages": "hi"}, version="v3")
run = graph.stream_v2({"messages": "hi"})
streams = list(run.messages)
assert len(streams) == 2
assert streams[0].node == "streaming_node"
@@ -722,7 +717,7 @@ class TestEndToEndV2Invoke:
.compile()
)
run = await graph.astream_events({"messages": "hi"}, version="v3")
run = await graph.astream_v2({"messages": "hi"})
streams = [s async for s in run.messages]
assert len(streams) == 1
assert isinstance(streams[0], AsyncChatModelStream)
@@ -737,7 +732,7 @@ class TestEndToEndV2Invoke:
class TestDirectMessagesModeStaysV1:
def test_direct_graph_stream_messages_yields_ai_message_chunks(self) -> None:
"""graph.stream(stream_mode="messages") must not leak v2 event dicts —
the v2 flag is only injected by stream_events(version="v3") / astream_events(version="v3")."""
the v2 flag is only injected by stream_v2 / astream_v2."""
model = GenericFakeChatModel(messages=iter(["legacy path"]))
def call_model(state: MessagesState) -> dict[str, Any]:
@@ -760,10 +755,8 @@ class TestDirectMessagesModeStaysV1:
== "legacy path"
)
def test_nested_graph_stream_messages_stays_v1_under_outer_stream_events_v3(
self,
) -> None:
"""An outer `stream_events(version="v3")` run must not flip an inner direct
def test_nested_graph_stream_messages_stays_v1_under_outer_stream_v2(self) -> None:
"""An outer `stream_v2()` run must not flip an inner direct
`stream_mode="messages"` call onto the v2 event protocol."""
model = GenericFakeChatModel(messages=iter(["nested legacy path"]))
@@ -814,7 +807,7 @@ class TestDirectMessagesModeStaysV1:
.compile()
)
result = outer.stream_events({}, version="v3").output
result = outer.stream_v2({}).output
assert result is not None
assert result["saw_only_chunks"] is True
@@ -4,7 +4,7 @@ Subscribes to `tasks` events and produces in-process `SubgraphRunStream`
handles backed by mini-muxes (built via `StreamMux._make_child`). The
synthetic-event tests isolate the inference / mini-mux wiring; the
real-graph tests exercise the end-to-end navigation path through
`stream_events(version="v3")`.
`stream_v2`.
"""
from __future__ import annotations
@@ -93,7 +93,7 @@ def _tasks_result(
def _native_factories() -> list[Any]:
"""Mirror the factory list `Pregel.stream_events(version="v3")` registers."""
"""Mirror the factory list `Pregel.stream_v2` registers."""
return [
ValuesTransformer,
MessagesTransformer,
@@ -159,13 +159,8 @@ def _subgraph_transformer(mux: StreamMux) -> SubgraphTransformer:
return transformer
def _unstamped(items):
"""Strip push stamps from a StreamChannel's internal buffer."""
return [item for _stamp, item in items]
def _drain_subgraphs(mux: StreamMux) -> list[SubgraphRunStream]:
return _unstamped(_subgraph_transformer(mux)._log._items)
return list(_subgraph_transformer(mux)._log._items)
def _child_mux(handle: SubgraphRunStream | AsyncSubgraphRunStream) -> StreamMux:
@@ -174,13 +169,13 @@ def _child_mux(handle: SubgraphRunStream | AsyncSubgraphRunStream) -> StreamMux:
def _event_items(mux: StreamMux) -> list[ProtocolEvent]:
return _unstamped(mux._events._items)
return list(mux._events._items)
def _lifecycle_payloads(mux: StreamMux) -> list[dict[str, Any]]:
lifecycle_t = mux.transformer_by_key("lifecycle")
assert isinstance(lifecycle_t, LifecycleTransformer)
return _unstamped(lifecycle_t._channel._items)
return list(lifecycle_t._channel._items)
# ---------------------------------------------------------------------------
@@ -252,7 +247,7 @@ def test_grandchild_discovered_via_child_mini_mux() -> None:
[child_handle] = _drain_subgraphs(mux)
assert child_handle.path == ("agent:abc",)
# The grandchild appears on the CHILD'S subgraphs projection.
grandchildren = _unstamped(child_handle.subgraphs._items)
grandchildren = list(child_handle.subgraphs._items)
assert len(grandchildren) == 1
assert grandchildren[0].path == ("agent:abc", "tool:def")
@@ -761,10 +756,10 @@ def _make_failing_nested() -> Any:
return outer_b.compile()
def test_stream_events_v3_real_graph_yields_subgraph_handles() -> None:
def test_stream_v2_real_graph_yields_subgraph_handles() -> None:
"""Iterating `run.subgraphs` yields handles for direct-child subgraphs."""
graph = _make_two_level_nested()
run = graph.stream_events({"value": "x", "items": []}, version="v3")
run = graph.stream_v2({"value": "x", "items": []})
handle_paths: list[tuple[str, ...]] = []
final_status: dict[tuple[str, ...], str] = {}
@@ -780,10 +775,10 @@ def test_stream_events_v3_real_graph_yields_subgraph_handles() -> None:
assert final_status[handle_paths[0]] == "completed"
def test_stream_events_v3_grandchild_visible_on_child_handle() -> None:
def test_stream_v2_grandchild_visible_on_child_handle() -> None:
"""Drilling into `handle.subgraphs` surfaces nested grandchildren."""
graph = _make_two_level_nested()
run = graph.stream_events({"value": "x", "items": []}, version="v3")
run = graph.stream_v2({"value": "x", "items": []})
grandchild_paths: list[tuple[str, ...]] = []
middle_path: tuple[str, ...] | None = None
@@ -810,7 +805,7 @@ def test_subgraph_output_stops_at_own_terminal_without_draining_siblings() -> No
inside the loop body misses its events.
"""
graph = _make_two_sibling_subgraphs()
run = graph.stream_events({"value": "x", "items": []}, version="v3")
run = graph.stream_v2({"value": "x", "items": []})
paths: list[tuple[str, ...]] = []
second_values: list[dict[str, Any]] = []
@@ -829,7 +824,7 @@ def test_subgraph_output_stops_at_own_terminal_without_draining_siblings() -> No
def test_aborted_subgraph_handle_does_not_fail_parent_forwarding() -> None:
graph = _make_two_sibling_subgraphs()
run = graph.stream_events({"value": "x", "items": []}, version="v3")
run = graph.stream_v2({"value": "x", "items": []})
seen: list[str | None] = []
for handle in run.subgraphs:
@@ -847,7 +842,7 @@ def test_aborted_subgraph_handle_does_not_fail_parent_forwarding() -> None:
def test_failed_subgraph_output_raises_terminal_error() -> None:
graph = _make_failing_nested()
run = graph.stream_events({"value": "x", "items": []}, version="v3")
run = graph.stream_v2({"value": "x", "items": []})
handle = next(iter(run.subgraphs))
with pytest.raises(RuntimeError, match="child boom"):
@@ -19,11 +19,9 @@ from typing_extensions import TypedDict, assert_type
from langgraph._internal._constants import INTERRUPT
from langgraph.constants import END, START
from langgraph.errors import GraphDrained
from langgraph.func import entrypoint
from langgraph.graph import StateGraph
from langgraph.graph.message import MessagesState
from langgraph.runtime import RunControl
from langgraph.types import (
CheckpointPayload,
CheckpointStreamPart,
@@ -231,32 +229,6 @@ class TestV2Stream:
for c in chunks:
_assert_stream_part_shape(c)
def test_stream_events_v3_accepts_control_for_drain(self) -> None:
class DrainState(TypedDict, total=False):
value: str
skipped: str
control = RunControl()
def first_node(state: DrainState) -> dict[str, str]:
control.request_drain("sigterm")
return {"value": "done"}
def second_node(state: DrainState) -> dict[str, str]:
return {"skipped": "nope"}
builder = StateGraph(DrainState)
builder.add_node("first", first_node)
builder.add_node("second", second_node)
builder.add_edge(START, "first")
builder.add_edge("first", "second")
builder.add_edge("second", END)
graph = builder.compile()
run = graph.stream_events({}, control=control, version="v3")
with pytest.raises(GraphDrained, match="sigterm"):
list(run.values)
def test_subgraphs_ns(self) -> None:
outer = _make_subgraph()
chunks = list(
@@ -1124,7 +1096,7 @@ class TestV2ValidationErrors:
_INVALID_INPUT: dict[str, Any] = {"value": [1, 2, 3], "items": []}
def test_stream_events_v3_pydantic_validation_error(self) -> None:
def test_stream_v2_pydantic_validation_error(self) -> None:
"""Invalid input to stream with v2 + pydantic state raises ValidationError."""
graph = _make_pydantic_graph()
with pytest.raises(ValidationError):
@@ -1,9 +1,9 @@
"""End-to-end tests exercising all stream_events(version="v3") projections together.
"""End-to-end tests exercising all stream_v2 projections together.
Each test builds a realistic graph (subgraphs, LLM calls, custom writers,
interrupts) and verifies that every projection values, messages, lifecycle,
subgraphs, raw events, output, interleave produces correct, consistent
results through a single stream_events(version="v3") / astream_events(version="v3") run.
results through a single stream_v2 / astream_v2 run.
"""
from __future__ import annotations
@@ -218,9 +218,9 @@ class _CounterTransformer(StreamTransformer):
class TestStreamV2E2ESync:
def test_all_projections_nested_graph(self) -> None:
"""Run a nested graph through stream_events(version="v3") and verify values + lifecycle."""
"""Run a nested graph through stream_v2 and verify values + lifecycle."""
graph = _make_nested_graph()
run = graph.stream_events({"value": "x", "items": []}, version="v3")
run = graph.stream_v2({"value": "x", "items": []})
values_snapshots: list[dict[str, Any]] = []
lifecycle_events: list[dict[str, Any]] = []
@@ -246,7 +246,7 @@ class TestStreamV2E2ESync:
def test_subgraph_handles_with_drill_down(self) -> None:
"""Subgraph handles yield and support values drill-down."""
graph = _make_nested_graph()
run = graph.stream_events({"value": "x", "items": []}, version="v3")
run = graph.stream_v2({"value": "x", "items": []})
handles = []
for handle in run.subgraphs:
@@ -270,7 +270,7 @@ class TestStreamV2E2ESync:
def test_raw_events_have_monotonic_seq(self) -> None:
"""Raw protocol events have monotonically increasing seq numbers."""
graph = _make_nested_graph()
run = graph.stream_events({"value": "x", "items": []}, version="v3")
run = graph.stream_v2({"value": "x", "items": []})
events = list(run)
assert len(events) > 0
@@ -285,15 +285,11 @@ class TestStreamV2E2ESync:
def test_output_matches_final_values_snapshot(self) -> None:
"""output property returns the same state as the last values snapshot."""
run1 = _make_nested_graph().stream_events(
{"value": "x", "items": []}, version="v3"
)
run1 = _make_nested_graph().stream_v2({"value": "x", "items": []})
snapshots = list(run1.values)
final_via_values = snapshots[-1]
run2 = _make_nested_graph().stream_events(
{"value": "x", "items": []}, version="v3"
)
run2 = _make_nested_graph().stream_v2({"value": "x", "items": []})
final_via_output = run2.output
assert final_via_values == final_via_output
@@ -301,7 +297,7 @@ class TestStreamV2E2ESync:
def test_context_manager_and_abort(self) -> None:
"""Context manager calls abort, marking the stream exhausted."""
graph = _make_nested_graph()
with graph.stream_events({"value": "x", "items": []}, version="v3") as run:
with graph.stream_v2({"value": "x", "items": []}) as run:
first_val = next(iter(run.values))
assert isinstance(first_val, dict)
assert run._exhausted is True
@@ -309,7 +305,7 @@ class TestStreamV2E2ESync:
def test_extensions_has_all_native_keys(self) -> None:
"""Extensions dict exposes all native projection keys."""
graph = _make_nested_graph()
run = graph.stream_events({"value": "x", "items": []}, version="v3")
run = graph.stream_v2({"value": "x", "items": []})
_ = run.output
assert "values" in run.extensions
@@ -331,7 +327,7 @@ class TestStreamV2E2EMessages:
def test_messages_projection_from_invoke(self) -> None:
"""Messages projection captures LLM calls via model.invoke() auto-routing."""
graph = _make_messages_graph()
run = graph.stream_events({"messages": "hi"}, version="v3")
run = graph.stream_v2({"messages": "hi"})
streams = list(run.messages)
assert len(streams) >= 1
@@ -354,7 +350,7 @@ class TestStreamV2E2EMessages:
.compile()
)
run = graph.stream_events({"messages": "go"}, version="v3")
run = graph.stream_v2({"messages": "go"})
(stream,) = list(run.messages)
assert "".join(stream.text) == "streamed answer"
@@ -372,7 +368,7 @@ class TestStreamV2E2EMessages:
.compile()
)
run = graph.stream_events({"messages": "hi"}, version="v3")
run = graph.stream_v2({"messages": "hi"})
(stream,) = list(run.messages)
assert stream.output.text == "hardcoded"
assert stream.message_id == "msg-1"
@@ -380,7 +376,7 @@ class TestStreamV2E2EMessages:
def test_root_messages_only_shows_root_scope(self) -> None:
"""Root messages projection doesn't surface subgraph-scoped messages."""
graph = _make_messages_subgraph()
run = graph.stream_events({"messages": ["hi"], "done": False}, version="v3")
run = graph.stream_v2({"messages": ["hi"], "done": False})
root_streams = list(run.messages)
# The message is emitted inside the subgraph, so the root
# messages projection (scoped to root namespace) doesn't see it.
@@ -389,7 +385,7 @@ class TestStreamV2E2EMessages:
def test_subgraph_handle_messages_drill_down(self) -> None:
"""Drilling into subgraph handle's messages surfaces subgraph messages."""
graph = _make_messages_subgraph()
run = graph.stream_events({"messages": ["hi"], "done": False}, version="v3")
run = graph.stream_v2({"messages": ["hi"], "done": False})
found_messages = False
for handle in run.subgraphs:
@@ -411,9 +407,8 @@ class TestStreamV2E2ECustom:
"""Custom StreamWriter events appear on the main log when a
transformer declares the custom mode."""
graph = _make_custom_writer_graph()
run = graph.stream_events(
run = graph.stream_v2(
{"value": "x", "items": []},
version="v3",
transformers=[_CustomPassthroughTransformer],
)
events = list(run)
@@ -425,7 +420,7 @@ class TestStreamV2E2ECustom:
def test_custom_events_suppressed_without_transformer(self) -> None:
"""Without a custom-mode transformer, custom events don't flow."""
graph = _make_custom_writer_graph()
run = graph.stream_events({"value": "x", "items": []}, version="v3")
run = graph.stream_v2({"value": "x", "items": []})
events = list(run)
custom = [e for e in events if e["method"] == "custom"]
assert custom == []
@@ -433,9 +428,8 @@ class TestStreamV2E2ECustom:
def test_custom_transformer_with_stream_channel(self) -> None:
"""A custom transformer with a StreamChannel produces extension data."""
graph = _make_nested_graph()
run = graph.stream_events(
run = graph.stream_v2(
{"value": "x", "items": []},
version="v3",
transformers=[_CounterTransformer],
)
@@ -450,9 +444,8 @@ class TestStreamV2E2ECustom:
def test_custom_channel_events_on_main_log(self) -> None:
"""StreamChannel auto-forward injects custom:<name> events into the main log."""
graph = _make_nested_graph()
run = graph.stream_events(
run = graph.stream_v2(
{"value": "x", "items": []},
version="v3",
transformers=[_CounterTransformer],
)
events = list(run)
@@ -471,7 +464,7 @@ class TestStreamV2E2EInterrupt:
"""Interrupted run has correct flags and interrupt payloads."""
graph = _make_interrupt_graph()
config: dict[str, Any] = {"configurable": {"thread_id": "int-1"}}
run = graph.stream_events({"value": "x", "items": []}, config, version="v3")
run = graph.stream_v2({"value": "x", "items": []}, config)
output = run.output
assert output is not None
@@ -484,7 +477,7 @@ class TestStreamV2E2EInterrupt:
"""Values snapshots captured before the interrupt reflect partial state."""
graph = _make_interrupt_graph()
config: dict[str, Any] = {"configurable": {"thread_id": "int-2"}}
run = graph.stream_events({"value": "x", "items": []}, config, version="v3")
run = graph.stream_v2({"value": "x", "items": []}, config)
snapshots = list(run.values)
assert len(snapshots) >= 1
@@ -501,14 +494,14 @@ class TestStreamV2E2EErrors:
def test_subgraph_error_propagates_through_output(self) -> None:
"""Error in a subgraph propagates through output."""
graph = _make_error_subgraph()
run = graph.stream_events({"value": "x", "items": []}, version="v3")
run = graph.stream_v2({"value": "x", "items": []})
with pytest.raises(ValueError, match="subgraph explosion"):
_ = run.output
def test_subgraph_error_propagates_through_raw_events(self) -> None:
graph = _make_error_subgraph()
run = graph.stream_events({"value": "x", "items": []}, version="v3")
run = graph.stream_v2({"value": "x", "items": []})
with pytest.raises(ValueError, match="subgraph explosion"):
list(run)
@@ -516,7 +509,7 @@ class TestStreamV2E2EErrors:
def test_error_subgraph_handle_status(self) -> None:
"""Subgraph handle surfaces the error status."""
graph = _make_error_subgraph()
run = graph.stream_events({"value": "x", "items": []}, version="v3")
run = graph.stream_v2({"value": "x", "items": []})
handle = next(iter(run.subgraphs))
with pytest.raises(RuntimeError, match="subgraph explosion"):
@@ -536,7 +529,7 @@ class TestStreamV2E2EAsync:
async def test_all_projections_async(self) -> None:
"""Async run exercises values projection."""
graph = _make_nested_graph()
run = await graph.astream_events({"value": "x", "items": []}, version="v3")
run = await graph.astream_v2({"value": "x", "items": []})
values_snapshots = [s async for s in run.values]
assert len(values_snapshots) >= 1
@@ -547,7 +540,7 @@ class TestStreamV2E2EAsync:
async def test_async_output(self) -> None:
"""Async output returns the final state."""
graph = _make_nested_graph()
run = await graph.astream_events({"value": "x", "items": []}, version="v3")
run = await graph.astream_v2({"value": "x", "items": []})
output = await run.output()
assert output is not None
assert output["value"] == "x_routed_processed"
@@ -557,7 +550,7 @@ class TestStreamV2E2EAsync:
async def test_async_raw_events(self) -> None:
"""Async raw event iteration yields well-formed ProtocolEvents."""
graph = _make_nested_graph()
run = await graph.astream_events({"value": "x", "items": []}, version="v3")
run = await graph.astream_v2({"value": "x", "items": []})
events = [e async for e in run]
assert len(events) > 0
seqs = [e["seq"] for e in events]
@@ -579,7 +572,7 @@ class TestStreamV2E2EAsync:
.compile()
)
run = await graph.astream_events({"messages": "hi"}, version="v3")
run = await graph.astream_v2({"messages": "hi"})
streams = [s async for s in run.messages]
assert len(streams) >= 1
for s in streams:
@@ -590,9 +583,7 @@ class TestStreamV2E2EAsync:
"""Async interrupted run has correct flags."""
graph = _make_interrupt_graph()
config: dict[str, Any] = {"configurable": {"thread_id": "async-int-1"}}
run = await graph.astream_events(
{"value": "x", "items": []}, config, version="v3"
)
run = await graph.astream_v2({"value": "x", "items": []}, config)
output = await run.output()
assert output is not None
@@ -602,14 +593,14 @@ class TestStreamV2E2EAsync:
async def test_async_error_propagation(self) -> None:
"""Async error from subgraph propagates through output."""
graph = _make_error_subgraph()
run = await graph.astream_events({"value": "x", "items": []}, version="v3")
run = await graph.astream_v2({"value": "x", "items": []})
with pytest.raises(ValueError, match="subgraph explosion"):
await run.output()
async def test_async_context_manager(self) -> None:
"""Async context manager calls abort on exit."""
graph = _make_nested_graph()
run = await graph.astream_events({"value": "x", "items": []}, version="v3")
run = await graph.astream_v2({"value": "x", "items": []})
async with run:
_ = await anext(aiter(run.values))
assert run._exhausted is True
@@ -617,7 +608,7 @@ class TestStreamV2E2EAsync:
async def test_async_extensions_present(self) -> None:
"""Async run has all native extensions."""
graph = _make_nested_graph()
run = await graph.astream_events({"value": "x", "items": []}, version="v3")
run = await graph.astream_v2({"value": "x", "items": []})
_ = await run.output()
assert "values" in run.extensions
assert "messages" in run.extensions
@@ -627,9 +618,8 @@ class TestStreamV2E2EAsync:
async def test_async_custom_transformer(self) -> None:
"""Async custom transformer with StreamChannel works."""
graph = _make_nested_graph()
run = await graph.astream_events(
run = await graph.astream_v2(
{"value": "x", "items": []},
version="v3",
transformers=[_CounterTransformer],
)
assert "counter" in run.extensions
@@ -649,7 +639,7 @@ class TestStreamV2E2ECombined:
def test_interleave_all_native_projections(self) -> None:
"""Interleave values + messages + lifecycle without deadlock."""
graph = _make_nested_graph()
run = graph.stream_events({"value": "x", "items": []}, version="v3")
run = graph.stream_v2({"value": "x", "items": []})
seen_names: set[str] = set()
for name, _item in run.interleave("values", "messages", "lifecycle"):
@@ -677,9 +667,8 @@ class TestStreamV2E2ECombined:
return True
graph = _make_nested_graph()
run = graph.stream_events(
run = graph.stream_v2(
{"value": "x", "items": []},
version="v3",
transformers=[_CounterTransformer, TagTransformer],
)
@@ -733,7 +722,7 @@ class TestStreamV2E2ECombined:
.compile()
)
run = outer.stream_events({"items": []}, version="v3")
run = outer.stream_v2({"items": []})
handles = []
for handle in run.subgraphs:
list(handle.values)
@@ -751,17 +740,13 @@ class TestStreamV2E2ECombined:
def test_lifecycle_matches_subgraph_handles(self) -> None:
"""Lifecycle events and subgraph handles agree on discovered subgraphs."""
run1 = _make_nested_graph().stream_events(
{"value": "x", "items": []}, version="v3"
)
run1 = _make_nested_graph().stream_v2({"value": "x", "items": []})
handle_paths: list[tuple[str, ...]] = []
for handle in run1.subgraphs:
list(handle.values)
handle_paths.append(handle.path)
run2 = _make_nested_graph().stream_events(
{"value": "x", "items": []}, version="v3"
)
run2 = _make_nested_graph().stream_v2({"value": "x", "items": []})
lifecycle = list(run2.lifecycle)
started_ns = [
@@ -788,9 +773,8 @@ class TestStreamV2E2ECombined:
.compile()
)
run = graph.stream_events(
run = graph.stream_v2(
{"messages": "hi"},
version="v3",
transformers=[_CounterTransformer],
)
+10 -10
View File
@@ -1348,7 +1348,7 @@ wheels = [
[[package]]
name = "langchain-core"
version = "1.4.0a2"
version = "1.3.2"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "jsonpatch" },
@@ -1361,26 +1361,26 @@ dependencies = [
{ name = "typing-extensions" },
{ name = "uuid-utils" },
]
sdist = { url = "https://files.pythonhosted.org/packages/3c/93/68bafa047f8e1770d0cf0f61d6c70889f1dec42ef6bd263540d916c421b9/langchain_core-1.4.0a2.tar.gz", hash = "sha256:b723c7961b615c7f2180ce2bcf352fdad8247bc51a60adecd3d97088235c120d", size = 916486, upload-time = "2026-05-01T15:02:19.029Z" }
sdist = { url = "https://files.pythonhosted.org/packages/a8/03/7219502e8ca728d65eb44d7a3eb60239230742a70dbfc9241b9bfd61c4ab/langchain_core-1.3.2.tar.gz", hash = "sha256:fd7a50b2f28ba561fd9d7f5d2760bc9e06cf00cdf820a3ccafe88a94ffa8d5b7", size = 911813, upload-time = "2026-04-24T15:49:23.699Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/4e/8e/933e0ba7ba0430ce264e36b178d581b255239ad45093872483142d93478c/langchain_core-1.4.0a2-py3-none-any.whl", hash = "sha256:a5c689f8404357df797120c012da7704144a953b2ae18f258df263301e7badd5", size = 546297, upload-time = "2026-05-01T15:02:17.731Z" },
{ url = "https://files.pythonhosted.org/packages/7d/d5/8fa4431007cbb7cfed7590f4d6a5dea3ad724f4174d248f6642ef5ce7d05/langchain_core-1.3.2-py3-none-any.whl", hash = "sha256:d44a66127f9f8db735bdfd0ab9661bccb47a97113cfd3f2d89c74864422b7274", size = 542390, upload-time = "2026-04-24T15:49:21.991Z" },
]
[[package]]
name = "langchain-protocol"
version = "0.0.14"
version = "0.0.12"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "typing-extensions" },
]
sdist = { url = "https://files.pythonhosted.org/packages/05/bf/efb5e2ed832e4d6d45590e25a9e5191986b291b543bc6a807b48bee070b0/langchain_protocol-0.0.14.tar.gz", hash = "sha256:bc1e8553122e6ede310280462d5813023a172ff2785ccbbdec54d43f3a15e5f2", size = 5862, upload-time = "2026-04-29T16:40:18.657Z" }
sdist = { url = "https://files.pythonhosted.org/packages/5c/51/1157009b6f94e6e58be58fa8b620187d657909a8b36a6bf5b0c52a2711f6/langchain_protocol-0.0.12.tar.gz", hash = "sha256:5e14c434290a705c9510fdb1a83ecf7561a5e6e0dfd053930ade80dba069269f", size = 6408, upload-time = "2026-04-25T01:05:01.489Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/c2/e9/06c47ecb2aff08f83dfa30058da3bf86be64862c19569043ed5331bbeecd/langchain_protocol-0.0.14-py3-none-any.whl", hash = "sha256:ffc35089779bd8ca217015180cef5e660fc3b074efdaa0f2e95df73583f1a047", size = 6984, upload-time = "2026-04-29T16:40:17.841Z" },
{ url = "https://files.pythonhosted.org/packages/95/82/3431e3061c917439589fa88a6b23c9bc0e154cba0f05d2e895a68c76ff74/langchain_protocol-0.0.12-py3-none-any.whl", hash = "sha256:402b61f42d4139692528cf37226c367bb6efc8ff8165b29380accb0abfece7b2", size = 6639, upload-time = "2026-04-25T01:05:00.487Z" },
]
[[package]]
name = "langgraph"
version = "1.2.0a3"
version = "1.1.10"
source = { editable = "." }
dependencies = [
{ name = "langchain-core" },
@@ -1452,7 +1452,7 @@ test = [
[package.metadata]
requires-dist = [
{ name = "langchain-core", specifier = ">=1.4.0a2,<2" },
{ name = "langchain-core", specifier = ">=1.3.2,<2" },
{ name = "langgraph-checkpoint", editable = "../checkpoint" },
{ name = "langgraph-prebuilt", editable = "../prebuilt" },
{ name = "langgraph-sdk", editable = "../sdk-py" },
@@ -1561,7 +1561,7 @@ wheels = [
[[package]]
name = "langgraph-checkpoint"
version = "4.1.0a3"
version = "4.0.3"
source = { editable = "../checkpoint" }
dependencies = [
{ name = "langchain-core" },
@@ -1609,7 +1609,7 @@ test = [
[[package]]
name = "langgraph-checkpoint-postgres"
version = "3.1.0a3"
version = "3.0.5"
source = { editable = "../checkpoint-postgres" }
dependencies = [
{ name = "langgraph-checkpoint" },
@@ -44,7 +44,7 @@ import inspect
import json
from collections.abc import Awaitable, Callable
from copy import copy, deepcopy
from dataclasses import dataclass, field, replace
from dataclasses import dataclass, replace
from types import UnionType
from typing import (
TYPE_CHECKING,
@@ -1723,9 +1723,9 @@ class ToolRuntime(_DirectlyInjectedToolArg, Generic[ContextT, StateT]):
context: ContextT
config: RunnableConfig
stream_writer: StreamWriter
tools: list[BaseTool]
tool_call_id: str | None
store: BaseStore | None
tools: list[BaseTool] = field(default_factory=list)
execution_info: ExecutionInfo | None = None
server_info: ServerInfo | None = None
@@ -30,11 +30,6 @@ from langgraph.prebuilt._tool_call_stream import ToolCallStream
TS = int(time.time() * 1000)
def _unstamped(items):
"""Strip push stamps from a StreamChannel's internal buffer."""
return [item for _stamp, item in items]
def _tool_event(
event: str,
tool_call_id: str,
@@ -100,7 +95,7 @@ class TestToolCallTransformerUnit:
input={"text": "hi"},
)
)
handles = _unstamped(transformer._log._items)
handles = list(transformer._log._items)
assert len(handles) == 1
h = handles[0]
assert isinstance(h, ToolCallStream)
@@ -116,7 +111,7 @@ class TestToolCallTransformerUnit:
mux.push(_tool_event("tool-output-delta", "tc1", delta="a"))
mux.push(_tool_event("tool-output-delta", "tc1", delta="b"))
stream = transformer._active["tc1"]
assert _unstamped(stream._output_deltas._items) == ["a", "b"]
assert list(stream._output_deltas._items) == ["a", "b"]
def test_finish_closes_stream(self) -> None:
mux, transformer = _mux()
@@ -147,17 +142,14 @@ class TestToolCallTransformerUnit:
mux.push(_tool_event("tool-output-delta", "a", delta="A1"))
mux.push(_tool_event("tool-output-delta", "b", delta="B1"))
mux.push(_tool_event("tool-output-delta", "a", delta="A2"))
assert _unstamped(transformer._active["a"]._output_deltas._items) == [
"A1",
"A2",
]
assert _unstamped(transformer._active["b"]._output_deltas._items) == ["B1"]
assert list(transformer._active["a"]._output_deltas._items) == ["A1", "A2"]
assert list(transformer._active["b"]._output_deltas._items) == ["B1"]
def test_tools_event_passes_through_main_log(self) -> None:
mux, transformer = _mux()
_subscribe(mux._events)
mux.push(_tool_event("tool-started", "tc1", tool_name="echo"))
kept = [e for e in _unstamped(mux._events._items) if e["method"] == "tools"]
kept = [e for e in mux._events._items if e["method"] == "tools"]
assert len(kept) == 1
@@ -202,9 +194,7 @@ class TestToolCallTransformerEndToEnd:
}
graph = _build_graph(caller, [streamer])
run = graph.stream_events(
{"messages": []}, transformers=[ToolCallTransformer], version="v3"
)
run = graph.stream_v2({"messages": []}, transformers=[ToolCallTransformer])
tool_calls: list[ToolCallStream] = []
for tc in run.tool_calls:
@@ -240,13 +230,11 @@ class TestToolCallTransformerEndToEnd:
# Without ToolCallTransformer, no tool_calls projection is
# exposed and no `tools` events flow through (required_stream_modes
# omits it).
run_no_tc = graph.stream_events({"messages": []}, version="v3")
run_no_tc = graph.stream_v2({"messages": []})
assert "tool_calls" not in run_no_tc._mux.extensions # type: ignore[attr-defined]
# With ToolCallTransformer, the projection is present.
run = graph.stream_events(
{"messages": []}, transformers=[ToolCallTransformer], version="v3"
)
run = graph.stream_v2({"messages": []}, transformers=[ToolCallTransformer])
assert "tool_calls" in run._mux.extensions # type: ignore[attr-defined]
# Drain so the run closes cleanly.
list(run.tool_calls)
@@ -273,8 +261,8 @@ class TestToolCallTransformerEndToEnd:
}
graph = _build_graph(caller, [astreamer])
run = await graph.astream_events(
{"messages": []}, version="v3", transformers=[ToolCallTransformer]
run = await graph.astream_v2(
{"messages": []}, transformers=[ToolCallTransformer]
)
collected: list[ToolCallStream] = []
@@ -303,9 +291,7 @@ class TestToolCallTransformerEndToEnd:
}
graph = _build_graph(caller, [boom])
run = graph.stream_events(
{"messages": []}, transformers=[ToolCallTransformer], version="v3"
)
run = graph.stream_v2({"messages": []}, transformers=[ToolCallTransformer])
collected: list[ToolCallStream] = []
with pytest.raises(ValueError, match="nope"):
-13
View File
@@ -2016,19 +2016,6 @@ async def test_tool_node_inject_runtime_dynamic_tool_via_wrap_tool_call_async()
assert tool_message.tool_call_id == "call_dynamic_2"
def test_tool_runtime_defaults_tools_to_empty_list() -> None:
runtime = ToolRuntime(
state={},
context=None,
config={},
stream_writer=lambda *args, **kwargs: None,
tool_call_id=None,
store=None,
)
assert runtime.tools == []
def test_tool_runtime_forwards_execution_info_server_info_and_tools() -> None:
"""Test that execution_info, server_info, and tools are forwarded from Runtime to ToolRuntime."""
from langgraph.runtime import ExecutionInfo, ServerInfo
+10 -10
View File
@@ -249,7 +249,7 @@ wheels = [
[[package]]
name = "langchain-core"
version = "1.4.0a2"
version = "1.3.2"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "jsonpatch" },
@@ -262,26 +262,26 @@ dependencies = [
{ name = "typing-extensions" },
{ name = "uuid-utils" },
]
sdist = { url = "https://files.pythonhosted.org/packages/3c/93/68bafa047f8e1770d0cf0f61d6c70889f1dec42ef6bd263540d916c421b9/langchain_core-1.4.0a2.tar.gz", hash = "sha256:b723c7961b615c7f2180ce2bcf352fdad8247bc51a60adecd3d97088235c120d", size = 916486, upload-time = "2026-05-01T15:02:19.029Z" }
sdist = { url = "https://files.pythonhosted.org/packages/a8/03/7219502e8ca728d65eb44d7a3eb60239230742a70dbfc9241b9bfd61c4ab/langchain_core-1.3.2.tar.gz", hash = "sha256:fd7a50b2f28ba561fd9d7f5d2760bc9e06cf00cdf820a3ccafe88a94ffa8d5b7", size = 911813, upload-time = "2026-04-24T15:49:23.699Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/4e/8e/933e0ba7ba0430ce264e36b178d581b255239ad45093872483142d93478c/langchain_core-1.4.0a2-py3-none-any.whl", hash = "sha256:a5c689f8404357df797120c012da7704144a953b2ae18f258df263301e7badd5", size = 546297, upload-time = "2026-05-01T15:02:17.731Z" },
{ url = "https://files.pythonhosted.org/packages/7d/d5/8fa4431007cbb7cfed7590f4d6a5dea3ad724f4174d248f6642ef5ce7d05/langchain_core-1.3.2-py3-none-any.whl", hash = "sha256:d44a66127f9f8db735bdfd0ab9661bccb47a97113cfd3f2d89c74864422b7274", size = 542390, upload-time = "2026-04-24T15:49:21.991Z" },
]
[[package]]
name = "langchain-protocol"
version = "0.0.14"
version = "0.0.12"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "typing-extensions" },
]
sdist = { url = "https://files.pythonhosted.org/packages/05/bf/efb5e2ed832e4d6d45590e25a9e5191986b291b543bc6a807b48bee070b0/langchain_protocol-0.0.14.tar.gz", hash = "sha256:bc1e8553122e6ede310280462d5813023a172ff2785ccbbdec54d43f3a15e5f2", size = 5862, upload-time = "2026-04-29T16:40:18.657Z" }
sdist = { url = "https://files.pythonhosted.org/packages/5c/51/1157009b6f94e6e58be58fa8b620187d657909a8b36a6bf5b0c52a2711f6/langchain_protocol-0.0.12.tar.gz", hash = "sha256:5e14c434290a705c9510fdb1a83ecf7561a5e6e0dfd053930ade80dba069269f", size = 6408, upload-time = "2026-04-25T01:05:01.489Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/c2/e9/06c47ecb2aff08f83dfa30058da3bf86be64862c19569043ed5331bbeecd/langchain_protocol-0.0.14-py3-none-any.whl", hash = "sha256:ffc35089779bd8ca217015180cef5e660fc3b074efdaa0f2e95df73583f1a047", size = 6984, upload-time = "2026-04-29T16:40:17.841Z" },
{ url = "https://files.pythonhosted.org/packages/95/82/3431e3061c917439589fa88a6b23c9bc0e154cba0f05d2e895a68c76ff74/langchain_protocol-0.0.12-py3-none-any.whl", hash = "sha256:402b61f42d4139692528cf37226c367bb6efc8ff8165b29380accb0abfece7b2", size = 6639, upload-time = "2026-04-25T01:05:00.487Z" },
]
[[package]]
name = "langgraph"
version = "1.2.0a3"
version = "1.1.10"
source = { editable = "../langgraph" }
dependencies = [
{ name = "langchain-core" },
@@ -294,7 +294,7 @@ dependencies = [
[package.metadata]
requires-dist = [
{ name = "langchain-core", specifier = ">=1.4.0a2,<2" },
{ name = "langchain-core", specifier = ">=1.3.2,<2" },
{ name = "langgraph-checkpoint", editable = "../checkpoint" },
{ name = "langgraph-prebuilt", editable = "." },
{ name = "langgraph-sdk", editable = "../sdk-py" },
@@ -365,7 +365,7 @@ test = [
[[package]]
name = "langgraph-checkpoint"
version = "4.1.0a3"
version = "4.0.3"
source = { editable = "../checkpoint" }
dependencies = [
{ name = "langchain-core" },
@@ -413,7 +413,7 @@ test = [
[[package]]
name = "langgraph-checkpoint-postgres"
version = "3.1.0a3"
version = "3.0.5"
source = { editable = "../checkpoint-postgres" }
dependencies = [
{ name = "langgraph-checkpoint" },
+9 -9
View File
@@ -266,7 +266,7 @@ wheels = [
[[package]]
name = "langchain-core"
version = "1.4.0a2"
version = "1.3.2"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "jsonpatch" },
@@ -279,26 +279,26 @@ dependencies = [
{ name = "typing-extensions" },
{ name = "uuid-utils" },
]
sdist = { url = "https://files.pythonhosted.org/packages/3c/93/68bafa047f8e1770d0cf0f61d6c70889f1dec42ef6bd263540d916c421b9/langchain_core-1.4.0a2.tar.gz", hash = "sha256:b723c7961b615c7f2180ce2bcf352fdad8247bc51a60adecd3d97088235c120d", size = 916486, upload-time = "2026-05-01T15:02:19.029Z" }
sdist = { url = "https://files.pythonhosted.org/packages/a8/03/7219502e8ca728d65eb44d7a3eb60239230742a70dbfc9241b9bfd61c4ab/langchain_core-1.3.2.tar.gz", hash = "sha256:fd7a50b2f28ba561fd9d7f5d2760bc9e06cf00cdf820a3ccafe88a94ffa8d5b7", size = 911813, upload-time = "2026-04-24T15:49:23.699Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/4e/8e/933e0ba7ba0430ce264e36b178d581b255239ad45093872483142d93478c/langchain_core-1.4.0a2-py3-none-any.whl", hash = "sha256:a5c689f8404357df797120c012da7704144a953b2ae18f258df263301e7badd5", size = 546297, upload-time = "2026-05-01T15:02:17.731Z" },
{ url = "https://files.pythonhosted.org/packages/7d/d5/8fa4431007cbb7cfed7590f4d6a5dea3ad724f4174d248f6642ef5ce7d05/langchain_core-1.3.2-py3-none-any.whl", hash = "sha256:d44a66127f9f8db735bdfd0ab9661bccb47a97113cfd3f2d89c74864422b7274", size = 542390, upload-time = "2026-04-24T15:49:21.991Z" },
]
[[package]]
name = "langchain-protocol"
version = "0.0.14"
version = "0.0.12"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "typing-extensions" },
]
sdist = { url = "https://files.pythonhosted.org/packages/05/bf/efb5e2ed832e4d6d45590e25a9e5191986b291b543bc6a807b48bee070b0/langchain_protocol-0.0.14.tar.gz", hash = "sha256:bc1e8553122e6ede310280462d5813023a172ff2785ccbbdec54d43f3a15e5f2", size = 5862, upload-time = "2026-04-29T16:40:18.657Z" }
sdist = { url = "https://files.pythonhosted.org/packages/5c/51/1157009b6f94e6e58be58fa8b620187d657909a8b36a6bf5b0c52a2711f6/langchain_protocol-0.0.12.tar.gz", hash = "sha256:5e14c434290a705c9510fdb1a83ecf7561a5e6e0dfd053930ade80dba069269f", size = 6408, upload-time = "2026-04-25T01:05:01.489Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/c2/e9/06c47ecb2aff08f83dfa30058da3bf86be64862c19569043ed5331bbeecd/langchain_protocol-0.0.14-py3-none-any.whl", hash = "sha256:ffc35089779bd8ca217015180cef5e660fc3b074efdaa0f2e95df73583f1a047", size = 6984, upload-time = "2026-04-29T16:40:17.841Z" },
{ url = "https://files.pythonhosted.org/packages/95/82/3431e3061c917439589fa88a6b23c9bc0e154cba0f05d2e895a68c76ff74/langchain_protocol-0.0.12-py3-none-any.whl", hash = "sha256:402b61f42d4139692528cf37226c367bb6efc8ff8165b29380accb0abfece7b2", size = 6639, upload-time = "2026-04-25T01:05:00.487Z" },
]
[[package]]
name = "langgraph"
version = "1.2.0a3"
version = "1.1.10"
source = { editable = "../langgraph" }
dependencies = [
{ name = "langchain-core" },
@@ -311,7 +311,7 @@ dependencies = [
[package.metadata]
requires-dist = [
{ name = "langchain-core", specifier = ">=1.4.0a2,<2" },
{ name = "langchain-core", specifier = ">=1.3.2,<2" },
{ name = "langgraph-checkpoint", editable = "../checkpoint" },
{ name = "langgraph-prebuilt", editable = "../prebuilt" },
{ name = "langgraph-sdk", editable = "." },
@@ -382,7 +382,7 @@ test = [
[[package]]
name = "langgraph-checkpoint"
version = "4.1.0a3"
version = "4.0.3"
source = { editable = "../checkpoint" }
dependencies = [
{ name = "langchain-core" },