Merge branch 'main' into vb/update-get-state

libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py

@@ -318,16 +318,6 @@ class PostgresSaver(BasePostgresSaver):
             task_id (str): Identifier for the task creating the writes.
         """
         with self._cursor(pipeline=True) as cur:
-            cur.execute(
-                self.DELETE_WRITES_SQL,
-                (
-                    config["configurable"]["thread_id"],
-                    config["configurable"]["checkpoint_ns"],
-                    config["configurable"]["checkpoint_id"],
-                    task_id,
-                    len(writes),
-                ),
-            )
             cur.executemany(
                 self.UPSERT_CHECKPOINT_WRITES_SQL,
                 self._dump_writes(

libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py

@@ -274,16 +274,6 @@ class AsyncPostgresSaver(BasePostgresSaver):
             task_id (str): Identifier for the task creating the writes.
         """
         async with self._cursor(pipeline=True) as cur:
-            await cur.execute(
-                self.DELETE_WRITES_SQL,
-                (
-                    config["configurable"]["thread_id"],
-                    config["configurable"]["checkpoint_ns"],
-                    config["configurable"]["checkpoint_id"],
-                    task_id,
-                    len(writes),
-                ),
-            )
             await cur.executemany(
                 self.UPSERT_CHECKPOINT_WRITES_SQL,
                 await asyncio.to_thread(

libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py

@@ -6,6 +6,7 @@ from langchain_core.runnables import RunnableConfig
 from psycopg.types.json import Jsonb
 
 from langgraph.checkpoint.base import (
+    WRITES_IDX_MAP,
     BaseCheckpointSaver,
     Checkpoint,
     EmptyChannelError,
@@ -105,15 +106,6 @@ UPSERT_CHECKPOINT_WRITES_SQL = """
     ON CONFLICT (thread_id, checkpoint_ns, checkpoint_id, task_id, idx) DO NOTHING
 """
 
-DELETE_WRITES_SQL = """
-    DELETE FROM checkpoint_writes
-    WHERE thread_id = %s
-    AND checkpoint_ns = %s
-    AND checkpoint_id = %s
-    AND task_id = %s
-    AND idx >= %s
-"""
-
 
 class BasePostgresSaver(BaseCheckpointSaver):
     SELECT_SQL = SELECT_SQL
@@ -121,7 +113,6 @@ class BasePostgresSaver(BaseCheckpointSaver):
     UPSERT_CHECKPOINT_BLOBS_SQL = UPSERT_CHECKPOINT_BLOBS_SQL
     UPSERT_CHECKPOINTS_SQL = UPSERT_CHECKPOINTS_SQL
     UPSERT_CHECKPOINT_WRITES_SQL = UPSERT_CHECKPOINT_WRITES_SQL
-    DELETE_WRITES_SQL = DELETE_WRITES_SQL
 
     jsonplus_serde = JsonPlusSerializer()
 
@@ -210,7 +201,7 @@ class BasePostgresSaver(BaseCheckpointSaver):
                 checkpoint_ns,
                 checkpoint_id,
                 task_id,
-                idx,
+                WRITES_IDX_MAP.get(channel, idx),
                 channel,
                 *self.serde.dumps_typed(value),
             )

libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/__init__.py

@@ -1,12 +1,13 @@
 import sqlite3
 import threading
-from contextlib import contextmanager
+from contextlib import closing, contextmanager
 from hashlib import md5
 from typing import Any, AsyncIterator, Dict, Iterator, Optional, Sequence, Tuple
 
 from langchain_core.runnables import RunnableConfig
 
 from langgraph.checkpoint.base import (
+    WRITES_IDX_MAP,
     BaseCheckpointSaver,
     ChannelVersions,
     Checkpoint,
@@ -318,9 +319,7 @@ class SqliteSaver(BaseCheckpointSaver):
         ORDER BY checkpoint_id DESC"""
         if limit:
             query += f" LIMIT {limit}"
-        with self.cursor(transaction=False) as cur, self.cursor(
-            transaction=False
-        ) as writes_cur:
+        with self.cursor(transaction=False) as cur, closing(self.conn.cursor()) as wcur:
             cur.execute(query, param_values)
             for (
                 thread_id,
@@ -331,13 +330,9 @@ class SqliteSaver(BaseCheckpointSaver):
                 checkpoint,
                 metadata,
             ) in cur:
-                writes_cur.execute(
+                wcur.execute(
                     "SELECT task_id, channel, type, value FROM writes WHERE thread_id = ? AND checkpoint_ns = ? AND checkpoint_id = ?",
-                    (
-                        thread_id,
-                        checkpoint_ns,
-                        checkpoint_id,
-                    ),
+                    (thread_id, checkpoint_ns, checkpoint_id),
                 )
                 yield CheckpointTuple(
                     {
@@ -362,7 +357,7 @@ class SqliteSaver(BaseCheckpointSaver):
                     ),
                     [
                         (task_id, channel, self.serde.loads_typed((type, value)))
-                        for task_id, channel, type, value in writes_cur
+                        for task_id, channel, type, value in wcur
                     ],
                 )
 
@@ -438,25 +433,15 @@ class SqliteSaver(BaseCheckpointSaver):
             task_id (str): Identifier for the task creating the writes.
         """
         with self.lock, self.cursor() as cur:
-            cur.execute(
-                "DELETE FROM writes WHERE thread_id = ? AND checkpoint_ns = ? AND checkpoint_id = ? AND task_id = ? AND idx >= ?",
-                (
-                    str(config["configurable"]["thread_id"]),
-                    str(config["configurable"]["checkpoint_ns"]),
-                    str(config["configurable"]["checkpoint_id"]),
-                    task_id,
-                    len(writes),
-                ),
-            )
             cur.executemany(
-                "INSERT OR REPLACE INTO writes (thread_id, checkpoint_ns, checkpoint_id, task_id, idx, channel, type, value) VALUES (?, ?, ?, ?, ?, ?, ?, ?)",
+                "INSERT OR IGNORE INTO writes (thread_id, checkpoint_ns, checkpoint_id, task_id, idx, channel, type, value) VALUES (?, ?, ?, ?, ?, ?, ?, ?)",
                 [
                     (
                         str(config["configurable"]["thread_id"]),
                         str(config["configurable"]["checkpoint_ns"]),
                         str(config["configurable"]["checkpoint_id"]),
                         task_id,
-                        idx,
+                        WRITES_IDX_MAP.get(channel, idx),
                         channel,
                         *self.serde.dumps_typed(value),
                     )

libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/aio.py

@@ -16,6 +16,7 @@ import aiosqlite
 from langchain_core.runnables import RunnableConfig
 
 from langgraph.checkpoint.base import (
+    WRITES_IDX_MAP,
     BaseCheckpointSaver,
     ChannelVersions,
     Checkpoint,
@@ -329,14 +330,14 @@ class AsyncSqliteSaver(BaseCheckpointSaver):
             AsyncIterator[CheckpointTuple]: An asynchronous iterator of matching checkpoint tuples.
         """
         await self.setup()
-        where, param_values = search_where(config, filter, before)
+        where, params = search_where(config, filter, before)
         query = f"""SELECT thread_id, checkpoint_ns, checkpoint_id, parent_checkpoint_id, type, checkpoint, metadata
         FROM checkpoints
         {where}
         ORDER BY checkpoint_id DESC"""
         if limit:
             query += f" LIMIT {limit}"
-        async with self.conn.execute(query, param_values) as cursor:
+        async with self.conn.execute(query, params) as cur, self.conn.cursor() as wcur:
             async for (
                 thread_id,
                 checkpoint_ns,
@@ -345,14 +346,10 @@ class AsyncSqliteSaver(BaseCheckpointSaver):
                 type,
                 checkpoint,
                 metadata,
-            ) in cursor:
-                writes_cur = await self.conn.execute(
+            ) in cur:
+                await wcur.execute(
                     "SELECT task_id, channel, type, value FROM writes WHERE thread_id = ? AND checkpoint_ns = ? AND checkpoint_id = ?",
-                    (
-                        thread_id,
-                        checkpoint_ns,
-                        checkpoint_id,
-                    ),
+                    (thread_id, checkpoint_ns, checkpoint_id),
                 )
                 yield CheckpointTuple(
                     {
@@ -377,7 +374,7 @@ class AsyncSqliteSaver(BaseCheckpointSaver):
                     ),
                     [
                         (task_id, channel, self.serde.loads_typed((type, value)))
-                        async for task_id, channel, type, value in writes_cur
+                        async for task_id, channel, type, value in wcur
                     ],
                 )
 
@@ -445,25 +442,15 @@ class AsyncSqliteSaver(BaseCheckpointSaver):
         """
         await self.setup()
         async with self.conn.cursor() as cur:
-            await cur.execute(
-                "DELETE FROM writes WHERE thread_id = ? AND checkpoint_ns = ? AND checkpoint_id = ? AND task_id = ? AND idx >= ?",
-                (
-                    str(config["configurable"]["thread_id"]),
-                    str(config["configurable"]["checkpoint_ns"]),
-                    str(config["configurable"]["checkpoint_id"]),
-                    task_id,
-                    len(writes),
-                ),
-            )
             await cur.executemany(
-                "INSERT OR REPLACE INTO writes (thread_id, checkpoint_ns, checkpoint_id, task_id, idx, channel, type, value) VALUES (?, ?, ?, ?, ?, ?, ?, ?)",
+                "INSERT OR IGNORE INTO writes (thread_id, checkpoint_ns, checkpoint_id, task_id, idx, channel, type, value) VALUES (?, ?, ?, ?, ?, ?, ?, ?)",
                 [
                     (
                         str(config["configurable"]["thread_id"]),
                         str(config["configurable"]["checkpoint_ns"]),
                         str(config["configurable"]["checkpoint_id"]),
                         task_id,
-                        idx,
+                        WRITES_IDX_MAP.get(channel, idx),
                         channel,
                         *self.serde.dumps_typed(value),
                     )

libs/checkpoint/langgraph/checkpoint/base/__init__.py

@@ -22,6 +22,7 @@ from langgraph.checkpoint.base.id import uuid6
 from langgraph.checkpoint.serde.base import SerializerProtocol, maybe_add_typed_methods
 from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
 from langgraph.checkpoint.serde.types import (
+    ERROR,
     ChannelProtocol,
     SendProtocol,
 )
@@ -98,6 +99,7 @@ class Checkpoint(TypedDict):
     Cleared by the next checkpoint."""
     current_tasks: Dict[str, TaskInfo]
     """Map from task ID to task info."""
+    # TODO remove this
 
 
 def empty_checkpoint() -> Checkpoint:
@@ -140,6 +142,8 @@ def create_checkpoint(
     else:
         values: dict[str, Any] = {}
         for k, v in channels.items():
+            if k not in checkpoint["channel_versions"]:
+                continue
             try:
                 values[k] = v.checkpoint()
             except EmptyChannelError:
@@ -437,3 +441,13 @@ def get_checkpoint_id(config: RunnableConfig) -> Optional[str]:
     return config["configurable"].get(
         "checkpoint_id", config["configurable"].get("thread_ts")
     )
+
+
+"""
+Mapping from error type to error index.
+Regular writes just map to their index in the list of writes being saved.
+Special writes (e.g. errors) map to negative indices, to avoid those writes from
+saving regular writes.
+Each Checkpointer implementation should use this mapping in put_writes.
+"""
+WRITES_IDX_MAP = {ERROR: -1}

libs/checkpoint/langgraph/checkpoint/memory/__init__.py

@@ -8,6 +8,7 @@ from typing import Any, AsyncIterator, Dict, Iterator, List, Optional, Tuple
 from langchain_core.runnables import RunnableConfig
 
 from langgraph.checkpoint.base import (
+    WRITES_IDX_MAP,
     BaseCheckpointSaver,
     ChannelVersions,
     Checkpoint,
@@ -52,6 +53,9 @@ class MemorySaver(
 
     # thread ID ->  checkpoint NS -> checkpoint ID -> checkpoint mapping
     storage: defaultdict[str, dict[str, dict[str, tuple[bytes, bytes, Optional[str]]]]]
+    writes: defaultdict[
+        tuple[str, str, str], dict[tuple[str, int], tuple[str, str, bytes]]
+    ]
 
     def __init__(
         self,
@@ -60,7 +64,7 @@ class MemorySaver(
     ) -> None:
         super().__init__(serde=serde)
         self.storage = defaultdict(lambda: defaultdict(dict))
-        self.writes = defaultdict(list)
+        self.writes = defaultdict(dict)
 
     def __enter__(self) -> "MemorySaver":
         return self
@@ -103,7 +107,7 @@ class MemorySaver(
         if checkpoint_id := get_checkpoint_id(config):
             if saved := self.storage[thread_id][checkpoint_ns].get(checkpoint_id):
                 checkpoint, metadata, parent_checkpoint_id = saved
-                writes = self.writes[(thread_id, checkpoint_ns, checkpoint_id)]
+                writes = self.writes[(thread_id, checkpoint_ns, checkpoint_id)].values()
                 return CheckpointTuple(
                     config=config,
                     checkpoint=self.serde.loads_typed(checkpoint),
@@ -125,7 +129,7 @@ class MemorySaver(
             if checkpoints := self.storage[thread_id][checkpoint_ns]:
                 checkpoint_id = max(checkpoints.keys())
                 checkpoint, metadata, parent_checkpoint_id = checkpoints[checkpoint_id]
-                writes = self.writes[(thread_id, checkpoint_ns, checkpoint_id)]
+                writes = self.writes[(thread_id, checkpoint_ns, checkpoint_id)].values()
                 return CheckpointTuple(
                     config={
                         "configurable": {
@@ -206,7 +210,8 @@ class MemorySaver(
                     elif limit is not None:
                         limit -= 1
 
-                    writes = self.writes[(thread_id, checkpoint_ns, checkpoint_id)]
+                    writes = self.writes[(thread_id, checkpoint_ns, checkpoint_id)].values()
+
                     yield CheckpointTuple(
                         config={
                             "configurable": {
@@ -216,9 +221,6 @@ class MemorySaver(
                             }
                         },
                         checkpoint=self.serde.loads_typed(checkpoint),
-                        pending_writes=[
-                            (id, c, self.serde.loads_typed(v)) for id, c, v in writes
-                        ],
                         metadata=metadata,
                         parent_config={
                             "configurable": {
@@ -229,6 +231,9 @@ class MemorySaver(
                         }
                         if parent_checkpoint_id
                         else None,
+                        pending_writes=[
+                            (id, c, self.serde.loads_typed(v)) for id, c, v in writes
+                        ],
                     )
 
     def put(
@@ -293,11 +298,10 @@ class MemorySaver(
         thread_id = config["configurable"]["thread_id"]
         checkpoint_ns = config["configurable"]["checkpoint_ns"]
         checkpoint_id = config["configurable"]["checkpoint_id"]
-        key = (thread_id, checkpoint_ns, checkpoint_id)
-        self.writes[key] = [w for w in self.writes[key] if w[0] != task_id]
-        self.writes[key].extend(
-            [(task_id, c, self.serde.dumps_typed(v)) for c, v in writes]
-        )
+        outer_key = (thread_id, checkpoint_ns, checkpoint_id)
+        for idx, (c, v) in enumerate(writes):
+            inner_key = (task_id, WRITES_IDX_MAP.get(c, idx))
+            self.writes[outer_key][inner_key] = (task_id, c, self.serde.dumps_typed(v))
 
     async def aget_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:
         """Asynchronous version of get_tuple.

libs/checkpoint/langgraph/checkpoint/serde/types.py

@@ -12,6 +12,8 @@ from typing import (
 from langchain_core.runnables import RunnableConfig
 from typing_extensions import Self
 
+ERROR = "__error__"
+
 Value = TypeVar("Value")
 Update = TypeVar("Update")
 C = TypeVar("C")

libs/langgraph/langgraph/channels/base.py

@@ -2,9 +2,9 @@ from abc import ABC, abstractmethod
 from contextlib import asynccontextmanager, contextmanager
 from typing import (
     Any,
-    AsyncGenerator,
-    Generator,
+    AsyncIterator,
     Generic,
+    Iterator,
     Optional,
     Sequence,
     TypeVar,
@@ -21,6 +21,8 @@ C = TypeVar("C")
 
 
 class BaseChannel(Generic[Value, Update, C], ABC):
+    key: str = ""
+
     @property
     @abstractmethod
     def ValueType(self) -> Any:
@@ -43,19 +45,35 @@ class BaseChannel(Generic[Value, Update, C], ABC):
     @abstractmethod
     def from_checkpoint(
         self, checkpoint: Optional[C], config: RunnableConfig
-    ) -> Generator[Self, None, None]:
+    ) -> Iterator[Self]:
         """Return a new identical channel, optionally initialized from a checkpoint.
         If the checkpoint contains complex data structures, they should be copied."""
 
+    @contextmanager
+    def from_checkpoint_named(
+        self, checkpoint: Optional[C], config: RunnableConfig
+    ) -> Iterator[Self]:
+        with self.from_checkpoint(checkpoint, config) as value:
+            value.key = self.key
+            yield value
+
     @asynccontextmanager
     async def afrom_checkpoint(
         self, checkpoint: Optional[C], config: RunnableConfig
-    ) -> AsyncGenerator[Self, None]:
+    ) -> AsyncIterator[Self]:
         """Return a new identical channel, optionally initialized from a checkpoint.
         If the checkpoint contains complex data structures, they should be copied."""
         with self.from_checkpoint(checkpoint, config) as value:
             yield value
 
+    @asynccontextmanager
+    async def afrom_checkpoint_named(
+        self, checkpoint: Optional[C], config: RunnableConfig
+    ) -> AsyncIterator[Self]:
+        async with self.afrom_checkpoint(checkpoint, config) as value:
+            value.key = self.key
+            yield value
+
     # state methods
 
     @abstractmethod

libs/langgraph/langgraph/channels/context.py

@@ -112,7 +112,9 @@ class Context(Generic[Value], BaseChannel[Value, None, None]):
 
     def update(self, values: Sequence[None]) -> bool:
         if values:
-            raise InvalidUpdateError("Context channel does not accept writes.")
+            raise InvalidUpdateError(
+                f"At key '{self.key}': Context channel does not accept writes."
+            )
         return False
 
     def get(self) -> Value:

libs/langgraph/langgraph/channels/dynamic_barrier_value.py

@@ -69,7 +69,7 @@ class DynamicBarrierValue(
         if wait_for_names := [v for v in values if isinstance(v, WaitForNames)]:
             if len(wait_for_names) > 1:
                 raise InvalidUpdateError(
-                    "Received multiple WaitForNames updates in the same step."
+                    f"At key '{self.key}': Received multiple WaitForNames updates in the same step."
                 )
             self.names = wait_for_names[0].names
             return True

libs/langgraph/langgraph/channels/ephemeral_value.py

@@ -58,7 +58,7 @@ class EphemeralValue(Generic[Value], BaseChannel[Value, Value, Value]):
                 return False
         if len(values) != 1 and self.guard:
             raise InvalidUpdateError(
-                "EphemeralValue can only receive one value per step."
+                f"At key '{self.key}': EphemeralValue(guard=True) can receive only one value per step. Use guard=False if you want to store any one of multiple values."
             )
 
         self.value = values[-1]

libs/langgraph/langgraph/channels/last_value.py

@@ -52,7 +52,9 @@ class LastValue(Generic[Value], BaseChannel[Value, Value, Value]):
         if len(values) == 0:
             return False
         if len(values) != 1:
-            raise InvalidUpdateError("LastValue can only receive one value per step.")
+            raise InvalidUpdateError(
+                f"At key '{self.key}': Can receive only one value per step. Use an Annotated key to handle multiple values."
+            )
 
         self.value = values[-1]
         return True

libs/langgraph/langgraph/channels/named_barrier_value.py

@@ -53,7 +53,9 @@ class NamedBarrierValue(Generic[Value], BaseChannel[Value, Value, set[Value]]):
                     self.seen.add(value)
                     updated = True
             else:
-                raise InvalidUpdateError(f"Value {value} not in {self.names}")
+                raise InvalidUpdateError(
+                    f"At key '{self.key}': Value {value} not in {self.names}"
+                )
         return updated
 
     def get(self) -> Value:

libs/langgraph/langgraph/channels/untracked_value.py

@@ -49,7 +49,7 @@ class UntrackedValue(Generic[Value], BaseChannel[Value, Value, Value]):
             return False
         if len(values) != 1 and self.guard:
             raise InvalidUpdateError(
-                "UntrackedValue can only receive one value per step."
+                f"At key '{self.key}': UntrackedValue(guard=True) can receive only one value per step. Use guard=False if you want to store any one of multiple values."
             )
 
         self.value = values[-1]

libs/langgraph/langgraph/errors.py

@@ -1,4 +1,4 @@
-from typing import Any
+from typing import Any, Sequence
 
 from langgraph.checkpoint.base import EmptyChannelError
 from langgraph.constants import Interrupt
@@ -32,7 +32,7 @@ class InvalidUpdateError(Exception):
 class GraphInterrupt(Exception):
     """Raised when a subgraph is interrupted."""
 
-    def __init__(self, interrupts: list[Interrupt]) -> None:
+    def __init__(self, interrupts: Sequence[Interrupt] = ()) -> None:
         super().__init__(interrupts)
 
 

libs/langgraph/langgraph/graph/graph.py

@@ -198,12 +198,14 @@ class Graph:
             raise ValueError("END cannot be a start node")
         if end_key == START:
             raise ValueError("START cannot be an end node")
-        if not self.support_multiple_edges and start_key in set(
+
+        # run this validation only for non-StateGraph graphs
+        if not hasattr(self, "channels") and start_key in set(
             start for start, _ in self.edges
         ):
             raise ValueError(
                 f"Already found path for node '{start_key}'.\n"
-                "For multiple edges, use StateGraph with an annotated state key."
+                "For multiple edges, use StateGraph with an Annotated state key."
             )
 
         self.edges.add((start_key, end_key))

libs/langgraph/langgraph/graph/state.py

@@ -1,3 +1,4 @@
+import inspect
 import logging
 import typing
 import warnings
@@ -196,10 +197,6 @@ class StateGraph(Graph):
                         )
                 else:
                     self.managed[key] = managed
-            if any(
-                isinstance(c, BinaryOperatorAggregate) for c in self.channels.values()
-            ):
-                self.support_multiple_edges = True
 
     @overload
     def add_node(
@@ -338,10 +335,13 @@ class StateGraph(Graph):
                 hints := get_type_hints(action.__call__) or get_type_hints(action)
             ):
                 if input is None:
-                    input_hint = hints[list(hints.keys())[0]]
-                    if isinstance(input_hint, type) and get_type_hints(input_hint):
-                        input = input_hint
-        except TypeError:
+                    first_parameter_name = next(
+                        iter(inspect.signature(action).parameters.keys())
+                    )
+                    if input_hint := hints.get(first_parameter_name):
+                        if isinstance(input_hint, type) and get_type_hints(input_hint):
+                            input = input_hint
+        except (TypeError, StopIteration):
             pass
         if input is not None:
             self._add_schema(input)
@@ -727,10 +727,15 @@ def _get_channel(
         else:
             raise ValueError(f"This {annotation} not allowed in this position")
     elif channel := _is_field_channel(annotation):
+        channel.key = name
         return channel
     elif channel := _is_field_binop(annotation):
+        channel.key = name
         return channel
-    return LastValue(annotation)
+
+    fallback = LastValue(annotation)
+    fallback.key = name
+    return fallback
 
 
 def _is_field_channel(typ: Type[Any]) -> Optional[BaseChannel]:

libs/langgraph/langgraph/pregel/__init__.py

@@ -71,23 +71,14 @@ from langgraph.constants import (
 )
 from langgraph.errors import GraphInterrupt, GraphRecursionError, InvalidUpdateError
 from langgraph.managed.base import ManagedValueSpec
-from langgraph.pregel.algo import (
-    apply_writes,
-    local_read,
-    prepare_next_tasks,
-    should_interrupt,
-)
+from langgraph.pregel.algo import apply_writes, local_read, prepare_next_tasks
 from langgraph.pregel.debug import (
-    map_debug_task_results,
     print_step_checkpoint,
     print_step_tasks,
     print_step_writes,
     tasks_w_writes,
 )
-from langgraph.pregel.io import (
-    map_output_updates,
-    read_channels,
-)
+from langgraph.pregel.io import read_channels
 from langgraph.pregel.loop import AsyncPregelLoop, SyncPregelLoop
 from langgraph.pregel.manager import AsyncChannelsManager, ChannelsManager
 from langgraph.pregel.read import PregelNode
@@ -503,6 +494,7 @@ class Pregel(
                 config=config,
                 metadata=None,
                 created_at=None,
+                parent_config=None,
                 tasks=(),
             )
 
@@ -578,6 +570,7 @@ class Pregel(
                 config=config,
                 metadata=None,
                 created_at=None,
+                parent_config=None,
                 tasks=(),
             )
 
@@ -723,7 +716,7 @@ class Pregel(
         # update channels
         with ChannelsManager(self.channels, checkpoint, config) as (
             channels,
-            managed,
+            _,
         ):
             # create task to run all writers of the chosen node
             writers = self.nodes[as_node].get_writers()
@@ -759,31 +752,6 @@ class Pregel(
                 checkpoint, channels, [task], self.checkpointer.get_next_version
             ), "Can't write to SharedValues from update_state"
             checkpoint = create_checkpoint(checkpoint, channels, step + 1)
-            # check interrupt before
-            if tasks := should_interrupt(
-                checkpoint,
-                self.interrupt_before_nodes,
-                prepare_next_tasks(
-                    checkpoint,
-                    self.nodes,
-                    channels,
-                    managed,
-                    config,
-                    step + 2,
-                    for_execution=False,
-                ),
-            ):
-                for t in tasks:
-                    self.checkpointer.put_writes(
-                        {
-                            "configurable": {
-                                **checkpoint_config["configurable"],
-                                "checkpoint_id": checkpoint["id"],
-                            }
-                        },
-                        [(INTERRUPT, Interrupt("before"))],
-                        t.id,
-                    )
             return self.checkpointer.put(
                 checkpoint_config,
                 checkpoint,
@@ -866,7 +834,7 @@ class Pregel(
         # update channels, acting as the chosen node
         async with AsyncChannelsManager(self.channels, checkpoint, config) as (
             channels,
-            managed,
+            _,
         ):
             # create task to run all writers of the chosen node
             writers = self.nodes[as_node].get_writers()
@@ -902,35 +870,6 @@ class Pregel(
                 checkpoint, channels, [task], self.checkpointer.get_next_version
             ), "Can't write to SharedValues from update_state"
             checkpoint = create_checkpoint(checkpoint, channels, step + 1)
-            # check interrupt before
-            if tasks := should_interrupt(
-                checkpoint,
-                self.interrupt_before_nodes,
-                prepare_next_tasks(
-                    checkpoint,
-                    self.nodes,
-                    channels,
-                    managed,
-                    config,
-                    step + 2,
-                    for_execution=False,
-                ),
-            ):
-                await asyncio.gather(
-                    *(
-                        self.checkpointer.aput_writes(
-                            {
-                                "configurable": {
-                                    **checkpoint_config["configurable"],
-                                    "checkpoint_id": checkpoint["id"],
-                                }
-                            },
-                            [(INTERRUPT, Interrupt("before"))],
-                            t.id,
-                        )
-                        for t in tasks
-                    )
-                )
             return await self.checkpointer.aput(
                 checkpoint_config,
                 checkpoint,
@@ -1111,6 +1050,7 @@ class Pregel(
                 nodes=self.nodes,
                 specs=self.channels,
                 output_keys=output_keys,
+                stream_keys=self.stream_channels_asis,
             ) as loop:
                 # Similarly to Bulk Synchronous Parallel / Pregel model
                 # computation proceeds in steps, while there are channel updates
@@ -1119,7 +1059,6 @@ class Pregel(
                 # with channel updates applied only at the transition between steps
                 while loop.tick(
                     input_keys=self.input_channels,
-                    stream_keys=self.stream_channels_asis,
                     interrupt_before=interrupt_before,
                     interrupt_after=interrupt_after,
                     manager=run_manager,
@@ -1188,26 +1127,17 @@ class Pregel(
                             else:
                                 # save task writes to checkpointer
                                 loop.put_writes(task.id, task.writes)
-                                # yield updates output for the finished task
-                                if "updates" in stream_modes:
-                                    yield from _with_mode(
-                                        "updates",
-                                        isinstance(stream_mode, list),
-                                        map_output_updates(output_keys, [task]),
-                                    )
-                                if "debug" in stream_modes:
-                                    yield from _with_mode(
-                                        "debug",
-                                        isinstance(stream_mode, list),
-                                        map_debug_task_results(
-                                            loop.step,
-                                            [task],
-                                            self.stream_channels_list,
-                                        ),
-                                    )
                         else:
                             # remove references to loop vars
                             del fut, task
+                        # emit output
+                        while loop.stream:
+                            mode, payload = loop.stream.popleft()
+                            if mode in stream_modes:
+                                if isinstance(stream_mode, list):
+                                    yield (mode, payload)
+                                else:
+                                    yield payload
                         if _should_stop_others(done):
                             break
 
@@ -1222,21 +1152,21 @@ class Pregel(
                             [w for t in loop.tasks for w in t.writes],
                             self.stream_channels_list,
                         )
-                # emit output
-                while loop.stream:
-                    mode, payload = loop.stream.popleft()
-                    if mode in stream_modes:
-                        if isinstance(stream_mode, list):
-                            yield (mode, payload)
-                        else:
-                            yield payload
-                # handle exit
-                if loop.status == "out_of_steps":
-                    raise GraphRecursionError(
-                        f"Recursion limit of {config['recursion_limit']} reached "
-                        "without hitting a stop condition. You can increase the "
-                        "limit by setting the `recursion_limit` config key."
-                    )
+            # emit output
+            while loop.stream:
+                mode, payload = loop.stream.popleft()
+                if mode in stream_modes:
+                    if isinstance(stream_mode, list):
+                        yield (mode, payload)
+                    else:
+                        yield payload
+            # handle exit
+            if loop.status == "out_of_steps":
+                raise GraphRecursionError(
+                    f"Recursion limit of {config['recursion_limit']} reached "
+                    "without hitting a stop condition. You can increase the "
+                    "limit by setting the `recursion_limit` config key."
+                )
             # set final channel values as run output
             run_manager.on_chain_end(loop.output)
         except BaseException as e:
@@ -1368,6 +1298,7 @@ class Pregel(
                 nodes=self.nodes,
                 specs=self.channels,
                 output_keys=output_keys,
+                stream_keys=self.stream_channels_asis,
             ) as loop:
                 aioloop = asyncio.get_event_loop()
                 # Similarly to Bulk Synchronous Parallel / Pregel model
@@ -1377,7 +1308,6 @@ class Pregel(
                 # with channel updates applied only at the transition between steps
                 while loop.tick(
                     input_keys=self.input_channels,
-                    stream_keys=self.stream_channels_asis,
                     interrupt_before=interrupt_before,
                     interrupt_after=interrupt_after,
                     manager=run_manager,
@@ -1448,28 +1378,17 @@ class Pregel(
                             else:
                                 # save task writes to checkpointer
                                 loop.put_writes(task.id, task.writes)
-                                # yield updates output for the finished task
-                                if "updates" in stream_modes:
-                                    for chunk in _with_mode(
-                                        "updates",
-                                        isinstance(stream_mode, list),
-                                        map_output_updates(output_keys, [task]),
-                                    ):
-                                        yield chunk
-                                if "debug" in stream_modes:
-                                    for chunk in _with_mode(
-                                        "debug",
-                                        isinstance(stream_mode, list),
-                                        map_debug_task_results(
-                                            loop.step,
-                                            [task],
-                                            self.stream_channels_list,
-                                        ),
-                                    ):
-                                        yield chunk
                         else:
                             # remove references to loop vars
                             del fut, task
+                        # emit output
+                        while loop.stream:
+                            mode, payload = loop.stream.popleft()
+                            if mode in stream_modes:
+                                if isinstance(stream_mode, list):
+                                    yield (mode, payload)
+                                else:
+                                    yield payload
                         if _should_stop_others(done):
                             break
 
@@ -1484,21 +1403,21 @@ class Pregel(
                             [w for t in loop.tasks for w in t.writes],
                             self.stream_channels_list,
                         )
-                # emit output
-                while loop.stream:
-                    mode, payload = loop.stream.popleft()
-                    if mode in stream_modes:
-                        if isinstance(stream_mode, list):
-                            yield (mode, payload)
-                        else:
-                            yield payload
-                # handle exit
-                if loop.status == "out_of_steps":
-                    raise GraphRecursionError(
-                        f"Recursion limit of {config['recursion_limit']} reached "
-                        "without hitting a stop condition. You can increase the "
-                        "limit by setting the `recursion_limit` config key."
-                    )
+            # emit output
+            while loop.stream:
+                mode, payload = loop.stream.popleft()
+                if mode in stream_modes:
+                    if isinstance(stream_mode, list):
+                        yield (mode, payload)
+                    else:
+                        yield payload
+            # handle exit
+            if loop.status == "out_of_steps":
+                raise GraphRecursionError(
+                    f"Recursion limit of {config['recursion_limit']} reached "
+                    "without hitting a stop condition. You can increase the "
+                    "limit by setting the `recursion_limit` config key."
+                )
             # set final channel values as run output
             await run_manager.on_chain_end(loop.output)
         except BaseException as e:

libs/langgraph/langgraph/pregel/algo.py

@@ -198,12 +198,7 @@ def apply_writes(
     updated_channels: set[str] = set()
     for chan, vals in pending_writes_by_channel.items():
         if chan in channels:
-            try:
-                updated = channels[chan].update(vals)
-            except InvalidUpdateError as e:
-                raise InvalidUpdateError(
-                    f"Invalid update for channel {chan} with values {vals}"
-                ) from e
+            updated = channels[chan].update(vals)
             if updated and get_next_version is not None:
                 checkpoint["channel_versions"][chan] = get_next_version(
                     max_version, channels[chan]

libs/langgraph/langgraph/pregel/debug.py

@@ -1,5 +1,6 @@
 import json
 from collections import defaultdict
+from dataclasses import asdict
 from datetime import datetime, timezone
 from pprint import pformat
 from typing import Any, Iterator, Literal, Mapping, Optional, Sequence, TypedDict, Union
@@ -25,6 +26,8 @@ class TaskPayload(TypedDict):
 class TaskResultPayload(TypedDict):
     id: str
     name: str
+    error: Optional[str]
+    interrupts: list[dict]
     result: list[tuple[str, Any]]
 
 
@@ -97,11 +100,14 @@ def map_debug_tasks(
 
 def map_debug_task_results(
     step: int,
-    tasks: list[PregelExecutableTask],
-    stream_channels_list: Sequence[str],
+    tasks: list[tuple[PregelExecutableTask, Sequence[tuple[str, Any]]]],
+    stream_keys: Union[str, Sequence[str]],
 ) -> Iterator[DebugOutputTaskResult]:
+    stream_channels_list = (
+        [stream_keys] if isinstance(stream_keys, str) else stream_keys
+    )
     ts = datetime.now(timezone.utc).isoformat()
-    for name, _, _, writes, config, _, _, _ in tasks:
+    for (name, _, _, _, config, _, _, _), writes in tasks:
         if config is not None and TAG_HIDDEN in config.get("tags", []):
             continue
 
@@ -116,7 +122,9 @@ def map_debug_task_results(
             "payload": {
                 "id": str(uuid5(TASK_NAMESPACE, json.dumps((name, step, metadata)))),
                 "name": name,
+                "error": next((w[1] for w in writes if w[0] == ERROR), None),
                 "result": [w for w in writes if w[0] in stream_channels_list],
+                "interrupts": [asdict(w[1]) for w in writes if w[0] == INTERRUPT],
             },
         }
 
@@ -150,7 +158,7 @@ def map_debug_checkpoint(
                 else {
                     "id": t.id,
                     "name": t.name,
-                    "interrupts": t.interrupts,
+                    "interrupts": tuple(asdict(i) for i in t.interrupts),
                 }
                 for t in tasks_w_writes(tasks, pending_writes)
             ],

libs/langgraph/langgraph/pregel/io.py

@@ -3,7 +3,7 @@ from typing import Any, Iterator, Mapping, Optional, Sequence, TypeVar, Union
 from langchain_core.runnables.utils import AddableDict
 
 from langgraph.channels.base import BaseChannel, EmptyChannelError
-from langgraph.constants import TAG_HIDDEN
+from langgraph.constants import ERROR, INTERRUPT, TAG_HIDDEN
 from langgraph.pregel.log import logger
 from langgraph.pregel.types import PregelExecutableTask
 
@@ -95,19 +95,22 @@ class AddableUpdatesDict(AddableDict):
 
 def map_output_updates(
     output_channels: Union[str, Sequence[str]],
-    tasks: list[PregelExecutableTask],
+    tasks: list[tuple[PregelExecutableTask, Sequence[tuple[str, Any]]]],
 ) -> Iterator[dict[str, Union[Any, dict[str, Any]]]]:
     """Map pending writes (a sequence of tuples (channel, value)) to output chunk."""
     output_tasks = [
-        t for t in tasks if not t.config or TAG_HIDDEN not in t.config.get("tags")
+        (t, ww)
+        for t, ww in tasks
+        if (not t.config or TAG_HIDDEN not in t.config.get("tags"))
+        and all(k not in (ERROR, INTERRUPT) for k, _ in ww)
     ]
     if not output_tasks:
         return
     if isinstance(output_channels, str):
         updated = [
             (task.name, value)
-            for task in output_tasks
-            for chan, value in task.writes
+            for task, writes in output_tasks
+            for chan, value in writes
             if chan == output_channels
         ]
     else:
@@ -116,10 +119,10 @@ def map_output_updates(
                 task.name,
                 {chan: value for chan, value in task.writes if chan in output_channels},
             )
-            for task in output_tasks
-            if any(chan in output_channels for chan, _ in task.writes)
+            for task, writes in output_tasks
+            if any(chan in output_channels for chan, _ in writes)
         ]
-    grouped = {t.name: [] for t in output_tasks}
+    grouped = {t.name: [] for t, _ in output_tasks}
     for node, value in updated:
         grouped[node].append(value)
     for node, value in grouped.items():

libs/langgraph/langgraph/pregel/loop.py

@@ -41,7 +41,6 @@ from langgraph.constants import (
     ERROR,
     INPUT,
     INTERRUPT,
-    Interrupt,
 )
 from langgraph.errors import EmptyInputError, GraphInterrupt
 from langgraph.managed.base import (
@@ -56,7 +55,11 @@ from langgraph.pregel.algo import (
     prepare_next_tasks,
     should_interrupt,
 )
-from langgraph.pregel.debug import map_debug_checkpoint, map_debug_tasks
+from langgraph.pregel.debug import (
+    map_debug_checkpoint,
+    map_debug_task_results,
+    map_debug_tasks,
+)
 from langgraph.pregel.executor import (
     AsyncBackgroundExecutor,
     BackgroundExecutor,
@@ -90,6 +93,7 @@ class PregelLoop:
     nodes: Mapping[str, PregelNode]
     specs: Mapping[str, Union[BaseChannel, ManagedValueSpec]]
     output_keys: Union[str, Sequence[str]]
+    stream_keys: Union[str, Sequence[str]]
     is_nested: bool
 
     checkpointer_get_next_version: Callable[[Optional[V]], V]
@@ -138,6 +142,7 @@ class PregelLoop:
         nodes: Mapping[str, PregelNode],
         specs: Mapping[str, Union[BaseChannel, ManagedValueSpec]],
         output_keys: Union[str, Sequence[str]],
+        stream_keys: Union[str, Sequence[str]],
     ) -> None:
         self.stream = deque()
         self.input = input
@@ -147,6 +152,7 @@ class PregelLoop:
         self.nodes = nodes
         self.specs = specs
         self.output_keys = output_keys
+        self.stream_keys = stream_keys
         self.is_nested = CONFIG_KEY_READ in self.config.get("configurable", {})
 
     def mark_tasks_scheduled(self, tasks: Sequence[PregelExecutableTask]) -> None:
@@ -155,6 +161,8 @@ class PregelLoop:
 
     def put_writes(self, task_id: str, writes: Sequence[tuple[str, Any]]) -> None:
         """Put writes for a task, to be read by the next tick."""
+        if not writes:
+            return
         self.checkpoint_pending_writes.extend((task_id, k, v) for k, v in writes)
         if self.checkpointer_put_writes is not None:
             self.submit(
@@ -172,12 +180,22 @@ class PregelLoop:
                 writes,
                 task_id,
             )
+        if task := next((t for t in self.tasks if t.id == task_id), None):
+            self.stream.extend(
+                ("updates", v)
+                for v in map_output_updates(self.output_keys, [(task, writes)])
+            )
+            self.stream.extend(
+                ("debug", v)
+                for v in map_debug_task_results(
+                    self.step, [(task, writes)], self.stream_keys
+                )
+            )
 
     def tick(
         self,
         *,
         input_keys: Union[str, Sequence[str]],
-        stream_keys: Union[str, Sequence[str]] = EMPTY_SEQ,
         interrupt_after: Sequence[str] = EMPTY_SEQ,
         interrupt_before: Sequence[str] = EMPTY_SEQ,
         manager: Union[None, AsyncParentRunManager, ParentRunManager] = None,
@@ -213,17 +231,18 @@ class PregelLoop:
             self._put_checkpoint(
                 {
                     "source": "loop",
-                    "writes": single(map_output_updates(self.output_keys, self.tasks)),
+                    "writes": single(
+                        map_output_updates(
+                            self.output_keys, [(t, t.writes) for t in self.tasks]
+                        )
+                    ),
                 }
             )
             # after execution, check if we should interrupt
-            if tasks := should_interrupt(self.checkpoint, interrupt_after, self.tasks):
+            if should_interrupt(self.checkpoint, interrupt_after, self.tasks):
                 self.status = "interrupt_after"
-                interrupts = [(t.id, Interrupt("after")) for t in tasks]
-                for tid, interrupt in interrupts:
-                    self.put_writes(tid, [(INTERRUPT, interrupt)])
                 if self.is_nested:
-                    raise GraphInterrupt([i[1] for i in interrupts])
+                    raise GraphInterrupt()
                 else:
                     return False
         else:
@@ -256,7 +275,7 @@ class PregelLoop:
                     self.step - 1,  # printing checkpoint for previous step
                     self.checkpoint_config,
                     self.channels,
-                    stream_keys,
+                    self.stream_keys,
                     self.checkpoint_metadata,
                     self.checkpoint,
                     self.tasks,
@@ -281,20 +300,16 @@ class PregelLoop:
         if all(task.writes for task in self.tasks):
             return self.tick(
                 input_keys=input_keys,
-                stream_keys=stream_keys,
                 interrupt_after=interrupt_after,
                 interrupt_before=interrupt_before,
                 manager=manager,
             )
 
         # before execution, check if we should interrupt
-        if tasks := should_interrupt(self.checkpoint, interrupt_before, self.tasks):
+        if should_interrupt(self.checkpoint, interrupt_before, self.tasks):
             self.status = "interrupt_before"
-            interrupts = [(t.id, Interrupt("before")) for t in tasks]
-            for tid, interrupt in interrupts:
-                self.put_writes(tid, [(INTERRUPT, interrupt)])
             if self.is_nested:
-                raise GraphInterrupt([i[1] for i in interrupts])
+                raise GraphInterrupt()
             else:
                 return False
 
@@ -435,6 +450,7 @@ class SyncPregelLoop(PregelLoop, ContextManager):
         nodes: Mapping[str, PregelNode],
         specs: Mapping[str, Union[BaseChannel, ManagedValueSpec]],
         output_keys: Union[str, Sequence[str]] = EMPTY_SEQ,
+        stream_keys: Union[str, Sequence[str]] = EMPTY_SEQ,
     ) -> None:
         super().__init__(
             input,
@@ -444,6 +460,7 @@ class SyncPregelLoop(PregelLoop, ContextManager):
             nodes=nodes,
             specs=specs,
             output_keys=output_keys,
+            stream_keys=stream_keys,
         )
         self.stack = ExitStack()
         if checkpointer:
@@ -522,6 +539,7 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager):
         nodes: Mapping[str, PregelNode],
         specs: Mapping[str, Union[BaseChannel, ManagedValueSpec]],
         output_keys: Union[str, Sequence[str]] = EMPTY_SEQ,
+        stream_keys: Union[str, Sequence[str]] = EMPTY_SEQ,
     ) -> None:
         super().__init__(
             input,
@@ -531,6 +549,7 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager):
             nodes=nodes,
             specs=specs,
             output_keys=output_keys,
+            stream_keys=stream_keys,
         )
         self.store = AsyncBatchedStore(self.store) if self.store else None
         self.stack = AsyncExitStack()

libs/langgraph/langgraph/pregel/manager.py

@@ -41,7 +41,7 @@ def ChannelsManager(
         yield (
             {
                 k: stack.enter_context(
-                    v.from_checkpoint(checkpoint["channel_values"].get(k), config)
+                    v.from_checkpoint_named(checkpoint["channel_values"].get(k), config)
                 )
                 for k, v in channel_specs.items()
             },
@@ -95,7 +95,9 @@ async def AsyncChannelsManager(
             # channels: enter each channel with checkpoint
             {
                 k: await stack.enter_async_context(
-                    v.afrom_checkpoint(checkpoint["channel_values"].get(k), config)
+                    v.afrom_checkpoint_named(
+                        checkpoint["channel_values"].get(k), config
+                    )
                 )
                 for k, v in channel_specs.items()
             },

libs/langgraph/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "langgraph"
-version = "0.2.9"
+version = "0.2.11"
 description = "Building stateful, multi-actor applications with LLMs"
 authors = []
 license = "MIT"

libs/langgraph/tests/any_str.py

@@ -1,3 +1,6 @@
+from typing import Any, Sequence
+
+
 class AnyStr(str):
     def __init__(self) -> None:
         super().__init__()
@@ -9,6 +12,17 @@ class AnyStr(str):
         return hash(str(self))
 
 
+class AnyVersion:
+    def __init__(self) -> None:
+        super().__init__()
+
+    def __eq__(self, other: object) -> bool:
+        return isinstance(other, (str, int, float))
+
+    def __hash__(self) -> int:
+        return hash(str(self))
+
+
 class ExceptionLike:
     def __init__(self, exc: Exception) -> None:
         self.exc = exc
@@ -22,3 +36,24 @@ class ExceptionLike:
 
     def __hash__(self) -> int:
         return hash((self.exc.__class__, str(self.exc)))
+
+    def __repr__(self) -> str:
+        return str(self.exc)
+
+
+class UnsortedSequence:
+    def __init__(self, *values: Any) -> None:
+        self.seq = values
+
+    def __eq__(self, value: object) -> bool:
+        return (
+            isinstance(value, Sequence)
+            and len(self.seq) == len(value)
+            and all(a in value for a in self.seq)
+        )
+
+    def __hash__(self) -> int:
+        return hash(frozenset(self.seq))
+
+    def __repr__(self) -> str:
+        return repr(self.seq)

libs/langgraph/tests/test_pregel.py

@@ -72,7 +72,7 @@ from langgraph.pregel import (
 from langgraph.pregel.retry import RetryPolicy
 from langgraph.pregel.types import PregelTask
 from langgraph.store.memory import MemoryStore
-from tests.any_str import AnyStr, ExceptionLike
+from tests.any_str import AnyStr, AnyVersion, ExceptionLike, UnsortedSequence
 from tests.fake_tracer import FakeTracer
 from tests.memory_assert import (
     MemorySaverAssertCheckpointMetadata,
@@ -203,6 +203,21 @@ def test_graph_validation() -> None:
     with pytest.raises(ValueError, match="Invalid reducer"):
         StateGraph(BadReducerState)
 
+    def node_b(state: State) -> State:
+        return {"hello": "world"}
+
+    builder = StateGraph(State)
+    builder.add_node("a", node_b)
+    builder.add_node("b", node_b)
+    builder.add_node("c", node_b)
+    builder.set_entry_point("a")
+    builder.add_edge("a", "b")
+    builder.add_edge("a", "c")
+    graph = builder.compile()
+
+    with pytest.raises(InvalidUpdateError, match="At key 'hello'"):
+        graph.invoke({"hello": "there"})
+
 
 def test_checkpoint_errors() -> None:
     class FaultyGetCheckpointer(MemorySaver):
@@ -774,7 +789,7 @@ def test_invoke_two_processes_in_out_interrupt(
         ),
     ]
 
-    # forking from any previous checkpoint w/out forking should do nothing
+    # re-running from any previous checkpoint w/out forking should do nothing
     assert [c for c in app.stream(None, history[0].config, stream_mode="updates")] == []
     assert [c for c in app.stream(None, history[1].config, stream_mode="updates")] == []
     assert [c for c in app.stream(None, history[2].config, stream_mode="updates")] == []
@@ -1038,6 +1053,8 @@ def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None:
                 "id": "2687f72c-e3a8-5f6f-9afa-047cbf24e923",
                 "name": "one",
                 "result": [("inbox", 3)],
+                "error": None,
+                "interrupts": [],
             },
         },
         {
@@ -1048,6 +1065,8 @@ def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None:
                 "id": "18f52f6a-828d-58a1-a501-53cc0c7af33e",
                 "name": "two",
                 "result": [("output", 13)],
+                "error": None,
+                "interrupts": [],
             },
         },
         {
@@ -1069,6 +1088,8 @@ def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None:
                 "id": "871d6e74-7bb3-565f-a4fe-cef4b8f19b62",
                 "name": "two",
                 "result": [("output", 4)],
+                "error": None,
+                "interrupts": [],
             },
         },
     ]
@@ -1196,6 +1217,21 @@ def test_invoke_two_processes_two_in_two_out_invalid(mocker: MockerFixture) -> N
         # LastValue channels can only be updated once per iteration
         app.invoke(2)
 
+    class State(TypedDict):
+        hello: str
+
+    def my_node(input: State) -> State:
+        return {"hello": "world"}
+
+    builder = StateGraph(State)
+    builder.add_node("one", my_node)
+    builder.add_node("two", my_node)
+    builder.set_conditional_entry_point(lambda _: ["one", "two"])
+
+    graph = builder.compile()
+    with pytest.raises(InvalidUpdateError, match="At key 'hello'"):
+        graph.invoke({"hello": "there"}, debug=True)
+
 
 def test_invoke_two_processes_two_in_two_out_valid(mocker: MockerFixture) -> None:
     add_one = mocker.Mock(side_effect=lambda x: x + 1)
@@ -1380,6 +1416,147 @@ def test_pending_writes_resume(
     # both the pending write and the new write were applied, 1 + 2 + 3 = 6
     assert graph.invoke(None, thread1) == {"value": 6}
 
+    # check all final checkpoints
+    checkpoints = [c for c in checkpointer.list(thread1)]
+    # we should have 3
+    assert len(checkpoints) == 3
+    # the last one not too interesting for this test
+    assert checkpoints[0] == CheckpointTuple(
+        config={
+            "configurable": {
+                "thread_id": "1",
+                "checkpoint_ns": "",
+                "checkpoint_id": AnyStr(),
+            }
+        },
+        checkpoint={
+            "v": 1,
+            "id": AnyStr(),
+            "ts": AnyStr(),
+            "current_tasks": {},
+            "pending_sends": [],
+            "versions_seen": {
+                "one": {
+                    "start:one": AnyVersion(),
+                },
+                "two": {
+                    "start:two": AnyVersion(),
+                },
+                "__input__": {},
+                "__start__": {
+                    "__start__": AnyVersion(),
+                },
+                "__interrupt__": {
+                    "value": AnyVersion(),
+                    "__start__": AnyVersion(),
+                    "start:one": AnyVersion(),
+                    "start:two": AnyVersion(),
+                },
+            },
+            "channel_versions": {
+                "one": AnyVersion(),
+                "two": AnyVersion(),
+                "value": AnyVersion(),
+                "__start__": AnyVersion(),
+                "start:one": AnyVersion(),
+                "start:two": AnyVersion(),
+            },
+            "channel_values": {"one": "one", "two": "two", "value": 6},
+        },
+        metadata={
+            "step": 1,
+            "source": "loop",
+            "writes": {"one": {"value": 2}, "two": {"value": 3}},
+        },
+        parent_config={
+            "configurable": {
+                "thread_id": "1",
+                "checkpoint_ns": "",
+                "checkpoint_id": checkpoints[1].config["configurable"]["checkpoint_id"],
+            }
+        },
+        pending_writes=[],
+    )
+    # the previous one we assert that pending writes contains both
+    # - original error
+    # - successful writes from resuming after preventing error
+    assert checkpoints[1] == CheckpointTuple(
+        config={
+            "configurable": {
+                "thread_id": "1",
+                "checkpoint_ns": "",
+                "checkpoint_id": AnyStr(),
+            }
+        },
+        checkpoint={
+            "v": 1,
+            "id": AnyStr(),
+            "ts": AnyStr(),
+            "current_tasks": {},
+            "pending_sends": [],
+            "versions_seen": {
+                "__input__": {},
+                "__start__": {
+                    "__start__": AnyVersion(),
+                },
+            },
+            "channel_versions": {
+                "value": AnyVersion(),
+                "__start__": AnyVersion(),
+                "start:one": AnyVersion(),
+                "start:two": AnyVersion(),
+            },
+            "channel_values": {
+                "value": 1,
+                "start:one": "__start__",
+                "start:two": "__start__",
+            },
+        },
+        metadata={"step": 0, "source": "loop", "writes": None},
+        parent_config={
+            "configurable": {
+                "thread_id": "1",
+                "checkpoint_ns": "",
+                "checkpoint_id": checkpoints[2].config["configurable"]["checkpoint_id"],
+            }
+        },
+        pending_writes=UnsortedSequence(
+            (AnyStr(), "one", "one"),
+            (AnyStr(), "value", 2),
+            (AnyStr(), "__error__", ExceptionLike(ConnectionError("I'm not good"))),
+            (AnyStr(), "two", "two"),
+            (AnyStr(), "value", 3),
+        ),
+    )
+    assert checkpoints[2] == CheckpointTuple(
+        config={
+            "configurable": {
+                "thread_id": "1",
+                "checkpoint_ns": "",
+                "checkpoint_id": AnyStr(),
+            }
+        },
+        checkpoint={
+            "v": 1,
+            "id": AnyStr(),
+            "ts": AnyStr(),
+            "current_tasks": {},
+            "pending_sends": [],
+            "versions_seen": {"__input__": {}},
+            "channel_versions": {
+                "__start__": AnyVersion(),
+            },
+            "channel_values": {"__start__": {"value": 1}},
+        },
+        metadata={"step": -1, "source": "input", "writes": {"value": 1}},
+        parent_config=None,
+        pending_writes=UnsortedSequence(
+            (AnyStr(), "value", 1),
+            (AnyStr(), "start:one", "__start__"),
+            (AnyStr(), "start:two", "__start__"),
+        ),
+    )
+
 
 def test_cond_edge_after_send() -> None:
     class Node:
@@ -1749,7 +1926,6 @@ def test_channel_enter_exit_timing(mocker: MockerFixture) -> None:
     assert cleanup.call_count == 0
     for i, chunk in enumerate(app.stream(2)):
         assert setup.call_count == 1, "Expected setup to be called once"
-        assert cleanup.call_count == 0, "Expected cleanup to not be called yet"
         if i == 0:
             assert chunk == {"inbox": [3]}
         elif i == 1:
@@ -2205,7 +2381,7 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None:
                 ),
             },
         },
-        tasks=(PregelTask(AnyStr(), "tools", interrupts=(Interrupt("before"),)),),
+        tasks=(PregelTask(AnyStr(), "tools"),),
         next=("tools",),
         config=app_w_interrupt.checkpointer.get_tuple(config).config,
         created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"],
@@ -2251,7 +2427,7 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None:
                 "input": "what is weather in sf",
             },
         },
-        tasks=(PregelTask(AnyStr(), "tools", interrupts=(Interrupt("before"),)),),
+        tasks=(PregelTask(AnyStr(), "tools"),),
         next=("tools",),
         config=app_w_interrupt.checkpointer.get_tuple(config).config,
         created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"],
@@ -2412,7 +2588,7 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None:
                 ),
             },
         },
-        tasks=(PregelTask(AnyStr(), "tools", interrupts=(Interrupt("before"),)),),
+        tasks=(PregelTask(AnyStr(), "tools"),),
         next=("tools",),
         config=app_w_interrupt.checkpointer.get_tuple(config).config,
         created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"],
@@ -3013,7 +3189,7 @@ def test_conditional_state_graph(
             ),
             "intermediate_steps": [],
         },
-        tasks=(PregelTask(AnyStr(), "tools", interrupts=(Interrupt("before"),)),),
+        tasks=(PregelTask(AnyStr(), "tools"),),
         next=("tools",),
         config=app_w_interrupt.checkpointer.get_tuple(config).config,
         created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"],
@@ -3053,7 +3229,7 @@ def test_conditional_state_graph(
             ),
             "intermediate_steps": [],
         },
-        tasks=(PregelTask(AnyStr(), "tools", interrupts=(Interrupt("before"),)),),
+        tasks=(PregelTask(AnyStr(), "tools"),),
         next=("tools",),
         config=app_w_interrupt.checkpointer.get_tuple(config).config,
         created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"],
@@ -3162,7 +3338,7 @@ def test_conditional_state_graph(
         values={
             "intermediate_steps": [],
         },
-        tasks=(PregelTask(AnyStr(), "agent", interrupts=(Interrupt("before"),)),),
+        tasks=(PregelTask(AnyStr(), "agent"),),
         next=("agent",),
         config=app_w_interrupt.checkpointer.get_tuple(config).config,
         created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"],
@@ -3187,7 +3363,7 @@ def test_conditional_state_graph(
             ),
             "intermediate_steps": [],
         },
-        tasks=(PregelTask(AnyStr(), "tools", interrupts=(Interrupt("before"),)),),
+        tasks=(PregelTask(AnyStr(), "tools"),),
         next=("tools",),
         config=app_w_interrupt.checkpointer.get_tuple(config).config,
         created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"],
@@ -3240,7 +3416,7 @@ def test_conditional_state_graph(
                 )
             ],
         },
-        tasks=(PregelTask(AnyStr(), "agent", interrupts=(Interrupt("before"),)),),
+        tasks=(PregelTask(AnyStr(), "agent"),),
         next=("agent",),
         config=app_w_interrupt.checkpointer.get_tuple(config).config,
         created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"],
@@ -4842,13 +5018,7 @@ def test_message_graph(
                 id="ai1",
             ),
         ],
-        tasks=(
-            PregelTask(
-                AnyStr(),
-                "tools",
-                interrupts=(Interrupt("before"),),
-            ),
-        ),
+        tasks=(PregelTask(AnyStr(), "tools"),),
         next=("tools",),
         config=app_w_interrupt.checkpointer.get_tuple(config).config,
         created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"],
@@ -4893,7 +5063,7 @@ def test_message_graph(
                 ],
             ),
         ],
-        tasks=(PregelTask(AnyStr(), "tools", interrupts=(Interrupt("before"),)),),
+        tasks=(PregelTask(AnyStr(), "tools"),),
         next=("tools",),
         config=app_w_interrupt.checkpointer.get_tuple(config).config,
         created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"],
@@ -4975,7 +5145,7 @@ def test_message_graph(
                 id="ai2",
             ),
         ],
-        tasks=(PregelTask(AnyStr(), "tools", interrupts=(Interrupt("before"),)),),
+        tasks=(PregelTask(AnyStr(), "tools"),),
         next=("tools",),
         config=app_w_interrupt.checkpointer.get_tuple(config).config,
         created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"],
@@ -5567,13 +5737,7 @@ def test_root_graph(
                 id="ai1",
             ),
         ],
-        tasks=(
-            PregelTask(
-                AnyStr(),
-                "tools",
-                interrupts=(Interrupt("before"),),
-            ),
-        ),
+        tasks=(PregelTask(AnyStr(), "tools"),),
         next=("tools",),
         config=app_w_interrupt.checkpointer.get_tuple(config).config,
         created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"],
@@ -5618,7 +5782,7 @@ def test_root_graph(
                 ],
             ),
         ],
-        tasks=(PregelTask(AnyStr(), "tools", interrupts=(Interrupt("before"),)),),
+        tasks=(PregelTask(AnyStr(), "tools"),),
         next=("tools",),
         config=app_w_interrupt.checkpointer.get_tuple(config).config,
         created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"],
@@ -5700,7 +5864,7 @@ def test_root_graph(
                 id="ai2",
             ),
         ],
-        tasks=(PregelTask(AnyStr(), "tools", interrupts=(Interrupt("before"),)),),
+        tasks=(PregelTask(AnyStr(), "tools"),),
         next=("tools",),
         config=app_w_interrupt.checkpointer.get_tuple(config).config,
         created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"],
@@ -6022,6 +6186,8 @@ def test_in_one_fan_out_out_one_graph_state() -> None:
                     "id": "592f3430-c17c-5d1c-831f-fecebb2c05bf",
                     "name": "rewrite_query",
                     "result": [("query", "query: what is weather in sf")],
+                    "error": None,
+                    "interrupts": [],
                 },
             },
         ),
@@ -6076,6 +6242,8 @@ def test_in_one_fan_out_out_one_graph_state() -> None:
                     "id": "96965ed0-2c10-52a1-86eb-081ba6de73b2",
                     "name": "retriever_two",
                     "result": [("docs", ["doc3", "doc4"])],
+                    "error": None,
+                    "interrupts": [],
                 },
             },
         ),
@@ -6093,6 +6261,8 @@ def test_in_one_fan_out_out_one_graph_state() -> None:
                     "id": "7db5e9d8-e132-5079-ab99-ced15e67d48b",
                     "name": "retriever_one",
                     "result": [("docs", ["doc1", "doc2"])],
+                    "error": None,
+                    "interrupts": [],
                 },
             },
         ),
@@ -6132,6 +6302,8 @@ def test_in_one_fan_out_out_one_graph_state() -> None:
                     "id": "8959fb57-d0f5-5725-9ac4-ec1c554fb0a0",
                     "name": "qa",
                     "result": [("answer", "doc1,doc2,doc3,doc4")],
+                    "error": None,
+                    "interrupts": [],
                 },
             },
         ),
@@ -6302,13 +6474,7 @@ def test_start_branch_then(snapshot: SnapshotAssertion) -> None:
         ]
         assert tool_two.get_state(thread1) == StateSnapshot(
             values={"my_key": "value ⛰️", "market": "DE"},
-            tasks=(
-                PregelTask(
-                    AnyStr(),
-                    "tool_two_slow",
-                    interrupts=(Interrupt("before"),),
-                ),
-            ),
+            tasks=(PregelTask(AnyStr(), "tool_two_slow"),),
             next=("tool_two_slow",),
             config=tool_two.checkpointer.get_tuple(thread1).config,
             created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"],
@@ -6342,13 +6508,7 @@ def test_start_branch_then(snapshot: SnapshotAssertion) -> None:
         }
         assert tool_two.get_state(thread2) == StateSnapshot(
             values={"my_key": "value", "market": "US"},
-            tasks=(
-                PregelTask(
-                    AnyStr(),
-                    "tool_two_fast",
-                    interrupts=(Interrupt("before"),),
-                ),
-            ),
+            tasks=(PregelTask(AnyStr(), "tool_two_fast"),),
             next=("tool_two_fast",),
             config=tool_two.checkpointer.get_tuple(thread2).config,
             created_at=tool_two.checkpointer.get_tuple(thread2).checkpoint["ts"],
@@ -6382,13 +6542,7 @@ def test_start_branch_then(snapshot: SnapshotAssertion) -> None:
         }
         assert tool_two.get_state(thread3) == StateSnapshot(
             values={"my_key": "value", "market": "US"},
-            tasks=(
-                PregelTask(
-                    AnyStr(),
-                    "tool_two_fast",
-                    interrupts=(Interrupt("before"),),
-                ),
-            ),
+            tasks=(PregelTask(AnyStr(), "tool_two_fast"),),
             next=("tool_two_fast",),
             config=tool_two.checkpointer.get_tuple(thread3).config,
             created_at=tool_two.checkpointer.get_tuple(thread3).checkpoint["ts"],
@@ -6399,13 +6553,7 @@ def test_start_branch_then(snapshot: SnapshotAssertion) -> None:
         tool_two.update_state(thread3, {"my_key": "key"})  # appends to my_key
         assert tool_two.get_state(thread3) == StateSnapshot(
             values={"my_key": "valuekey", "market": "US"},
-            tasks=(
-                PregelTask(
-                    AnyStr(),
-                    "tool_two_fast",
-                    interrupts=(Interrupt("before"),),
-                ),
-            ),
+            tasks=(PregelTask(AnyStr(), "tool_two_fast"),),
             next=("tool_two_fast",),
             config=tool_two.checkpointer.get_tuple(thread3).config,
             created_at=tool_two.checkpointer.get_tuple(thread3).checkpoint["ts"],
@@ -6549,6 +6697,8 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None:
                     "id": "7b7b0713-e958-5d07-803c-c9910a7cc162",
                     "name": "prepare",
                     "result": [("my_key", " prepared")],
+                    "error": None,
+                    "interrupts": [],
                 },
             },
             {
@@ -6601,6 +6751,8 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None:
                     "id": "dd9f2fa5-ccfa-5d12-81ec-942563056a08",
                     "name": "tool_two_slow",
                     "result": [("my_key", " slow")],
+                    "error": None,
+                    "interrupts": [],
                 },
             },
             {
@@ -6651,6 +6803,8 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None:
                     "id": "9b590c54-15ef-54b1-83a7-140d27b0bc52",
                     "name": "finish",
                     "result": [("my_key", " finished")],
+                    "error": None,
+                    "interrupts": [],
                 },
             },
             {
@@ -6700,13 +6854,7 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None:
         }
         assert tool_two.get_state(thread1) == StateSnapshot(
             values={"my_key": "value prepared", "market": "DE"},
-            tasks=(
-                PregelTask(
-                    AnyStr(),
-                    "tool_two_slow",
-                    interrupts=(Interrupt("before"),),
-                ),
-            ),
+            tasks=(PregelTask(AnyStr(), "tool_two_slow"),),
             next=("tool_two_slow",),
             config=tool_two.checkpointer.get_tuple(thread1).config,
             created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"],
@@ -6744,13 +6892,7 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None:
         }
         assert tool_two.get_state(thread2) == StateSnapshot(
             values={"my_key": "value prepared", "market": "US"},
-            tasks=(
-                PregelTask(
-                    AnyStr(),
-                    "tool_two_fast",
-                    interrupts=(Interrupt("before"),),
-                ),
-            ),
+            tasks=(PregelTask(AnyStr(), "tool_two_fast"),),
             next=("tool_two_fast",),
             config=tool_two.checkpointer.get_tuple(thread2).config,
             created_at=tool_two.checkpointer.get_tuple(thread2).checkpoint["ts"],
@@ -6797,7 +6939,7 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None:
                 "my_key": "value prepared slow",
                 "market": "DE",
             },
-            tasks=(PregelTask(AnyStr(), "finish", interrupts=(Interrupt("before"),)),),
+            tasks=(PregelTask(AnyStr(), "finish"),),
             next=("finish",),
             config=tool_two.checkpointer.get_tuple(thread1).config,
             created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"],
@@ -6816,7 +6958,7 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None:
                 "my_key": "value prepared slower",
                 "market": "DE",
             },
-            tasks=(PregelTask(AnyStr(), "finish", interrupts=(Interrupt("before"),)),),
+            tasks=(PregelTask(AnyStr(), "finish"),),
             next=("finish",),
             config=tool_two.checkpointer.get_tuple(thread1).config,
             created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"],
@@ -7073,7 +7215,7 @@ def test_in_one_fan_out_state_graph_waiting_edge(snapshot: SnapshotAssertion) ->
             "query": "analyzed: query: what is weather in sf",
             "docs": ["doc1", "doc2", "doc3", "doc4", "doc5"],
         },
-        tasks=(PregelTask(AnyStr(), "qa", interrupts=(Interrupt("before"),)),),
+        tasks=(PregelTask(AnyStr(), "qa"),),
         next=("qa",),
         config=app_w_interrupt.checkpointer.get_tuple(config).config,
         created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"],

libs/langgraph/tests/test_pregel_async.py

@@ -67,7 +67,7 @@ from langgraph.pregel import (
 from langgraph.pregel.retry import RetryPolicy
 from langgraph.pregel.types import PregelTask
 from langgraph.store.memory import MemoryStore
-from tests.any_str import AnyStr, ExceptionLike
+from tests.any_str import AnyStr, AnyVersion, ExceptionLike, UnsortedSequence
 from tests.fake_tracer import FakeTracer
 from tests.memory_assert import (
     MemorySaverAssertCheckpointMetadata,
@@ -212,7 +212,15 @@ async def test_node_cancellation_on_other_node_exception() -> None:
     assert inner_task_cancelled
 
 
-async def test_dynamic_interrupt(snapshot: SnapshotAssertion) -> None:
+@pytest.mark.parametrize(
+    "checkpointer_name",
+    ["memory", "sqlite_aio", "postgres_aio", "postgres_aio_pipe"],
+)
+async def test_dynamic_interrupt(
+    checkpointer_name: str, snapshot: SnapshotAssertion, request: pytest.FixtureRequest
+) -> None:
+    checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}")
+
     class State(TypedDict):
         my_key: Annotated[str, operator.add]
         market: str
@@ -250,51 +258,49 @@ async def test_dynamic_interrupt(snapshot: SnapshotAssertion) -> None:
         "market": "US",
     }
 
-    async with AsyncSqliteSaver.from_conn_string(":memory:") as saver:
-        tool_two = tool_two_graph.compile(checkpointer=saver)
+    tool_two = tool_two_graph.compile(checkpointer=checkpointer)
 
-        # missing thread_id
-        with pytest.raises(ValueError, match="thread_id"):
-            await tool_two.ainvoke({"my_key": "value", "market": "DE"})
+    # missing thread_id
+    with pytest.raises(ValueError, match="thread_id"):
+        await tool_two.ainvoke({"my_key": "value", "market": "DE"})
 
-        thread1 = {"configurable": {"thread_id": "1"}}
-        # stop when about to enter node
-        assert await tool_two.ainvoke(
-            {"my_key": "value ⛰️", "market": "DE"}, thread1
-        ) == {
-            "my_key": "value ⛰️",
-            "market": "DE",
-        }
-        assert [c.metadata async for c in tool_two.checkpointer.alist(thread1)] == [
-            {
-                "source": "loop",
-                "step": 0,
-                "writes": None,
-            },
-            {
-                "source": "input",
-                "step": -1,
-                "writes": {"my_key": "value ⛰️", "market": "DE"},
-            },
-        ]
-        tup = await tool_two.checkpointer.aget_tuple(thread1)
-        assert await tool_two.aget_state(thread1) == StateSnapshot(
-            values={"my_key": "value ⛰️", "market": "DE"},
-            next=("tool_two",),
-            tasks=(
-                PregelTask(
-                    AnyStr(),
-                    "tool_two",
-                    interrupts=(Interrupt("during", "Just because..."),),
-                ),
+    thread1 = {"configurable": {"thread_id": "1"}}
+    # stop when about to enter node
+    assert await tool_two.ainvoke({"my_key": "value ⛰️", "market": "DE"}, thread1) == {
+        "my_key": "value ⛰️",
+        "market": "DE",
+    }
+    assert [c.metadata async for c in tool_two.checkpointer.alist(thread1)] == [
+        {
+            "source": "loop",
+            "step": 0,
+            "writes": None,
+        },
+        {
+            "source": "input",
+            "step": -1,
+            "writes": {"my_key": "value ⛰️", "market": "DE"},
+        },
+    ]
+    tup = await tool_two.checkpointer.aget_tuple(thread1)
+    assert await tool_two.aget_state(thread1) == StateSnapshot(
+        values={"my_key": "value ⛰️", "market": "DE"},
+        next=("tool_two",),
+        tasks=(
+            PregelTask(
+                AnyStr(),
+                "tool_two",
+                interrupts=(Interrupt("during", "Just because..."),),
             ),
-            config=tup.config,
-            created_at=tup.checkpoint["ts"],
-            metadata={"source": "loop", "step": 0, "writes": None},
-            parent_config=[
-                c async for c in tool_two.checkpointer.alist(thread1, limit=2)
-            ][-1].config,
-        )
+        ),
+        config=tup.config,
+        created_at=tup.checkpoint["ts"],
+        metadata={"source": "loop", "step": 0, "writes": None},
+        parent_config=[c async for c in tool_two.checkpointer.alist(thread1, limit=2)][
+            -1
+        ].config,
+    )
+    # TODO use aget_state_history
 
 
 @pytest.mark.parametrize(
@@ -1283,6 +1289,8 @@ async def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None:
                 "id": "2687f72c-e3a8-5f6f-9afa-047cbf24e923",
                 "name": "one",
                 "result": [("inbox", 3)],
+                "error": None,
+                "interrupts": [],
             },
         },
         {
@@ -1293,6 +1301,8 @@ async def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None:
                 "id": "18f52f6a-828d-58a1-a501-53cc0c7af33e",
                 "name": "two",
                 "result": [("output", 13)],
+                "error": None,
+                "interrupts": [],
             },
         },
         {
@@ -1314,6 +1324,8 @@ async def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None:
                 "id": "871d6e74-7bb3-565f-a4fe-cef4b8f19b62",
                 "name": "two",
                 "result": [("output", 4)],
+                "error": None,
+                "interrupts": [],
             },
         },
     ]
@@ -1608,6 +1620,10 @@ async def test_pending_writes_resume(
     error_write = next(w for w in checkpoint.pending_writes if w[1] == ERROR)
     assert error_write[0] != non_error_writes[0][0]
 
+    # TODO arguably this shouldn't even run the failed task again,
+    # and should require empty update_state (ie new checkpoint_id)
+    # in order to try again
+
     # resume execution
     with pytest.raises(ValueError, match="I'm not good"):
         await graph.ainvoke(None, thread1)
@@ -1626,6 +1642,147 @@ async def test_pending_writes_resume(
     # both the pending write and the new write were applied, 1 + 2 + 3 = 6
     assert await graph.ainvoke(None, thread1) == {"value": 6}
 
+    # check all final checkpoints
+    checkpoints = [c async for c in checkpointer.alist(thread1)]
+    # we should have 3
+    assert len(checkpoints) == 3
+    # the last one not too interesting for this test
+    assert checkpoints[0] == CheckpointTuple(
+        config={
+            "configurable": {
+                "thread_id": "1",
+                "checkpoint_ns": "",
+                "checkpoint_id": AnyStr(),
+            }
+        },
+        checkpoint={
+            "v": 1,
+            "id": AnyStr(),
+            "ts": AnyStr(),
+            "current_tasks": {},
+            "pending_sends": [],
+            "versions_seen": {
+                "one": {
+                    "start:one": AnyVersion(),
+                },
+                "two": {
+                    "start:two": AnyVersion(),
+                },
+                "__input__": {},
+                "__start__": {
+                    "__start__": AnyVersion(),
+                },
+                "__interrupt__": {
+                    "value": AnyVersion(),
+                    "__start__": AnyVersion(),
+                    "start:one": AnyVersion(),
+                    "start:two": AnyVersion(),
+                },
+            },
+            "channel_versions": {
+                "one": AnyVersion(),
+                "two": AnyVersion(),
+                "value": AnyVersion(),
+                "__start__": AnyVersion(),
+                "start:one": AnyVersion(),
+                "start:two": AnyVersion(),
+            },
+            "channel_values": {"one": "one", "two": "two", "value": 6},
+        },
+        metadata={
+            "step": 1,
+            "source": "loop",
+            "writes": {"one": {"value": 2}, "two": {"value": 3}},
+        },
+        parent_config={
+            "configurable": {
+                "thread_id": "1",
+                "checkpoint_ns": "",
+                "checkpoint_id": checkpoints[1].config["configurable"]["checkpoint_id"],
+            }
+        },
+        pending_writes=[],
+    )
+    # the previous one we assert that pending writes contains both
+    # - original error
+    # - successful writes from resuming after preventing error
+    assert checkpoints[1] == CheckpointTuple(
+        config={
+            "configurable": {
+                "thread_id": "1",
+                "checkpoint_ns": "",
+                "checkpoint_id": AnyStr(),
+            }
+        },
+        checkpoint={
+            "v": 1,
+            "id": AnyStr(),
+            "ts": AnyStr(),
+            "current_tasks": {},
+            "pending_sends": [],
+            "versions_seen": {
+                "__input__": {},
+                "__start__": {
+                    "__start__": AnyVersion(),
+                },
+            },
+            "channel_versions": {
+                "value": AnyVersion(),
+                "__start__": AnyVersion(),
+                "start:one": AnyVersion(),
+                "start:two": AnyVersion(),
+            },
+            "channel_values": {
+                "value": 1,
+                "start:one": "__start__",
+                "start:two": "__start__",
+            },
+        },
+        metadata={"step": 0, "source": "loop", "writes": None},
+        parent_config={
+            "configurable": {
+                "thread_id": "1",
+                "checkpoint_ns": "",
+                "checkpoint_id": checkpoints[2].config["configurable"]["checkpoint_id"],
+            }
+        },
+        pending_writes=UnsortedSequence(
+            (AnyStr(), "one", "one"),
+            (AnyStr(), "value", 2),
+            (AnyStr(), "__error__", ExceptionLike(ValueError("I'm not good"))),
+            (AnyStr(), "two", "two"),
+            (AnyStr(), "value", 3),
+        ),
+    )
+    assert checkpoints[2] == CheckpointTuple(
+        config={
+            "configurable": {
+                "thread_id": "1",
+                "checkpoint_ns": "",
+                "checkpoint_id": AnyStr(),
+            }
+        },
+        checkpoint={
+            "v": 1,
+            "id": AnyStr(),
+            "ts": AnyStr(),
+            "current_tasks": {},
+            "pending_sends": [],
+            "versions_seen": {"__input__": {}},
+            "channel_versions": {
+                "__start__": AnyVersion(),
+            },
+            "channel_values": {"__start__": {"value": 1}},
+        },
+        metadata={"step": -1, "source": "input", "writes": {"value": 1}},
+        parent_config=None,
+        pending_writes=UnsortedSequence(
+            (AnyStr(), "value", 1),
+            (AnyStr(), "start:one", "__start__"),
+            (AnyStr(), "start:two", "__start__"),
+        ),
+    )
+
 
 async def test_cond_edge_after_send() -> None:
     class Node:
@@ -1983,7 +2140,6 @@ async def test_channel_enter_exit_timing(mocker: MockerFixture) -> None:
         assert setup_sync.call_count == 0, "Sync context manager should not be used"
         assert cleanup_sync.call_count == 0, "Sync context manager should not be used"
         assert setup_async.call_count == 1, "Expected setup to be called once"
-        assert cleanup_async.call_count == 0, "Expected cleanup to not be called yet"
         if i == 0:
             assert chunk == {"inbox": [3]}
         elif i == 1:
@@ -2516,13 +2672,7 @@ async def test_conditional_graph() -> None:
                 ),
             },
         },
-        tasks=(
-            PregelTask(
-                AnyStr(),
-                "tools",
-                interrupts=(Interrupt("before"),),
-            ),
-        ),
+        tasks=(PregelTask(AnyStr(), "tools"),),
         next=("tools",),
         config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config,
         created_at=(await app_w_interrupt.checkpointer.aget_tuple(config)).checkpoint[
@@ -2572,13 +2722,7 @@ async def test_conditional_graph() -> None:
                 "input": "what is weather in sf",
             },
         },
-        tasks=(
-            PregelTask(
-                AnyStr(),
-                "tools",
-                interrupts=(Interrupt("before"),),
-            ),
-        ),
+        tasks=(PregelTask(AnyStr(), "tools"),),
         next=("tools",),
         config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config,
         created_at=(await app_w_interrupt.checkpointer.aget_tuple(config)).checkpoint[
@@ -2750,13 +2894,7 @@ async def test_conditional_graph() -> None:
                 ),
             },
         },
-        tasks=(
-            PregelTask(
-                AnyStr(),
-                "tools",
-                interrupts=(Interrupt("before"),),
-            ),
-        ),
+        tasks=(PregelTask(AnyStr(), "tools"),),
         next=("tools",),
         config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config,
         created_at=(await app_w_interrupt.checkpointer.aget_tuple(config)).checkpoint[
@@ -3324,13 +3462,7 @@ async def test_conditional_graph_state(mocker: MockerFixture) -> None:
             ),
             "intermediate_steps": [],
         },
-        tasks=(
-            PregelTask(
-                AnyStr(),
-                "tools",
-                interrupts=(Interrupt("before"),),
-            ),
-        ),
+        tasks=(PregelTask(AnyStr(), "tools"),),
         next=("tools",),
         config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config,
         created_at=(await app_w_interrupt.checkpointer.aget_tuple(config)).checkpoint[
@@ -3374,7 +3506,7 @@ async def test_conditional_graph_state(mocker: MockerFixture) -> None:
             ),
             "intermediate_steps": [],
         },
-        tasks=(PregelTask(AnyStr(), "tools", interrupts=(Interrupt("before"),)),),
+        tasks=(PregelTask(AnyStr(), "tools"),),
         next=("tools",),
         config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config,
         created_at=(await app_w_interrupt.checkpointer.aget_tuple(config)).checkpoint[
@@ -4749,6 +4881,8 @@ async def test_in_one_fan_out_out_one_graph_state() -> None:
                     "id": "592f3430-c17c-5d1c-831f-fecebb2c05bf",
                     "name": "rewrite_query",
                     "result": [("query", "query: what is weather in sf")],
+                    "error": None,
+                    "interrupts": [],
                 },
             },
         ),
@@ -4803,6 +4937,8 @@ async def test_in_one_fan_out_out_one_graph_state() -> None:
                     "id": "96965ed0-2c10-52a1-86eb-081ba6de73b2",
                     "name": "retriever_two",
                     "result": [("docs", ["doc3", "doc4"])],
+                    "error": None,
+                    "interrupts": [],
                 },
             },
         ),
@@ -4820,6 +4956,8 @@ async def test_in_one_fan_out_out_one_graph_state() -> None:
                     "id": "7db5e9d8-e132-5079-ab99-ced15e67d48b",
                     "name": "retriever_one",
                     "result": [("docs", ["doc1", "doc2"])],
+                    "error": None,
+                    "interrupts": [],
                 },
             },
         ),
@@ -4859,6 +4997,8 @@ async def test_in_one_fan_out_out_one_graph_state() -> None:
                     "id": "8959fb57-d0f5-5725-9ac4-ec1c554fb0a0",
                     "name": "qa",
                     "result": [("answer", "doc1,doc2,doc3,doc4")],
+                    "error": None,
+                    "interrupts": [],
                 },
             },
         ),
@@ -4949,13 +5089,7 @@ async def test_start_branch_then() -> None:
         ]
         assert await tool_two.aget_state(thread1) == StateSnapshot(
             values={"my_key": "value", "market": "DE"},
-            tasks=(
-                PregelTask(
-                    AnyStr(),
-                    "tool_two_slow",
-                    interrupts=(Interrupt("before"),),
-                ),
-            ),
+            tasks=(PregelTask(AnyStr(), "tool_two_slow"),),
             next=("tool_two_slow",),
             config=(await tool_two.checkpointer.aget_tuple(thread1)).config,
             created_at=(await tool_two.checkpointer.aget_tuple(thread1)).checkpoint[
@@ -4997,13 +5131,7 @@ async def test_start_branch_then() -> None:
         }
         assert await tool_two.aget_state(thread2) == StateSnapshot(
             values={"my_key": "value", "market": "US"},
-            tasks=(
-                PregelTask(
-                    AnyStr(),
-                    "tool_two_fast",
-                    interrupts=(Interrupt("before"),),
-                ),
-            ),
+            tasks=(PregelTask(AnyStr(), "tool_two_fast"),),
             next=("tool_two_fast",),
             config=(await tool_two.checkpointer.aget_tuple(thread2)).config,
             created_at=(await tool_two.checkpointer.aget_tuple(thread2)).checkpoint[
@@ -5045,13 +5173,7 @@ async def test_start_branch_then() -> None:
         }
         assert await tool_two.aget_state(thread3) == StateSnapshot(
             values={"my_key": "value", "market": "US"},
-            tasks=(
-                PregelTask(
-                    AnyStr(),
-                    "tool_two_fast",
-                    interrupts=(Interrupt("before"),),
-                ),
-            ),
+            tasks=(PregelTask(AnyStr(), "tool_two_fast"),),
             next=("tool_two_fast",),
             config=(await tool_two.checkpointer.aget_tuple(thread3)).config,
             created_at=(await tool_two.checkpointer.aget_tuple(thread3)).checkpoint[
@@ -5066,13 +5188,7 @@ async def test_start_branch_then() -> None:
         await tool_two.aupdate_state(thread3, {"my_key": "key"})  # appends to my_key
         assert await tool_two.aget_state(thread3) == StateSnapshot(
             values={"my_key": "valuekey", "market": "US"},
-            tasks=(
-                PregelTask(
-                    AnyStr(),
-                    "tool_two_fast",
-                    interrupts=(Interrupt("before"),),
-                ),
-            ),
+            tasks=(PregelTask(AnyStr(), "tool_two_fast"),),
             next=("tool_two_fast",),
             config=(await tool_two.checkpointer.aget_tuple(thread3)).config,
             created_at=(await tool_two.checkpointer.aget_tuple(thread3)).checkpoint[
@@ -5223,6 +5339,8 @@ async def test_branch_then() -> None:
                     "id": "7b7b0713-e958-5d07-803c-c9910a7cc162",
                     "name": "prepare",
                     "result": [("my_key", " prepared")],
+                    "error": None,
+                    "interrupts": [],
                 },
             },
             {
@@ -5275,6 +5393,8 @@ async def test_branch_then() -> None:
                     "id": "dd9f2fa5-ccfa-5d12-81ec-942563056a08",
                     "name": "tool_two_slow",
                     "result": [("my_key", " slow")],
+                    "error": None,
+                    "interrupts": [],
                 },
             },
             {
@@ -5325,6 +5445,8 @@ async def test_branch_then() -> None:
                     "id": "9b590c54-15ef-54b1-83a7-140d27b0bc52",
                     "name": "finish",
                     "result": [("my_key", " finished")],
+                    "error": None,
+                    "interrupts": [],
                 },
             },
             {
@@ -5368,19 +5490,125 @@ async def test_branch_then() -> None:
 
         thread1 = {"configurable": {"thread_id": "1"}}
         # stop when about to enter node
-        assert await tool_two.ainvoke({"my_key": "value", "market": "DE"}, thread1) == {
-            "my_key": "value prepared",
-            "market": "DE",
-        }
+        assert [
+            c
+            async for c in tool_two.astream(
+                {"my_key": "value", "market": "DE"}, thread1, stream_mode="debug"
+            )
+        ] == [
+            {
+                "type": "checkpoint",
+                "timestamp": AnyStr(),
+                "step": -1,
+                "payload": {
+                    "config": {
+                        "tags": [],
+                        "metadata": {"thread_id": "1"},
+                        "callbacks": None,
+                        "recursion_limit": 25,
+                        "configurable": {
+                            "thread_id": "1",
+                            "checkpoint_ns": "",
+                            "checkpoint_id": AnyStr(),
+                        },
+                    },
+                    "values": {"my_key": ""},
+                    "metadata": {
+                        "source": "input",
+                        "step": -1,
+                        "writes": {"my_key": "value", "market": "DE"},
+                    },
+                    "next": ["__start__"],
+                    "tasks": [{"id": AnyStr(), "name": "__start__", "interrupts": ()}],
+                },
+            },
+            {
+                "type": "checkpoint",
+                "timestamp": AnyStr(),
+                "step": 0,
+                "payload": {
+                    "config": {
+                        "tags": [],
+                        "metadata": {"thread_id": "1"},
+                        "callbacks": None,
+                        "recursion_limit": 25,
+                        "configurable": {
+                            "thread_id": "1",
+                            "checkpoint_ns": "",
+                            "checkpoint_id": AnyStr(),
+                        },
+                    },
+                    "values": {
+                        "my_key": "value",
+                        "market": "DE",
+                    },
+                    "metadata": {
+                        "source": "loop",
+                        "step": 0,
+                        "writes": None,
+                    },
+                    "next": ["prepare"],
+                    "tasks": [{"id": AnyStr(), "name": "prepare", "interrupts": ()}],
+                },
+            },
+            {
+                "type": "task",
+                "timestamp": AnyStr(),
+                "step": 1,
+                "payload": {
+                    "id": "ca572c3b-b805-5fc6-a19e-3d79f52dde70",
+                    "name": "prepare",
+                    "input": {"my_key": "value", "market": "DE"},
+                    "triggers": ["start:prepare"],
+                },
+            },
+            {
+                "type": "task_result",
+                "timestamp": AnyStr(),
+                "step": 1,
+                "payload": {
+                    "id": "ca572c3b-b805-5fc6-a19e-3d79f52dde70",
+                    "name": "prepare",
+                    "result": [("my_key", " prepared")],
+                    "error": None,
+                    "interrupts": [],
+                },
+            },
+            {
+                "type": "checkpoint",
+                "timestamp": AnyStr(),
+                "step": 1,
+                "payload": {
+                    "config": {
+                        "tags": [],
+                        "metadata": {"thread_id": "1"},
+                        "callbacks": None,
+                        "recursion_limit": 25,
+                        "configurable": {
+                            "thread_id": "1",
+                            "checkpoint_ns": "",
+                            "checkpoint_id": AnyStr(),
+                        },
+                    },
+                    "values": {
+                        "my_key": "value prepared",
+                        "market": "DE",
+                    },
+                    "metadata": {
+                        "source": "loop",
+                        "step": 1,
+                        "writes": {"prepare": {"my_key": " prepared"}},
+                    },
+                    "next": ["tool_two_slow"],
+                    "tasks": [
+                        {"id": AnyStr(), "name": "tool_two_slow", "interrupts": ()}
+                    ],
+                },
+            },
+        ]
         assert await tool_two.aget_state(thread1) == StateSnapshot(
             values={"my_key": "value prepared", "market": "DE"},
-            tasks=(
-                PregelTask(
-                    AnyStr(),
-                    "tool_two_slow",
-                    interrupts=(Interrupt("before"),),
-                ),
-            ),
+            tasks=(PregelTask(AnyStr(), "tool_two_slow"),),
             next=("tool_two_slow",),
             config=(await tool_two.checkpointer.aget_tuple(thread1)).config,
             created_at=(await tool_two.checkpointer.aget_tuple(thread1)).checkpoint[
@@ -5426,13 +5654,7 @@ async def test_branch_then() -> None:
         }
         assert await tool_two.aget_state(thread2) == StateSnapshot(
             values={"my_key": "value prepared", "market": "US"},
-            tasks=(
-                PregelTask(
-                    AnyStr(),
-                    "tool_two_fast",
-                    interrupts=(Interrupt("before"),),
-                ),
-            ),
+            tasks=(PregelTask(AnyStr(), "tool_two_fast"),),
             next=("tool_two_fast",),
             config=(await tool_two.checkpointer.aget_tuple(thread2)).config,
             created_at=(await tool_two.checkpointer.aget_tuple(thread2)).checkpoint[

libs/langgraph/tests/test_state.py

@@ -2,10 +2,11 @@ from typing import Annotated as Annotated2
 from typing import Any
 
 import pytest
+from langchain_core.runnables import RunnableConfig
 from pydantic.v1 import BaseModel
 from typing_extensions import Annotated, TypedDict
 
-from langgraph.graph.state import _warn_invalid_state_schema
+from langgraph.graph.state import StateGraph, _warn_invalid_state_schema
 
 
 class State(BaseModel):
@@ -46,3 +47,42 @@ def test_doesnt_warn_valid_schema(schema: Any):
     # Assert the function does not raise a warning
     with pytest.warns(None):
         _warn_invalid_state_schema(schema)
+
+
+def test_state_schema_with_type_hint():
+    class InputState(TypedDict):
+        question: str
+
+    class OutputState(TypedDict):
+        input_state: InputState
+
+    def complete_hint(state: InputState) -> OutputState:
+        return {"input_state": state}
+
+    def miss_first_hint(state, config: RunnableConfig) -> OutputState:
+        return {"input_state": state}
+
+    def only_return_hint(state, config) -> OutputState:
+        return {"input_state": state}
+
+    def miss_all_hint(state, config):
+        return {"input_state": state}
+
+    graph = StateGraph(input=InputState, output=OutputState)
+    actions = [complete_hint, miss_first_hint, only_return_hint, miss_all_hint]
+
+    for action in actions:
+        graph.add_node(action)
+
+    graph.set_entry_point(actions[0].__name__)
+    for i in range(len(actions) - 1):
+        graph.add_edge(actions[i].__name__, actions[i + 1].__name__)
+    graph.set_finish_point(actions[-1].__name__)
+
+    graph = graph.compile()
+
+    input_state = InputState(question="Hello World!")
+    output_state = OutputState(input_state=input_state)
+    for i, c in enumerate(graph.stream(input_state, stream_mode="updates")):
+        node_name = actions[i].__name__
+        assert c[node_name] == output_state