cli: 0.3.3 (#5100)

Let you set "pip_installer": "pip" (or uv) to handle corner cases in install compatibilities

docs/docs/concepts/persistence.md

@@ -473,6 +473,23 @@ If the checkpointer is used with asynchronous graph execution (i.e. executing th
 When checkpointers save the graph state, they need to serialize the channel values in the state. This is done using serializer objects.
 `langgraph_checkpoint` defines [protocol][langgraph.checkpoint.serde.base.SerializerProtocol] for implementing serializers provides a default implementation ([JsonPlusSerializer][langgraph.checkpoint.serde.jsonplus.JsonPlusSerializer]) that handles a wide variety of types, including LangChain and LangGraph primitives, datetimes, enums and more.
 
+#### Serialization with `pickle`
+
+The default serializer, [`JsonPlusSerializer`][langgraph.checkpoint.serde.jsonplus.JsonPlusSerializer], uses ormsgpack and JSON under the hood, which is not suitable for all types of objects.
+
+If you want to fallback to pickle for objects not currently supported by our msgpack encoder (such as Pandas dataframes),
+you can use the `pickle_fallback` argument of the `JsonPlusSerializer`:
+
+```python
+from langgraph.checkpoint.memory import MemorySaver
+from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
+
+# ... Define the graph ...
+graph.compile(
+    checkpointer=MemorySaver(serde=JsonPlusSerializer(pickle_fallback=True))
+)
+```
+
 #### Encryption
 
 Checkpointers can optionally encrypt all persisted state. To enable this, pass an instance of [`EncryptedSerializer`][langgraph.checkpoint.serde.encrypted.EncryptedSerializer] to the `serde` argument of any `BaseCheckpointSaver` implementation. The easiest way to create an encrypted serializer is via [`from_pycryptodome_aes`][langgraph.checkpoint.serde.encrypted.EncryptedSerializer.from_pycryptodome_aes], which reads the AES key from the `LANGGRAPH_AES_KEY` environment variable (or accepts a `key` argument):

libs/checkpoint-postgres/uv.lock

@@ -327,6 +327,7 @@ dev = [
     { name = "dataclasses-json" },
     { name = "mypy" },
     { name = "numpy" },
+    { name = "pandas" },
     { name = "pytest" },
     { name = "pytest-asyncio" },
     { name = "pytest-mock" },

libs/checkpoint-sqlite/uv.lock

@@ -339,6 +339,7 @@ dev = [
     { name = "dataclasses-json" },
     { name = "mypy" },
     { name = "numpy" },
+    { name = "pandas" },
     { name = "pytest" },
     { name = "pytest-asyncio" },
     { name = "pytest-mock" },

libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py

@@ -475,6 +475,7 @@ def _msgpack_default(obj: Any) -> str | ormsgpack.Ext:
             buf = obj.tobytes(order="A")
             meta = (obj.dtype.str, obj.shape, order, buf)
             return ormsgpack.Ext(EXT_NUMPY_ARRAY, _msgpack_enc(meta))
+
     elif isinstance(obj, BaseException):
         return repr(obj)
     else:

libs/checkpoint/pyproject.toml

@@ -30,6 +30,7 @@ dev = [
   "mypy",
   "dataclasses-json",
   "numpy",
+  "pandas",
 ]
 
 [tool.hatch.build.targets.wheel]

libs/checkpoint/tests/test_jsonplus.py

@@ -12,6 +12,7 @@ from zoneinfo import ZoneInfo
 
 import dataclasses_json
 import numpy as np
+import pandas as pd
 import pytest
 from pydantic import BaseModel, SecretStr
 from pydantic.v1 import BaseModel as BaseModelV1
@@ -332,19 +333,139 @@ def test_serde_jsonplus_numpy_array_json_hook(arr: np.ndarray) -> None:
     assert result == arr.tolist()
 
 
-def test_loads_cannot_find() -> None:
-    serde = JsonPlusSerializer()
+@pytest.mark.parametrize(
+    "df",
+    [
+        pd.DataFrame(),
+        pd.DataFrame({"int_col": [1, 2, 3]}),
+        pd.DataFrame({"float_col": [1.1, 2.2, 3.3]}),
+        pd.DataFrame({"str_col": ["a", "b", "c"]}),
+        pd.DataFrame({"bool_col": [True, False, True]}),
+        pd.DataFrame(
+            {
+                "datetime_col": [
+                    datetime(2024, 1, 1),
+                    datetime(2024, 1, 2),
+                    datetime(2024, 1, 3),
+                ]
+            }
+        ),
+        pd.DataFrame(
+            {
+                "int_col": [1, 2, 3],
+                "float_col": [1.1, 2.2, 3.3],
+                "str_col": ["a", "b", "c"],
+            }
+        ),
+        pd.DataFrame(
+            {
+                "int_col": [1, 2, None],
+                "float_col": [1.1, None, 3.3],
+                "str_col": ["a", None, "c"],
+            }
+        ),
+        pd.DataFrame({"cat_col": pd.Categorical(["a", "b", "a", "c"])}),
+        pd.DataFrame(
+            {
+                "int8": pd.array([1, 2, 3], dtype="int8"),
+                "int16": pd.array([10, 20, 30], dtype="int16"),
+                "int32": pd.array([100, 200, 300], dtype="int32"),
+                "int64": pd.array([1000, 2000, 3000], dtype="int64"),
+                "float32": pd.array([1.1, 2.2, 3.3], dtype="float32"),
+                "float64": pd.array([10.1, 20.2, 30.3], dtype="float64"),
+            }
+        ),
+        pd.DataFrame({"value": [1, 2, 3]}, index=["x", "y", "z"]),
+        pd.DataFrame(
+            [[1, 2, 3, 4]],
+            columns=pd.MultiIndex.from_tuples(
+                [("A", "X"), ("A", "Y"), ("B", "X"), ("B", "Y")]
+            ),
+        ),
+        pd.DataFrame(
+            {"value": [1, 2, 3]}, index=pd.date_range("2024-01-01", periods=3, freq="D")
+        ),
+        pd.DataFrame(
+            {
+                "col1": range(1000),
+                "col2": [f"str_{i}" for i in range(1000)],
+                "col3": np.random.rand(1000),
+            }
+        ),
+        pd.DataFrame(
+            {"tz_datetime": pd.date_range("2024-01-01", periods=3, freq="D", tz="UTC")}
+        ),
+        pd.DataFrame({"timedelta": pd.to_timedelta([1, 2, 3], unit="D")}),
+        pd.DataFrame({"period": pd.period_range("2024-01", periods=3, freq="M")}),
+        pd.DataFrame({"interval": pd.interval_range(start=0, end=3, periods=3)}),
+        pd.DataFrame({"unicode": ["Hello 🌍", "Python 🐍", "Data 📊"]}),
+        pd.DataFrame({"mixed": [1, "string", [1, 2, 3], {"key": "value"}]}),
+        pd.DataFrame({"a": [1], "b": ["test"], "c": [3.14]}),
+        pd.DataFrame({"single": [42]}),
+        pd.DataFrame(
+            {
+                "small": [sys.float_info.min, 0, sys.float_info.max],
+                "large_int": [-(2**63), 0, 2**63 - 1],
+            }
+        ),
+        pd.DataFrame({"special_strings": ["", "null", "None", "NaN", "inf", "-inf"]}),
+        pd.DataFrame({"bytes_col": [b"hello", b"world", b"\x00\x01\x02"]}),
+    ],
+)
+def test_serde_jsonplus_pandas_dataframe(df: pd.DataFrame) -> None:
+    serde = JsonPlusSerializer(pickle_fallback=True)
 
-    dumped = (
-        "json",
-        b'{"lc": 2, "type": "constructor", "id": ["tests", "test_jsonplus", "MyPydanticccc"], "method": null, "args": [], "kwargs": {"foo": "foo", "bar": 1}}',
-    )
+    dumped = serde.dumps_typed(df)
+    assert dumped[0] == "pickle"
+    result = serde.loads_typed(dumped)
+    assert result.equals(df)
 
-    assert serde.loads_typed(dumped) is None, "Should return None if cannot find class"
 
-    dumped = (
-        "json",
-        b'{"lc": 2, "type": "constructor", "id": ["tests", "test_jsonpluss", "MyPydantic"], "method": null, "args": [], "kwargs": {"foo": "foo", "bar": 1}}',
-    )
+@pytest.mark.parametrize(
+    "series",
+    [
+        pd.Series([]),
+        pd.Series([1, 2, 3]),
+        pd.Series([1.1, 2.2, 3.3]),
+        pd.Series(["a", "b", "c"]),
+        pd.Series([True, False, True]),
+        pd.Series([datetime(2024, 1, 1), datetime(2024, 1, 2), datetime(2024, 1, 3)]),
+        pd.Series([1, 2, None]),
+        pd.Series([1.1, None, 3.3]),
+        pd.Series(["a", None, "c"]),
+        pd.Series(pd.Categorical(["a", "b", "a", "c"])),
+        pd.Series([1, 2, 3], dtype="int8"),
+        pd.Series([10, 20, 30], dtype="int16"),
+        pd.Series([100, 200, 300], dtype="int32"),
+        pd.Series([1000, 2000, 3000], dtype="int64"),
+        pd.Series([1.1, 2.2, 3.3], dtype="float32"),
+        pd.Series([10.1, 20.2, 30.3], dtype="float64"),
+        pd.Series([1, 2, 3], index=["x", "y", "z"]),
+        pd.Series([1, 2, 3], index=pd.date_range("2024-01-01", periods=3, freq="D")),
+        pd.Series(range(1000)),
+        pd.Series(pd.date_range("2024-01-01", periods=3, freq="D", tz="UTC")),
+        pd.Series(pd.to_timedelta([1, 2, 3], unit="D")),
+        pd.Series(pd.period_range("2024-01", periods=3, freq="M")),
+        pd.Series(pd.interval_range(start=0, end=3, periods=3)),
+        pd.Series(["Hello 🌍", "Python 🐍", "Data 📊"]),
+        pd.Series([1, "string", [1, 2, 3], {"key": "value"}]),
+        pd.Series([42], name="single"),
+        pd.Series([sys.float_info.min, 0, sys.float_info.max]),
+        pd.Series([-(2**63), 0, 2**63 - 1]),
+        pd.Series(["", "null", "None", "NaN", "inf", "-inf"]),
+        pd.Series([b"hello", b"world", b"\x00\x01\x02"]),
+        pd.Series([1, 2, 3], name="named_series"),
+        pd.Series(
+            [10, 20],
+            index=pd.MultiIndex.from_tuples([("a", 1), ("b", 2)], names=["x", "y"]),
+        ),
+    ],
+)
+def test_serde_jsonplus_pandas_series(series: pd.Series) -> None:
+    serde = JsonPlusSerializer(pickle_fallback=True)
+    dumped = serde.dumps_typed(series)
 
-    assert serde.loads_typed(dumped) is None, "Should return None if cannot find module"
+    assert dumped[0] == "pickle"
+    result = serde.loads_typed(dumped)
+
+    assert result.equals(series)

libs/cli/langgraph_cli/config.py

@@ -383,6 +383,14 @@ class Config(TypedDict, total=False):
     Only relevant if Python dependencies are installed via pip. If omitted, default pip settings are used.
     """
 
+    pip_installer: Optional[str]
+    """Optional. Python package installer to use ('auto', 'pip', 'uv').
+    
+    - 'auto' (default): Use uv for supported base images, otherwise pip
+    - 'pip': Force use of pip regardless of base image support
+    - 'uv': Force use of uv (will fail if base image doesn't support it)
+    """
+
     dockerfile_lines: list[str]
     """Optional. Additional Docker instructions that will be appended to your base Dockerfile.
     
@@ -536,6 +544,7 @@ def validate_config(config: Config) -> Config:
         "node_version": node_version,
         "python_version": python_version,
         "pip_config_file": config.get("pip_config_file"),
+        "pip_installer": config.get("pip_installer", "auto"),
         "_INTERNAL_docker_tag": config.get("_INTERNAL_docker_tag"),
         "base_image": config.get("base_image"),
         "image_distro": image_distro,
@@ -600,6 +609,13 @@ def validate_config(config: Config) -> Config:
                 "Must be either 'debian' or 'wolfi'."
             )
 
+    if pip_installer := config.get("pip_installer"):
+        if pip_installer not in ["auto", "pip", "uv"]:
+            raise click.UsageError(
+                f"Invalid pip_installer: '{pip_installer}'. "
+                "Must be 'auto', 'pip', or 'uv'."
+            )
+
     # Validate auth config
     if auth_conf := config.get("auth"):
         if "path" in auth_conf:
@@ -1114,12 +1130,21 @@ def python_config_to_docker(
     base_image: str,
 ) -> tuple[str, dict[str, str]]:
     """Generate a Dockerfile from the configuration."""
-    if _image_supports_uv(base_image):
+    pip_installer = config.get("pip_installer", "auto")
+
+    if pip_installer == "uv":
         install_cmd = "uv pip install --system"
         uv_removal = "RUN uv pip uninstall --system pip setuptools wheel && rm /usr/bin/uv /usr/bin/uvx"
-    else:
+    elif pip_installer == "pip":
         install_cmd = "pip install"
         uv_removal = ""
+    else:
+        if _image_supports_uv(base_image):
+            install_cmd = "uv pip install --system"
+            uv_removal = "RUN uv pip uninstall --system pip setuptools wheel && rm /usr/bin/uv /usr/bin/uvx"
+        else:
+            install_cmd = "pip install"
+            uv_removal = ""
 
     # configure pip
     pip_install = f"PYTHONDONTWRITEBYTECODE=1 {install_cmd} --no-cache-dir -c /api/constraints.txt"

libs/cli/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "langgraph-cli"
-version = "0.3.1"
+version = "0.3.3"
 description = "CLI for interacting with LangGraph API"
 authors = []
 requires-python = ">=3.9"

libs/cli/schemas/schema.json

@@ -134,6 +134,17 @@
               ],
               "description": "Optional. Linux distribution for the base image.\n\nMust be either 'debian' or 'wolfi'. If omitted, defaults to 'debian'.\n"
             },
+            "pip_installer": {
+              "anyOf": [
+                {
+                  "type": "string"
+                },
+                {
+                  "type": "null"
+                }
+              ],
+              "description": "Optional. Python package installer to use ('auto', 'pip', 'uv').\n\n"
+            },
             "store": {
               "anyOf": [
                 {
@@ -287,6 +298,17 @@
               ],
               "description": "Optional. Linux distribution for the base image.\n\nMust be either 'debian' or 'wolfi'. If omitted, defaults to 'debian'.\n"
             },
+            "pip_installer": {
+              "anyOf": [
+                {
+                  "type": "string"
+                },
+                {
+                  "type": "null"
+                }
+              ],
+              "description": "Optional. Python package installer to use ('auto', 'pip', 'uv').\n\n"
+            },
             "store": {
               "anyOf": [
                 {

libs/cli/schemas/schema.v0.json

@@ -134,6 +134,17 @@
               ],
               "description": "Optional. Linux distribution for the base image.\n\nMust be either 'debian' or 'wolfi'. If omitted, defaults to 'debian'.\n"
             },
+            "pip_installer": {
+              "anyOf": [
+                {
+                  "type": "string"
+                },
+                {
+                  "type": "null"
+                }
+              ],
+              "description": "Optional. Python package installer to use ('auto', 'pip', 'uv').\n\n"
+            },
             "store": {
               "anyOf": [
                 {
@@ -287,6 +298,17 @@
               ],
               "description": "Optional. Linux distribution for the base image.\n\nMust be either 'debian' or 'wolfi'. If omitted, defaults to 'debian'.\n"
             },
+            "pip_installer": {
+              "anyOf": [
+                {
+                  "type": "string"
+                },
+                {
+                  "type": "null"
+                }
+              ],
+              "description": "Optional. Python package installer to use ('auto', 'pip', 'uv').\n\n"
+            },
             "store": {
               "anyOf": [
                 {

libs/cli/tests/unit_tests/test_config.py

@@ -1,3 +1,4 @@
+import copy
 import json
 import os
 import pathlib
@@ -40,6 +41,7 @@ def test_validate_config():
         "python_version": "3.11",
         "node_version": None,
         "pip_config_file": None,
+        "pip_installer": "auto",
         "image_distro": "debian",
         "dockerfile_lines": [],
         "env": {},
@@ -61,6 +63,7 @@ def test_validate_config():
         "python_version": "3.12",
         "node_version": None,
         "pip_config_file": "pipconfig.txt",
+        "pip_installer": "auto",
         "image_distro": "debian",
         "dockerfile_lines": ["ARG meow"],
         "dependencies": [".", "langchain"],
@@ -216,6 +219,74 @@ def test_validate_config_image_distro():
     assert config["image_distro"] == "debian"
 
 
+def test_validate_config_pip_installer():
+    """Test validation of pip_installer field."""
+    # Valid pip_installer values should work
+    config = validate_config(
+        {
+            "python_version": "3.11",
+            "dependencies": ["."],
+            "graphs": {"agent": "./agent.py:graph"},
+            "pip_installer": "auto",
+        }
+    )
+    assert config["pip_installer"] == "auto"
+
+    config = validate_config(
+        {
+            "python_version": "3.11",
+            "dependencies": ["."],
+            "graphs": {"agent": "./agent.py:graph"},
+            "pip_installer": "pip",
+        }
+    )
+    assert config["pip_installer"] == "pip"
+
+    config = validate_config(
+        {
+            "python_version": "3.11",
+            "dependencies": ["."],
+            "graphs": {"agent": "./agent.py:graph"},
+            "pip_installer": "uv",
+        }
+    )
+    assert config["pip_installer"] == "uv"
+
+    # Missing pip_installer should default to "auto"
+    config = validate_config(
+        {
+            "python_version": "3.11",
+            "dependencies": ["."],
+            "graphs": {"agent": "./agent.py:graph"},
+        }
+    )
+    assert config["pip_installer"] == "auto"
+
+    # Invalid pip_installer values should raise error
+    with pytest.raises(click.UsageError) as exc_info:
+        validate_config(
+            {
+                "python_version": "3.11",
+                "dependencies": ["."],
+                "graphs": {"agent": "./agent.py:graph"},
+                "pip_installer": "conda",
+            }
+        )
+    assert "Invalid pip_installer: 'conda'" in str(exc_info.value)
+    assert "Must be 'auto', 'pip', or 'uv'" in str(exc_info.value)
+
+    with pytest.raises(click.UsageError) as exc_info:
+        validate_config(
+            {
+                "python_version": "3.11",
+                "dependencies": ["."],
+                "graphs": {"agent": "./agent.py:graph"},
+                "pip_installer": "invalid",
+            }
+        )
+    assert "Invalid pip_installer: 'invalid'" in str(exc_info.value)
+
+
 def test_validate_config_file():
     with tempfile.TemporaryDirectory() as tmpdir:
         tmpdir_path = pathlib.Path(tmpdir)
@@ -799,6 +870,61 @@ WORKDIR /deps/__outer_unit_tests/unit_tests"""
     assert additional_contexts == {}
 
 
+def test_config_to_docker_pip_installer():
+    """Test that pip_installer setting affects the generated Dockerfile."""
+    graphs = {"agent": "./graphs/agent.py:graph"}
+    base_config = {
+        "python_version": "3.11",
+        "dependencies": ["."],
+        "graphs": graphs,
+    }
+
+    # Test default (auto) behavior with UV-supporting image
+    config_auto = validate_config(
+        {**copy.deepcopy(base_config), "pip_installer": "auto"}
+    )
+    docker_auto, _ = config_to_docker(
+        PATH_TO_CONFIG, config_auto, "langchain/langgraph-api:0.2.47"
+    )
+    assert "uv pip install --system" in docker_auto
+    assert "rm /usr/bin/uv /usr/bin/uvx" in docker_auto
+
+    # Test explicit pip setting
+    config_pip = validate_config({**copy.deepcopy(base_config), "pip_installer": "pip"})
+    docker_pip, _ = config_to_docker(
+        PATH_TO_CONFIG, config_pip, "langchain/langgraph-api:0.2.47"
+    )
+    assert "uv pip install --system" not in docker_pip
+    assert "pip install" in docker_pip
+    assert "rm /usr/bin/uv" not in docker_pip
+
+    # Test explicit uv setting
+    config_uv = validate_config({**copy.deepcopy(base_config), "pip_installer": "uv"})
+    docker_uv, _ = config_to_docker(
+        PATH_TO_CONFIG, config_uv, "langchain/langgraph-api:0.2.47"
+    )
+    assert "uv pip install --system" in docker_uv
+    assert "rm /usr/bin/uv /usr/bin/uvx" in docker_uv
+
+    # Test auto behavior with older image (should use pip)
+    config_auto_old = validate_config(
+        {**copy.deepcopy(base_config), "pip_installer": "auto"}
+    )
+    docker_auto_old, _ = config_to_docker(
+        PATH_TO_CONFIG, config_auto_old, "langchain/langgraph-api:0.2.46"
+    )
+    assert "uv pip install --system" not in docker_auto_old
+    assert "pip install" in docker_auto_old
+    assert "rm /usr/bin/uv" not in docker_auto_old
+
+    # Test that missing pip_installer defaults to auto behavior
+    config_default = validate_config(copy.deepcopy(base_config))
+    docker_default, _ = config_to_docker(
+        PATH_TO_CONFIG, config_default, "langchain/langgraph-api:0.2.47"
+    )
+    assert "uv pip install --system" in docker_default
+
+
 # config_to_compose
 def test_config_to_compose_simple_config():
     graphs = {"agent": "./agent.py:graph"}

libs/cli/uv.lock

@@ -501,7 +501,7 @@ wheels = [
 
 [[package]]
 name = "langgraph-cli"
-version = "0.3.1"
+version = "0.3.2"
 source = { editable = "." }
 dependencies = [
     { name = "click" },

libs/langgraph/langgraph/constants.py

@@ -39,8 +39,6 @@ ERROR = sys.intern("__error__")
 # for errors raised by nodes
 NO_WRITES = sys.intern("__no_writes__")
 # marker to signal node didn't write anything
-SCHEDULED = sys.intern("__scheduled__")
-# marker to signal node was scheduled (in distributed mode)
 TASKS = sys.intern("__pregel_tasks")
 # for Send objects returned by nodes/edges, corresponds to PUSH below
 RETURN = sys.intern("__return__")
@@ -71,13 +69,6 @@ CONFIG_KEY_RESUMING = sys.intern("__pregel_resuming")
 # holds a boolean indicating if subgraphs should resume from a previous checkpoint
 CONFIG_KEY_TASK_ID = sys.intern("__pregel_task_id")
 # holds the task ID for the current task
-CONFIG_KEY_DEDUPE_TASKS = sys.intern("__pregel_dedupe_tasks")
-# holds a boolean indicating if tasks should be deduplicated (for distributed mode)
-CONFIG_KEY_ENSURE_LATEST = sys.intern("__pregel_ensure_latest")
-# holds a boolean indicating whether to assert the requested checkpoint is the latest
-# (for distributed mode)
-CONFIG_KEY_DELEGATE = sys.intern("__pregel_delegate")
-# holds a boolean indicating whether to delegate subgraphs (for distributed mode)
 CONFIG_KEY_THREAD_ID = sys.intern("thread_id")
 # holds the thread ID for the current invocation
 CONFIG_KEY_CHECKPOINT_MAP = sys.intern("checkpoint_map")
@@ -121,7 +112,6 @@ RESERVED = {
     RESUME,
     ERROR,
     NO_WRITES,
-    SCHEDULED,
     # reserved config.configurable keys
     CONFIG_KEY_SEND,
     CONFIG_KEY_READ,
@@ -132,9 +122,6 @@ RESERVED = {
     CONFIG_KEY_CHECKPOINT_MAP,
     CONFIG_KEY_RESUMING,
     CONFIG_KEY_TASK_ID,
-    CONFIG_KEY_DEDUPE_TASKS,
-    CONFIG_KEY_ENSURE_LATEST,
-    CONFIG_KEY_DELEGATE,
     CONFIG_KEY_CHECKPOINT_MAP,
     CONFIG_KEY_CHECKPOINT_ID,
     CONFIG_KEY_CHECKPOINT_NS,

libs/langgraph/langgraph/errors.py

@@ -78,13 +78,6 @@ class NodeInterrupt(GraphInterrupt):
         super().__init__([Interrupt(value=value)])
 
 
-class GraphDelegate(GraphBubbleUp):
-    """Raised when a graph is delegated (for distributed mode)."""
-
-    def __init__(self, *args: dict[str, Any]) -> None:
-        super().__init__(*args)
-
-
 class ParentCommand(GraphBubbleUp):
     args: tuple[Command]
 
@@ -102,9 +95,3 @@ class TaskNotFound(Exception):
     """Raised when the executor is unable to find a task (for distributed mode)."""
 
     pass
-
-
-class CheckpointNotLatest(Exception):
-    """Raised when the checkpoint is not the latest version (for distributed mode)."""
-
-    pass

libs/langgraph/langgraph/graph/state.py

@@ -849,13 +849,6 @@ class StateGraph(Generic[StateT, InputT, OutputT]):
             builder=self,
             schema_to_mapper={},
             config_type=self.config_schema,
-            input_model=(
-                self.input_schema
-                if len(self.channels) > 1
-                and isclass(self.input_schema)
-                and issubclass(self.input_schema, BaseModel)
-                else None
-            ),
             nodes={},
             channels={
                 **self.channels,
@@ -1006,10 +999,7 @@ class CompiledStateGraph(
             if input_schema in self.schema_to_mapper:
                 mapper = self.schema_to_mapper[input_schema]
             else:
-                mapper = _pick_mapper(
-                    input_channels,
-                    input_schema,
-                )
+                mapper = _pick_mapper(input_channels, input_schema)
                 self.schema_to_mapper[input_schema] = mapper
 
             branch_channel = CHANNEL_BRANCH_TO.format(key)

libs/langgraph/langgraph/pregel/__init__.py

@@ -60,7 +60,6 @@ from langgraph.constants import (
     NS_SEP,
     NULL_TASK_ID,
     PUSH,
-    SCHEDULED,
     TASKS,
 )
 from langgraph.errors import (
@@ -590,8 +589,6 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou
 
     config_type: type[Any] | None = None
 
-    input_model: type[BaseModel] | None = None
-
     config: RunnableConfig | None = None
 
     name: str = "LangGraph"
@@ -619,7 +616,6 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou
         retry_policy: RetryPolicy | Sequence[RetryPolicy] = (),
         cache_policy: CachePolicy | None = None,
         config_type: type[Any] | None = None,
-        input_model: type[BaseModel] | None = None,
         config: RunnableConfig | None = None,
         trigger_to_nodes: Mapping[str, Sequence[str]] | None = None,
         name: str = "LangGraph",
@@ -651,7 +647,6 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou
         )
         self.cache_policy = cache_policy
         self.config_type = config_type
-        self.input_model = input_model
         self.config = config
         self.trigger_to_nodes = trigger_to_nodes or {}
         self.name = name
@@ -750,6 +745,7 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou
         validate_graph(
             self.nodes,
             {k: v for k, v in self.channels.items() if isinstance(v, BaseChannel)},
+            {k: v for k, v in self.channels.items() if not isinstance(v, BaseChannel)},
             self.input_channels,
             self.output_channels,
             self.stream_channels,
@@ -788,8 +784,6 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou
                 return channel.UpdateType
 
     def get_input_schema(self, config: RunnableConfig | None = None) -> type[BaseModel]:
-        if self.input_model is not None:
-            return self.input_model
         config = merge_configs(self.config, config)
         if isinstance(self.input_channels, str):
             return super().get_input_schema(config)
@@ -1008,7 +1002,7 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou
             )
         if apply_pending_writes and saved.pending_writes:
             for tid, k, v in saved.pending_writes:
-                if k in (ERROR, INTERRUPT, SCHEDULED):
+                if k in (ERROR, INTERRUPT):
                     continue
                 if tid not in next_tasks:
                     continue
@@ -1127,7 +1121,7 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou
             )
         if apply_pending_writes and saved.pending_writes:
             for tid, k, v in saved.pending_writes:
-                if k in (ERROR, INTERRUPT, SCHEDULED):
+                if k in (ERROR, INTERRUPT):
                     continue
                 if tid not in next_tasks:
                     continue
@@ -1466,7 +1460,7 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou
                         )
                     # apply writes from tasks that already ran
                     for tid, k, v in saved.pending_writes or []:
-                        if k in (ERROR, INTERRUPT, SCHEDULED):
+                        if k in (ERROR, INTERRUPT):
                             continue
                         if tid not in next_tasks:
                             continue
@@ -1630,7 +1624,7 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou
                     )
                 # apply writes
                 for tid, k, v in saved.pending_writes:
-                    if k in (ERROR, INTERRUPT, SCHEDULED):
+                    if k in (ERROR, INTERRUPT):
                         continue
                     if tid not in next_tasks:
                         continue
@@ -1886,7 +1880,7 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou
                         )
                     # apply writes from tasks that already ran
                     for tid, k, v in saved.pending_writes or []:
-                        if k in (ERROR, INTERRUPT, SCHEDULED):
+                        if k in (ERROR, INTERRUPT):
                             continue
                         if tid not in next_tasks:
                             continue
@@ -2049,7 +2043,7 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou
                         self.trigger_to_nodes,
                     )
                 for tid, k, v in saved.pending_writes:
-                    if k in (ERROR, INTERRUPT, SCHEDULED):
+                    if k in (ERROR, INTERRUPT):
                         continue
                     if tid not in next_tasks:
                         continue
@@ -2404,7 +2398,6 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou
                 config[CONF][CONFIG_KEY_CHECKPOINT_DURING] = checkpoint_during
             with SyncPregelLoop(
                 input,
-                input_model=self.input_model,
                 stream=StreamProtocol(stream.put, stream_modes),
                 config=config,
                 store=store,
@@ -2413,6 +2406,7 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou
                 nodes=self.nodes,
                 specs=self.channels,
                 output_keys=output_keys,
+                input_keys=self.input_channels,
                 stream_keys=self.stream_channels_asis,
                 interrupt_before=interrupt_before_,
                 interrupt_after=interrupt_after_,
@@ -2467,7 +2461,7 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou
                 # Channel updates from step N are only visible in step N+1
                 # channels are guaranteed to be immutable for the duration of the step,
                 # with channel updates applied only at the transition between steps.
-                while loop.tick(input_keys=self.input_channels):
+                while loop.tick():
                     for task in loop.match_cached_writes():
                         loop.output_writes(task.id, task.writes, cached=True)
                     for _ in runner.tick(
@@ -2478,6 +2472,7 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou
                     ):
                         # emit output
                         yield from output()
+                    loop.after_tick()
             # emit output
             yield from output()
             # handle exit
@@ -2647,7 +2642,6 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou
                 config[CONF][CONFIG_KEY_CHECKPOINT_DURING] = checkpoint_during
             async with AsyncPregelLoop(
                 input,
-                input_model=self.input_model,
                 stream=StreamProtocol(stream.put_nowait, stream_modes),
                 config=config,
                 store=store,
@@ -2656,6 +2650,7 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou
                 nodes=self.nodes,
                 specs=self.channels,
                 output_keys=output_keys,
+                input_keys=self.input_channels,
                 stream_keys=self.stream_channels_asis,
                 interrupt_before=interrupt_before_,
                 interrupt_after=interrupt_after_,
@@ -2701,7 +2696,7 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou
                 # channel updates from step N are only visible in step N+1
                 # channels are guaranteed to be immutable for the duration of the step,
                 # with channel updates applied only at the transition between steps
-                while loop.tick(input_keys=self.input_channels):
+                while loop.tick():
                     for task in await loop.amatch_cached_writes():
                         loop.output_writes(task.id, task.writes, cached=True)
                     async for _ in runner.atick(
@@ -2713,6 +2708,7 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou
                         # emit output
                         for o in output():
                             yield o
+                    loop.after_tick()
             # emit output
             for o in output():
                 yield o

libs/langgraph/langgraph/pregel/algo.py

@@ -1034,7 +1034,7 @@ def _proc_input(
             else:
                 return MISSING
         else:
-            val = managed[proc.channels].get(scratchpad)
+            return MISSING
     else:
         raise RuntimeError(
             f"Invalid channels type, expected list or dict, got {proc.channels}"

libs/langgraph/langgraph/pregel/loop.py

@@ -3,7 +3,6 @@ from __future__ import annotations
 import asyncio
 import binascii
 import concurrent.futures
-import dataclasses
 from collections import defaultdict, deque
 from collections.abc import Iterator, Mapping, Sequence
 from contextlib import (
@@ -25,7 +24,6 @@ from typing import (
 
 from langchain_core.callbacks import AsyncParentRunManager, ParentRunManager
 from langchain_core.runnables import RunnableConfig
-from pydantic import BaseModel
 from typing_extensions import ParamSpec, Self
 
 from langgraph.cache.base import BaseCache
@@ -46,9 +44,6 @@ from langgraph.constants import (
     CONFIG_KEY_CHECKPOINT_ID,
     CONFIG_KEY_CHECKPOINT_MAP,
     CONFIG_KEY_CHECKPOINT_NS,
-    CONFIG_KEY_DEDUPE_TASKS,
-    CONFIG_KEY_DELEGATE,
-    CONFIG_KEY_ENSURE_LATEST,
     CONFIG_KEY_RESUME_MAP,
     CONFIG_KEY_RESUMING,
     CONFIG_KEY_SCRATCHPAD,
@@ -60,17 +55,15 @@ from langgraph.constants import (
     INPUT,
     INTERRUPT,
     MISSING,
+    NS_END,
     NS_SEP,
     NULL_TASK_ID,
     PUSH,
     RESUME,
-    SCHEDULED,
     TAG_HIDDEN,
 )
 from langgraph.errors import (
-    CheckpointNotLatest,
     EmptyInputError,
-    GraphDelegate,
     GraphInterrupt,
 )
 from langgraph.managed.base import (
@@ -132,9 +125,7 @@ from langgraph.utils.config import patch_configurable
 V = TypeVar("V")
 P = ParamSpec("P")
 
-INPUT_DONE = object()
-INPUT_RESUMING = object()
-INPUT_SHOULD_VALIDATE = object()
+
 WritesT = Sequence[tuple[str, Any]]
 
 
@@ -155,11 +146,11 @@ class PregelLoop:
     stop: int
 
     input: Any | None
-    input_model: type[BaseModel] | None
     cache: BaseCache[WritesT] | None
     checkpointer: BaseCheckpointSaver | None
     nodes: Mapping[str, PregelNode]
     specs: Mapping[str, BaseChannel | ManagedValueSpec]
+    input_keys: str | Sequence[str]
     output_keys: str | Sequence[str]
     stream_keys: str | Sequence[str]
     skip_done_tasks: bool
@@ -202,11 +193,16 @@ class PregelLoop:
     prev_checkpoint_config: RunnableConfig | None
 
     status: Literal[
-        "pending", "done", "interrupt_before", "interrupt_after", "out_of_steps"
+        "input",
+        "pending",
+        "done",
+        "interrupt_before",
+        "interrupt_after",
+        "out_of_steps",
     ]
     tasks: dict[str, PregelExecutableTask]
-    to_interrupt: list[PregelExecutableTask]
     output: None | dict[str, Any] | Any = None
+    updated_channels: set[str] | None = None
 
     # public
 
@@ -221,13 +217,13 @@ class PregelLoop:
         checkpointer: BaseCheckpointSaver | None,
         nodes: Mapping[str, PregelNode],
         specs: Mapping[str, BaseChannel | ManagedValueSpec],
+        input_keys: str | Sequence[str],
         output_keys: str | Sequence[str],
         stream_keys: str | Sequence[str],
         trigger_to_nodes: Mapping[str, Sequence[str]],
         interrupt_after: All | Sequence[str] = EMPTY_SEQ,
         interrupt_before: All | Sequence[str] = EMPTY_SEQ,
         manager: None | AsyncParentRunManager | ParentRunManager = None,
-        input_model: type[BaseModel] | None = None,
         debug: bool = False,
         migrate_checkpoint: Callable[[Checkpoint], None] | None = None,
         retry_policy: Sequence[RetryPolicy] = (),
@@ -240,21 +236,18 @@ class PregelLoop:
         self.step = 0
         self.stop = 0
         self.input = input
-        self.input_model = input_model
         self.checkpointer = checkpointer
         self.cache = cache
         self.nodes = nodes
         self.specs = specs
+        self.input_keys = input_keys
         self.output_keys = output_keys
         self.stream_keys = stream_keys
         self.interrupt_after = interrupt_after
         self.interrupt_before = interrupt_before
         self.manager = manager
         self.is_nested = CONFIG_KEY_TASK_ID in self.config.get(CONF, {})
-        self.skip_done_tasks = (
-            CONFIG_KEY_CHECKPOINT_ID not in config[CONF]
-            or CONFIG_KEY_DEDUPE_TASKS in config[CONF]
-        )
+        self.skip_done_tasks = CONFIG_KEY_CHECKPOINT_ID not in config[CONF]
         self._migrate_checkpoint = migrate_checkpoint
         self.trigger_to_nodes = trigger_to_nodes
         self.retry_policy = retry_policy
@@ -264,9 +257,7 @@ class PregelLoop:
         if self.stream is not None and CONFIG_KEY_STREAM in config[CONF]:
             self.stream = DuplexStream(self.stream, config[CONF][CONFIG_KEY_STREAM])
         scratchpad: PregelScratchpad | None = config[CONF].get(CONFIG_KEY_SCRATCHPAD)
-        if not self.config[CONF].get(CONFIG_KEY_DELEGATE) and isinstance(
-            scratchpad, PregelScratchpad
-        ):
+        if isinstance(scratchpad, PregelScratchpad):
             # if count is > 0, append to checkpoint_ns
             # if count is 0, leave as is
             if cnt := scratchpad.subgraph_counter():
@@ -404,12 +395,6 @@ class PregelLoop:
         self, task: PregelExecutableTask, write_idx: int, call: Call | None = None
     ) -> PregelExecutableTask | None:
         """Accept a PUSH from a task, potentially returning a new task to start."""
-        # don't start if we should interrupt *after* the original task
-        if self.interrupt_after and should_interrupt(
-            self.checkpoint, self.interrupt_after, [task]
-        ):
-            self.to_interrupt.append(task)
-            return
         checkpoint_id_bytes = binascii.unhexlify(self.checkpoint["id"].replace("-", ""))
         null_version = checkpoint_null_version(self.checkpoint)
         if pushed := cast(
@@ -435,12 +420,6 @@ class PregelLoop:
                 cache_policy=self.cache_policy,
             ),
         ):
-            # don't start if we should interrupt *before* the new task
-            if self.interrupt_before and should_interrupt(
-                self.checkpoint, self.interrupt_before, [pushed]
-            ):
-                self.to_interrupt.append(pushed)
-                return
             # produce debug output
             self._emit("debug", map_debug_tasks, self.step, [pushed])
             # debug flag
@@ -454,11 +433,7 @@ class PregelLoop:
             # return the new task, to be started if not run before
             return pushed
 
-    def tick(
-        self,
-        *,
-        input_keys: str | Sequence[str],
-    ) -> bool:
+    def tick(self) -> bool:
         """Execute a single iteration of the Pregel loop.
 
         Args:
@@ -467,72 +442,6 @@ class PregelLoop:
         Returns:
             True if more iterations are needed.
         """
-        if self.status != "pending":
-            raise RuntimeError("Cannot tick when status is no longer 'pending'")
-
-        updated_channels: set[str] | None = None
-
-        if self.input not in (INPUT_DONE, INPUT_RESUMING, INPUT_SHOULD_VALIDATE):
-            updated_channels = self._first(input_keys=input_keys)
-        elif self.to_interrupt:
-            # if we need to interrupt, do so
-            self.status = "interrupt_before"
-            raise GraphInterrupt()
-        elif all(task.writes for task in self.tasks.values()):
-            # finish superstep
-            writes = [w for t in self.tasks.values() for w in t.writes]
-            # debug flag
-            if self.debug:
-                print_step_writes(
-                    self.step,
-                    writes,
-                    (
-                        [self.stream_keys]
-                        if isinstance(self.stream_keys, str)
-                        else self.stream_keys
-                    ),
-                )
-            # all tasks have finished
-            updated_channels = apply_writes(
-                self.checkpoint,
-                self.channels,
-                self.tasks.values(),
-                self.checkpointer_get_next_version,
-                self.trigger_to_nodes,
-            )
-            # validate input if requested
-            if self.input is INPUT_SHOULD_VALIDATE:
-                self.input = INPUT_DONE
-                # validate
-                cast(type[BaseModel], self.input_model)(
-                    **read_channels(self.channels, self.stream_keys)
-                )
-            # produce values output
-            if not updated_channels.isdisjoint(
-                (self.output_keys,)
-                if isinstance(self.output_keys, str)
-                else self.output_keys
-            ):
-                self._emit(
-                    "values", map_output_values, self.output_keys, writes, self.channels
-                )
-            # clear pending writes
-            self.checkpoint_pending_writes.clear()
-            # "not skip_done_tasks" only applies to first tick after resuming
-            self.skip_done_tasks = True
-            # save checkpoint
-            self._put_checkpoint({"source": "loop"})
-            # after execution, check if we should interrupt
-            if self.interrupt_after and should_interrupt(
-                self.checkpoint, self.interrupt_after, self.tasks.values()
-            ):
-                self.status = "interrupt_after"
-                raise GraphInterrupt()
-
-            # unset resuming flag
-            self.config[CONF].pop(CONFIG_KEY_RESUMING, None)
-        else:
-            return False
 
         # check if iteration limit is reached
         if self.step > self.stop:
@@ -554,11 +463,10 @@ class PregelLoop:
             store=self.store,
             checkpointer=self.checkpointer,
             trigger_to_nodes=self.trigger_to_nodes,
-            updated_channels=updated_channels,
+            updated_channels=self.updated_channels,
             retry_policy=self.retry_policy,
             cache_policy=self.cache_policy,
         )
-        self.to_interrupt = []
 
         # produce debug output
         if self._checkpointer_put_after_previous is not None:
@@ -588,26 +496,10 @@ class PregelLoop:
             self.status = "done"
             return False
 
-        # check if we should delegate (used by subgraphs in distributed mode)
-        if self.config[CONF].get(CONFIG_KEY_DELEGATE):
-            assert self.input is INPUT_RESUMING
-            raise GraphDelegate(
-                {
-                    "config": patch_configurable(
-                        self.config, {CONFIG_KEY_DELEGATE: False}
-                    ),
-                    "input": None,
-                }
-            )
-
         # if there are pending writes from a previous loop, apply them
         if self.skip_done_tasks and self.checkpoint_pending_writes:
             self._match_writes(self.tasks)
 
-        # if all tasks have finished, re-tick
-        if all(task.writes for task in self.tasks.values()):
-            return self.tick(input_keys=input_keys)
-
         # before execution, check if we should interrupt
         if self.interrupt_before and should_interrupt(
             self.checkpoint, self.interrupt_before, self.tasks.values()
@@ -629,6 +521,52 @@ class PregelLoop:
 
         return True
 
+    def after_tick(self) -> None:
+        # finish superstep
+        writes = [w for t in self.tasks.values() for w in t.writes]
+        # debug flag
+        if self.debug:
+            print_step_writes(
+                self.step,
+                writes,
+                (
+                    [self.stream_keys]
+                    if isinstance(self.stream_keys, str)
+                    else self.stream_keys
+                ),
+            )
+        # all tasks have finished
+        self.updated_channels = apply_writes(
+            self.checkpoint,
+            self.channels,
+            self.tasks.values(),
+            self.checkpointer_get_next_version,
+            self.trigger_to_nodes,
+        )
+        # produce values output
+        if not self.updated_channels.isdisjoint(
+            (self.output_keys,)
+            if isinstance(self.output_keys, str)
+            else self.output_keys
+        ):
+            self._emit(
+                "values", map_output_values, self.output_keys, writes, self.channels
+            )
+        # clear pending writes
+        self.checkpoint_pending_writes.clear()
+        # "not skip_done_tasks" only applies to first tick after resuming
+        self.skip_done_tasks = True
+        # save checkpoint
+        self._put_checkpoint({"source": "loop"})
+        # after execution, check if we should interrupt
+        if self.interrupt_after and should_interrupt(
+            self.checkpoint, self.interrupt_after, self.tasks.values()
+        ):
+            self.status = "interrupt_after"
+            raise GraphInterrupt()
+        # unset resuming flag
+        self.config[CONF].pop(CONFIG_KEY_RESUMING, None)
+
     def match_cached_writes(self) -> Sequence[PregelExecutableTask]:
         raise NotImplementedError
 
@@ -642,14 +580,7 @@ class PregelLoop:
             if k in (ERROR, INTERRUPT, RESUME):
                 continue
             if task := tasks.get(tid):
-                if k == SCHEDULED:
-                    if v == max(
-                        self.checkpoint["versions_seen"].get(INTERRUPT, {}).values(),
-                        default=None,
-                    ):
-                        self.tasks[tid] = dataclasses.replace(task, scheduled=True)
-                else:
-                    task.writes.append((k, v))
+                task.writes.append((k, v))
 
     def _first(self, *, input_keys: str | Sequence[str]) -> set[str] | None:
         # resuming from previous checkpoint requires
@@ -715,21 +646,8 @@ class PregelLoop:
             self._emit(
                 "values", map_output_values, self.output_keys, True, self.channels
             )
-            # set flag
-            self.input = INPUT_RESUMING
         # map inputs to channel updates
         elif input_writes := deque(map_input(input_keys, self.input)):
-            # TODO shouldn't these writes be passed to put_writes too?
-            # check if we should delegate (used by subgraphs in distributed mode)
-            if self.config[CONF].get(CONFIG_KEY_DELEGATE):
-                raise GraphDelegate(
-                    {
-                        "config": patch_configurable(
-                            self.config, {CONFIG_KEY_DELEGATE: False}
-                        ),
-                        "input": self.input,
-                    }
-                )
             # discard any unfinished tasks from previous checkpoint
             discard_tasks = prepare_next_tasks(
                 self.checkpoint,
@@ -758,24 +676,15 @@ class PregelLoop:
             )
             # save input checkpoint
             self._put_checkpoint({"source": "input"})
-            # set flag
-            if (
-                self.input_model is not None
-                and not isinstance(self.input, self.input_model)
-                and not isinstance(self.stream_keys, str)
-            ):
-                self.input = INPUT_SHOULD_VALIDATE
-            else:
-                self.input = INPUT_DONE
         elif CONFIG_KEY_RESUMING not in configurable:
             raise EmptyInputError(f"Received no input for {input_keys}")
-        else:
-            self.input = INPUT_DONE
         # update config
         if not self.is_nested:
             self.config = patch_configurable(
                 self.config, {CONFIG_KEY_RESUMING: is_resuming}
             )
+        # set flag
+        self.status = "pending"
         return updated_channels
 
     def _put_checkpoint(self, metadata: CheckpointMetadata) -> None:
@@ -868,7 +777,14 @@ class PregelLoop:
         traceback: TracebackType | None,
     ) -> bool | None:
         # persist current checkpoint and writes
-        if not self.checkpoint_during:
+        if not self.checkpoint_during and (
+            # if it's a top graph
+            not self.is_nested
+            # or a nested graph with error or interrupt
+            or exc_value is not None
+            # or a nested graph with checkpointer=True
+            or all(NS_END not in part for part in self.checkpoint_ns)
+        ):
             self._put_checkpoint(self.checkpoint_metadata)
             self._put_pending_writes()
         # suppress interrupt
@@ -992,9 +908,9 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager):
         manager: None | AsyncParentRunManager | ParentRunManager = None,
         interrupt_after: All | Sequence[str] = EMPTY_SEQ,
         interrupt_before: All | Sequence[str] = EMPTY_SEQ,
+        input_keys: str | Sequence[str] = EMPTY_SEQ,
         output_keys: str | Sequence[str] = EMPTY_SEQ,
         stream_keys: str | Sequence[str] = EMPTY_SEQ,
-        input_model: type[BaseModel] | None = None,
         debug: bool = False,
         migrate_checkpoint: Callable[[Checkpoint], None] | None = None,
         retry_policy: Sequence[RetryPolicy] = (),
@@ -1003,7 +919,6 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager):
     ) -> None:
         super().__init__(
             input,
-            input_model=input_model,
             stream=stream,
             config=config,
             checkpointer=checkpointer,
@@ -1011,6 +926,7 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager):
             store=store,
             nodes=nodes,
             specs=specs,
+            input_keys=input_keys,
             output_keys=output_keys,
             stream_keys=stream_keys,
             interrupt_after=interrupt_after,
@@ -1097,25 +1013,7 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager):
     # context manager
 
     def __enter__(self) -> Self:
-        if self.config.get(CONF, {}).get(
-            CONFIG_KEY_ENSURE_LATEST
-        ) and self.checkpoint_config[CONF].get(CONFIG_KEY_CHECKPOINT_ID):
-            if self.checkpointer is None:
-                raise RuntimeError(
-                    "Cannot ensure latest checkpoint without checkpointer"
-                )
-            saved = self.checkpointer.get_tuple(
-                patch_configurable(
-                    self.checkpoint_config, {CONFIG_KEY_CHECKPOINT_ID: None}
-                )
-            )
-            if (
-                saved is None
-                or saved.checkpoint["id"]
-                != self.checkpoint_config[CONF][CONFIG_KEY_CHECKPOINT_ID]
-            ):
-                raise CheckpointNotLatest
-        elif self.checkpointer:
+        if self.checkpointer:
             saved = self.checkpointer.get_tuple(self.checkpoint_config)
         else:
             saved = None
@@ -1149,10 +1047,11 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager):
             self.specs, self.checkpoint
         )
         self.stack.push(self._suppress_interrupt)
-        self.status = "pending"
+        self.status = "input"
         self.step = self.checkpoint_metadata["step"] + 1
         self.stop = self.step + self.config["recursion_limit"] + 1
         self.checkpoint_previous_versions = self.checkpoint["channel_versions"].copy()
+        self.updated_channels = self._first(input_keys=self.input_keys)
 
         return self
 
@@ -1182,9 +1081,9 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager):
         interrupt_after: All | Sequence[str] = EMPTY_SEQ,
         interrupt_before: All | Sequence[str] = EMPTY_SEQ,
         manager: None | AsyncParentRunManager | ParentRunManager = None,
+        input_keys: str | Sequence[str] = EMPTY_SEQ,
         output_keys: str | Sequence[str] = EMPTY_SEQ,
         stream_keys: str | Sequence[str] = EMPTY_SEQ,
-        input_model: type[BaseModel] | None = None,
         debug: bool = False,
         migrate_checkpoint: Callable[[Checkpoint], None] | None = None,
         retry_policy: Sequence[RetryPolicy] = (),
@@ -1193,7 +1092,6 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager):
     ) -> None:
         super().__init__(
             input,
-            input_model=input_model,
             stream=stream,
             config=config,
             checkpointer=checkpointer,
@@ -1201,6 +1099,7 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager):
             store=store,
             nodes=nodes,
             specs=specs,
+            input_keys=input_keys,
             output_keys=output_keys,
             stream_keys=stream_keys,
             interrupt_after=interrupt_after,
@@ -1290,25 +1189,7 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager):
     # context manager
 
     async def __aenter__(self) -> Self:
-        if self.config.get(CONF, {}).get(
-            CONFIG_KEY_ENSURE_LATEST
-        ) and self.checkpoint_config[CONF].get(CONFIG_KEY_CHECKPOINT_ID):
-            if self.checkpointer is None:
-                raise RuntimeError(
-                    "Cannot ensure latest checkpoint without checkpointer"
-                )
-            saved = await self.checkpointer.aget_tuple(
-                patch_configurable(
-                    self.checkpoint_config, {CONFIG_KEY_CHECKPOINT_ID: None}
-                )
-            )
-            if (
-                saved is None
-                or saved.checkpoint["id"]
-                != self.checkpoint_config[CONF][CONFIG_KEY_CHECKPOINT_ID]
-            ):
-                raise CheckpointNotLatest
-        elif self.checkpointer:
+        if self.checkpointer:
             saved = await self.checkpointer.aget_tuple(self.checkpoint_config)
         else:
             saved = None
@@ -1344,11 +1225,11 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager):
             self.specs, self.checkpoint
         )
         self.stack.push(self._suppress_interrupt)
-        self.status = "pending"
+        self.status = "input"
         self.step = self.checkpoint_metadata["step"] + 1
         self.stop = self.step + self.config["recursion_limit"] + 1
-
         self.checkpoint_previous_versions = self.checkpoint["channel_versions"].copy()
+        self.updated_channels = self._first(input_keys=self.input_keys)
 
         return self
 

libs/langgraph/langgraph/pregel/read.py

@@ -12,12 +12,11 @@ from langchain_core.runnables import Runnable, RunnableConfig
 
 from langgraph.constants import CONF, CONFIG_KEY_READ
 from langgraph.pregel.protocol import PregelProtocol
-from langgraph.pregel.retry import RetryPolicy
 from langgraph.pregel.utils import find_subgraph_pregel
 from langgraph.pregel.write import ChannelWrite
-from langgraph.types import CachePolicy
+from langgraph.types import CachePolicy, RetryPolicy
 from langgraph.utils.config import merge_configs
-from langgraph.utils.runnable import RunnableCallable, RunnableSeq, coerce_to_runnable
+from langgraph.utils.runnable import RunnableCallable, RunnableSeq
 
 READ_TYPE = Callable[[Union[str, Sequence[str]], bool], Union[Any, dict[str, Any]]]
 INPUT_CACHE_KEY_TYPE = tuple[Callable[..., Any], tuple[str, ...]]
@@ -96,7 +95,7 @@ class ChannelRead(RunnableCallable):
 DEFAULT_BOUND = RunnableCallable(lambda input: input)
 
 
-class PregelNode(Runnable):
+class PregelNode:
     """A node in a Pregel graph. This won't be invoked as a runnable by the graph
     itself, but instead acts as a container for the components necessary to make
     a PregelExecutableTask for a node."""
@@ -227,38 +226,6 @@ class PregelNode(Runnable):
             else (self.channels,),
         )
 
-    def __or__(
-        self,
-        other: Runnable[Any, Any]
-        | Callable[[Any], Any]
-        | Mapping[str, Runnable[Any, Any] | Callable[[Any], Any]],
-    ) -> PregelNode:
-        if isinstance(other, Runnable) and ChannelWrite.is_writer(other):
-            return self.copy(update=dict(writers=[*self.writers, other]))
-        elif self.bound is DEFAULT_BOUND:
-            return self.copy(
-                update=dict(bound=coerce_to_runnable(other, name=None, trace=True))
-            )
-        else:
-            return self.copy(update=dict(bound=RunnableSeq(self.bound, other)))
-
-    def pipe(
-        self,
-        *others: Runnable[Any, Any] | Callable[[Any], Any],
-        name: str | None = None,
-    ) -> PregelNode:
-        for other in others:
-            self = self | other
-        return self
-
-    def __ror__(
-        self,
-        other: Runnable[Any, Any]
-        | Callable[[Any], Any]
-        | Mapping[str, Runnable[Any, Any] | Callable[[Any], Any]],
-    ) -> PregelNode:
-        raise NotImplementedError()
-
     def invoke(
         self,
         input: Any,

libs/langgraph/langgraph/pregel/validate.py

@@ -5,6 +5,7 @@ from typing import Any
 
 from langgraph.channels.base import BaseChannel
 from langgraph.constants import RESERVED
+from langgraph.managed.base import ManagedValueMapping
 from langgraph.pregel.read import PregelNode
 from langgraph.types import All
 
@@ -12,6 +13,7 @@ from langgraph.types import All
 def validate_graph(
     nodes: Mapping[str, PregelNode],
     channels: dict[str, BaseChannel],
+    managed: ManagedValueMapping,
     input_channels: str | Sequence[str],
     output_channels: str | Sequence[str],
     stream_channels: str | Sequence[str] | None,
@@ -20,14 +22,30 @@ def validate_graph(
 ) -> None:
     for chan in channels:
         if chan in RESERVED:
-            raise ValueError(f"Channel names {chan} are reserved")
+            raise ValueError(f"Channel name '{chan}' is reserved")
+    for name in managed:
+        if name in RESERVED:
+            raise ValueError(f"Managed name '{name}' is reserved")
 
     subscribed_channels = set[str]()
     for name, node in nodes.items():
         if name in RESERVED:
-            raise ValueError(f"Node names {RESERVED} are reserved")
+            raise ValueError(f"Node name '{name}' is reserved")
         if isinstance(node, PregelNode):
             subscribed_channels.update(node.triggers)
+            if isinstance(node.channels, str):
+                if node.channels not in channels:
+                    raise ValueError(
+                        f"Node {name} reads channel '{node.channels}' "
+                        f"not in known channels: '{repr(sorted(channels))[:100]}'"
+                    )
+            else:
+                for chan in node.channels:
+                    if chan not in channels and chan not in managed:
+                        raise ValueError(
+                            f"Node {name} reads channel '{chan}' "
+                            f"not in known channels: '{repr(sorted(channels))[:100]}'"
+                        )
         else:
             raise TypeError(
                 f"Invalid node type {type(node)}, expected PregelNode or NodeBuilder"

libs/langgraph/langgraph/types.py

@@ -203,7 +203,6 @@ class PregelExecutableTask:
     cache_key: CacheKey | None
     id: str
     path: tuple[str | int | tuple, ...]
-    scheduled: bool = False
     writers: Sequence[Runnable] = ()
     subgraphs: Sequence[PregelProtocol] = ()
 

libs/langgraph/tests/test_large_cases.py

@@ -4213,44 +4213,6 @@ def test_doubly_nested_graph_state(
     # get child graph history
     child_history = list(app.get_state_history(outer_history[1].tasks[0].state))
     assert child_history == [
-        StateSnapshot(
-            values={"my_key": "hi my value here and there"},
-            next=(),
-            config={
-                "configurable": {
-                    "thread_id": "1",
-                    "checkpoint_ns": AnyStr("child:"),
-                    "checkpoint_id": AnyStr(),
-                    "checkpoint_map": AnyDict(
-                        {"": AnyStr(), AnyStr("child:"): AnyStr()}
-                    ),
-                }
-            },
-            metadata={
-                "source": "loop",
-                "step": 1,
-                "parents": {"": AnyStr()},
-                "thread_id": "1",
-                "langgraph_node": "child",
-                "langgraph_path": [PULL, AnyStr("child")],
-                "langgraph_step": 2,
-                "langgraph_triggers": ["branch:to:child"],
-                "langgraph_checkpoint_ns": AnyStr("child:"),
-            },
-            created_at=AnyStr(),
-            parent_config={
-                "configurable": {
-                    "thread_id": "1",
-                    "checkpoint_ns": AnyStr("child:"),
-                    "checkpoint_id": AnyStr(),
-                    "checkpoint_map": AnyDict(
-                        {"": AnyStr(), AnyStr("child:"): AnyStr()}
-                    ),
-                }
-            },
-            tasks=(),
-            interrupts=(),
-        ),
         StateSnapshot(
             values={"my_key": "hi my value"},
             next=("child_1",),
@@ -4295,62 +4257,8 @@ def test_doubly_nested_graph_state(
         ),
     ]
     # get grandchild graph history
-    grandchild_history = list(app.get_state_history(child_history[1].tasks[0].state))
+    grandchild_history = list(app.get_state_history(child_history[0].tasks[0].state))
     assert grandchild_history == [
-        StateSnapshot(
-            values={"my_key": "hi my value here and there"},
-            next=(),
-            config={
-                "configurable": {
-                    "thread_id": "1",
-                    "checkpoint_ns": AnyStr(),
-                    "checkpoint_id": AnyStr(),
-                    "checkpoint_map": AnyDict(
-                        {
-                            "": AnyStr(),
-                            AnyStr("child:"): AnyStr(),
-                            AnyStr(re.compile(r"child:.+|child1:")): AnyStr(),
-                        }
-                    ),
-                }
-            },
-            metadata={
-                "source": "loop",
-                "step": 2,
-                "parents": AnyDict(
-                    {
-                        "": AnyStr(),
-                        AnyStr("child:"): AnyStr(),
-                    }
-                ),
-                "thread_id": "1",
-                "langgraph_checkpoint_ns": AnyStr("child:"),
-                "langgraph_node": "child_1",
-                "langgraph_path": [
-                    PULL,
-                    AnyStr("child_1"),
-                ],
-                "langgraph_step": 1,
-                "langgraph_triggers": ["branch:to:child_1"],
-            },
-            created_at=AnyStr(),
-            parent_config={
-                "configurable": {
-                    "thread_id": "1",
-                    "checkpoint_ns": AnyStr(),
-                    "checkpoint_id": AnyStr(),
-                    "checkpoint_map": AnyDict(
-                        {
-                            "": AnyStr(),
-                            AnyStr("child:"): AnyStr(),
-                            AnyStr(re.compile(r"child:.+|child1:")): AnyStr(),
-                        }
-                    ),
-                }
-            },
-            tasks=(),
-            interrupts=(),
-        ),
         StateSnapshot(
             values={"my_key": "hi my value here"},
             next=("grandchild_2",),
@@ -4418,7 +4326,7 @@ def test_send_to_nested_graphs(sync_checkpointer: BaseCheckpointSaver) -> None:
         return {"subject": f"{subject} - hohoho"}
 
     # subgraph
-    subgraph = StateGraph(JokeState, output=OverallState)
+    subgraph = StateGraph(JokeState, output_schema=OverallState)
     subgraph.add_node("edit", edit)
     subgraph.add_node(
         "generate", lambda state: {"jokes": [f"Joke about {state['subject']}"]}

libs/langgraph/tests/test_large_cases_async.py

@@ -3028,44 +3028,6 @@ async def test_doubly_nested_graph_state(
         c async for c in app.aget_state_history(outer_history[1].tasks[0].state)
     ]
     assert child_history == [
-        StateSnapshot(
-            values={"my_key": "hi my value here and there"},
-            next=(),
-            config={
-                "configurable": {
-                    "thread_id": "1",
-                    "checkpoint_ns": AnyStr("child:"),
-                    "checkpoint_id": AnyStr(),
-                    "checkpoint_map": AnyDict(
-                        {"": AnyStr(), AnyStr("child:"): AnyStr()}
-                    ),
-                }
-            },
-            metadata={
-                "source": "loop",
-                "step": 1,
-                "parents": {"": AnyStr()},
-                "thread_id": "1",
-                "langgraph_node": "child",
-                "langgraph_path": [PULL, AnyStr("child")],
-                "langgraph_step": 2,
-                "langgraph_triggers": ["branch:to:child"],
-                "langgraph_checkpoint_ns": AnyStr("child:"),
-            },
-            created_at=AnyStr(),
-            parent_config={
-                "configurable": {
-                    "thread_id": "1",
-                    "checkpoint_ns": AnyStr("child:"),
-                    "checkpoint_id": AnyStr(),
-                    "checkpoint_map": AnyDict(
-                        {"": AnyStr(), AnyStr("child:"): AnyStr()}
-                    ),
-                }
-            },
-            tasks=(),
-            interrupts=(),
-        ),
         StateSnapshot(
             values={"my_key": "hi my value"},
             next=("child_1",),
@@ -3111,65 +3073,9 @@ async def test_doubly_nested_graph_state(
     ]
     # get grandchild graph history
     grandchild_history = [
-        c async for c in app.aget_state_history(child_history[1].tasks[0].state)
+        c async for c in app.aget_state_history(child_history[0].tasks[0].state)
     ]
     assert grandchild_history == [
-        StateSnapshot(
-            values={"my_key": "hi my value here and there"},
-            next=(),
-            config={
-                "configurable": {
-                    "thread_id": "1",
-                    "checkpoint_ns": AnyStr(),
-                    "checkpoint_id": AnyStr(),
-                    "checkpoint_map": AnyDict(
-                        {
-                            "": AnyStr(),
-                            AnyStr("child:"): AnyStr(),
-                            AnyStr(re.compile(r"child:.+|child1:")): AnyStr(),
-                        }
-                    ),
-                }
-            },
-            metadata={
-                "source": "loop",
-                "step": 2,
-                "parents": AnyDict(
-                    {
-                        "": AnyStr(),
-                        AnyStr("child:"): AnyStr(),
-                    }
-                ),
-                "thread_id": "1",
-                "langgraph_checkpoint_ns": AnyStr("child:"),
-                "langgraph_node": "child_1",
-                "langgraph_path": [
-                    PULL,
-                    AnyStr("child_1"),
-                ],
-                "langgraph_step": 1,
-                "langgraph_triggers": [
-                    "branch:to:child_1",
-                ],
-            },
-            created_at=AnyStr(),
-            parent_config={
-                "configurable": {
-                    "thread_id": "1",
-                    "checkpoint_ns": AnyStr(),
-                    "checkpoint_id": AnyStr(),
-                    "checkpoint_map": AnyDict(
-                        {
-                            "": AnyStr(),
-                            AnyStr("child:"): AnyStr(),
-                            AnyStr(re.compile(r"child:.+|child1:")): AnyStr(),
-                        }
-                    ),
-                }
-            },
-            tasks=(),
-            interrupts=(),
-        ),
         StateSnapshot(
             values={"my_key": "hi my value here"},
             next=("grandchild_2",),
@@ -3239,7 +3145,7 @@ async def test_send_to_nested_graphs(async_checkpointer: BaseCheckpointSaver) ->
         return {"subject": f"{subject} - hohoho"}
 
     # subgraph
-    subgraph = StateGraph(JokeState, output=OverallState)
+    subgraph = StateGraph(JokeState, output_schema=OverallState)
     subgraph.add_node("edit", edit)
     subgraph.add_node(
         "generate", lambda state: {"jokes": [f"Joke about {state['subject']}"]}

libs/langgraph/tests/test_pregel.py

@@ -3276,6 +3276,57 @@ def test_subgraph_checkpoint_true(
         ),
     ]
 
+    checkpoints = list(app.get_state_history(config))
+    if checkpoint_during:
+        assert len(checkpoints) == 4
+    else:
+        assert len(checkpoints) == 1
+
+
+def test_subgraph_checkpoint_during_false_inherited() -> None:
+    sync_checkpointer = InMemorySaver()
+
+    class InnerState(TypedDict):
+        my_key: Annotated[str, operator.add]
+        my_other_key: str
+
+    def inner_1(state: InnerState):
+        return {"my_key": " got here", "my_other_key": state["my_key"]}
+
+    def inner_2(state: InnerState):
+        return {"my_key": " and there"}
+
+    inner = StateGraph(InnerState)
+    inner.add_node("inner_1", inner_1)
+    inner.add_node("inner_2", inner_2)
+    inner.add_edge("inner_1", "inner_2")
+    inner.set_entry_point("inner_1")
+    inner.set_finish_point("inner_2")
+
+    class State(TypedDict):
+        my_key: str
+
+    inner_app = inner.compile(checkpointer=sync_checkpointer)
+    graph = StateGraph(State)
+    graph.add_node("inner", inner_app)
+    graph.add_edge(START, "inner")
+    graph.add_conditional_edges(
+        "inner", lambda s: "inner" if s["my_key"].count("there") < 2 else END
+    )
+    app = graph.compile(checkpointer=sync_checkpointer)
+    for checkpoint_during in [True, False]:
+        thread_id = str(uuid.uuid4())
+        config = {"configurable": {"thread_id": thread_id}}
+        app.invoke(
+            {"my_key": ""}, config, subgraphs=True, checkpoint_during=checkpoint_during
+        )
+        if checkpoint_during:
+            checkpoints = list(sync_checkpointer.list(config))
+            assert len(checkpoints) == 12
+        else:
+            checkpoints = list(sync_checkpointer.list(config))
+            assert len(checkpoints) == 1
+
 
 def test_subgraph_checkpoint_true_interrupt(
     sync_checkpointer: BaseCheckpointSaver, checkpoint_during: bool
@@ -4300,7 +4351,7 @@ def test_store_injected(
     builder = StateGraph(State)
     builder.add_node("node", Node())
     builder.add_edge("__start__", "node")
-    N = 500
+    N = 50
     M = 1
 
     for i in range(N):
@@ -4575,11 +4626,14 @@ def test_debug_nested_subgraphs(
 
         return clean_config
 
-    for checkpoint_events, checkpoint_history in zip(
-        stream_ns.values(), history_ns.values()
+    for checkpoint_events, checkpoint_history, ns in zip(
+        stream_ns.values(), history_ns.values(), stream_ns.keys()
     ):
         if not checkpoint_during:
             checkpoint_events = checkpoint_events[-1:]
+            if ns:  # Save no checkpoints for subgraphs when checkpoint_during=False
+                assert not checkpoint_history
+                continue
         assert len(checkpoint_events) == len(checkpoint_history)
         for stream, history in zip(checkpoint_events, checkpoint_history):
             assert stream["values"] == history.values

libs/langgraph/tests/test_pregel_async.py

@@ -1338,7 +1338,7 @@ async def test_node_schemas_custom_output() -> None:
             "now": 123,
         }
 
-    builder = StateGraph(State, output=Output)
+    builder = StateGraph(State, output_schema=Output)
     builder.add_node("a", node_a)
     builder.add_node("b", node_b)
     builder.add_node("c", node_c)
@@ -1353,7 +1353,7 @@ async def test_node_schemas_custom_output() -> None:
         "messages": [_AnyIdHumanMessage(content="hello")],
     }
 
-    builder = StateGraph(State, output=Output)
+    builder = StateGraph(State, output_schema=Output)
     builder.add_node("a", node_a)
     builder.add_node("b", node_b)
     builder.add_node("c", node_c)
@@ -5029,6 +5029,51 @@ async def test_subgraph_checkpoint_true(
     ]
 
 
+async def test_subgraph_checkpoint_during_false_inherited() -> None:
+    async_checkpointer = InMemorySaver()
+
+    class InnerState(TypedDict):
+        my_key: Annotated[str, operator.add]
+        my_other_key: str
+
+    def inner_1(state: InnerState):
+        return {"my_key": " got here", "my_other_key": state["my_key"]}
+
+    def inner_2(state: InnerState):
+        return {"my_key": " and there"}
+
+    inner = StateGraph(InnerState)
+    inner.add_node("inner_1", inner_1)
+    inner.add_node("inner_2", inner_2)
+    inner.add_edge("inner_1", "inner_2")
+    inner.set_entry_point("inner_1")
+    inner.set_finish_point("inner_2")
+
+    class State(TypedDict):
+        my_key: str
+
+    inner_app = inner.compile(checkpointer=async_checkpointer)
+    graph = StateGraph(State)
+    graph.add_node("inner", inner_app)
+    graph.add_edge(START, "inner")
+    graph.add_conditional_edges(
+        "inner", lambda s: "inner" if s["my_key"].count("there") < 2 else END
+    )
+    app = graph.compile(checkpointer=async_checkpointer)
+    for checkpoint_during in [True, False]:
+        thread_id = str(uuid.uuid4())
+        config = {"configurable": {"thread_id": thread_id}}
+        await app.ainvoke(
+            {"my_key": ""}, config, subgraphs=True, checkpoint_during=checkpoint_during
+        )
+        if checkpoint_during:
+            checkpoints = list(async_checkpointer.list(config))
+            assert len(checkpoints) == 12
+        else:
+            checkpoints = list(async_checkpointer.list(config))
+            assert len(checkpoints) == 1
+
+
 @NEEDS_CONTEXTVARS
 async def test_subgraph_checkpoint_true_interrupt(
     async_checkpointer: BaseCheckpointSaver, checkpoint_during: bool
@@ -5736,7 +5781,7 @@ async def test_store_injected_async(
     builder.add_edge("__start__", "node")
     builder.add_edge("node", "other_node")
 
-    N = 500
+    N = 50
     M = 1
 
     for i in range(N):
@@ -6007,11 +6052,14 @@ async def test_debug_nested_subgraphs(
 
         return clean_config
 
-    for checkpoint_events, checkpoint_history in zip(
-        stream_ns.values(), history_ns.values()
+    for checkpoint_events, checkpoint_history, ns in zip(
+        stream_ns.values(), history_ns.values(), stream_ns.keys()
     ):
         if not checkpoint_during:
             checkpoint_events = checkpoint_events[-1:]
+            if ns:  # Save no checkpoints for subgraphs when checkpoint_during=False
+                assert not checkpoint_history
+                continue
         assert len(checkpoint_events) == len(checkpoint_history)
         for stream, history in zip(checkpoint_events, checkpoint_history):
             assert stream["values"] == history.values
@@ -6982,14 +7030,17 @@ async def test_multiple_subgraphs(async_checkpointer: BaseCheckpointSaver) -> No
         return {"result": state["a"] + state["b"]}
 
     add_subgraph = (
-        StateGraph(State, output=Output).add_node(add).add_edge(START, "add").compile()
+        StateGraph(State, output_schema=Output)
+        .add_node(add)
+        .add_edge(START, "add")
+        .compile()
     )
 
     async def multiply(state):
         return {"result": state["a"] * state["b"]}
 
     multiply_subgraph = (
-        StateGraph(State, output=Output)
+        StateGraph(State, output_schema=Output)
         .add_node(multiply)
         .add_edge(START, "multiply")
         .compile()
@@ -7002,7 +7053,7 @@ async def test_multiple_subgraphs(async_checkpointer: BaseCheckpointSaver) -> No
         return another_result
 
     parent_call_same_subgraph = (
-        StateGraph(State, output=Output)
+        StateGraph(State, output_schema=Output)
         .add_node(call_same_subgraph)
         .add_edge(START, "call_same_subgraph")
         .compile(checkpointer=async_checkpointer)
@@ -7026,7 +7077,7 @@ async def test_multiple_subgraphs(async_checkpointer: BaseCheckpointSaver) -> No
         }
 
     parent_call_multiple_subgraphs = (
-        StateGraph(State, output=Output)
+        StateGraph(State, output_schema=Output)
         .add_node(call_multiple_subgraphs)
         .add_edge(START, "call_multiple_subgraphs")
         .compile(checkpointer=async_checkpointer)
@@ -7104,14 +7155,17 @@ async def test_multiple_subgraphs_mixed_entrypoint(
         return {"result": state["a"] + state["b"]}
 
     add_subgraph = (
-        StateGraph(State, output=Output).add_node(add).add_edge(START, "add").compile()
+        StateGraph(State, output_schema=Output)
+        .add_node(add)
+        .add_edge(START, "add")
+        .compile()
     )
 
     async def multiply(state):
         return {"result": state["a"] * state["b"]}
 
     multiply_subgraph = (
-        StateGraph(State, output=Output)
+        StateGraph(State, output_schema=Output)
         .add_node(multiply)
         .add_edge(START, "multiply")
         .compile()
@@ -7181,7 +7235,7 @@ async def test_multiple_subgraphs_mixed_state_graph(
         return {"result": another_result}
 
     parent_call_same_subgraph = (
-        StateGraph(State, output=Output)
+        StateGraph(State, output_schema=Output)
         .add_node(call_same_subgraph)
         .add_edge(START, "call_same_subgraph")
         .compile(checkpointer=async_checkpointer)
@@ -7205,7 +7259,7 @@ async def test_multiple_subgraphs_mixed_state_graph(
         }
 
     parent_call_multiple_subgraphs = (
-        StateGraph(State, output=Output)
+        StateGraph(State, output_schema=Output)
         .add_node(call_multiple_subgraphs)
         .add_edge(START, "call_multiple_subgraphs")
         .compile(checkpointer=async_checkpointer)

libs/langgraph/tests/test_state.py

@@ -92,7 +92,7 @@ def test_state_schema_with_type_hint():
             assert state.pop("foo") == "bar"
             return {"input_state": state}
 
-    graph = StateGraph(InputState, output=OutputState)
+    graph = StateGraph(InputState, output_schema=OutputState)
     actions = [
         complete_hint,
         miss_first_hint,

libs/langgraph/uv.lock

@@ -1329,6 +1329,7 @@ dev = [
     { name = "dataclasses-json" },
     { name = "mypy" },
     { name = "numpy" },
+    { name = "pandas" },
     { name = "pytest" },
     { name = "pytest-asyncio" },
     { name = "pytest-mock" },

libs/prebuilt/tests/test_react_agent.py

@@ -1089,14 +1089,17 @@ def test_react_with_subgraph_tools(
         return {"result": state["a"] + state["b"]}
 
     add_subgraph = (
-        StateGraph(State, output=Output).add_node(add).add_edge(START, "add").compile()
+        StateGraph(State, output_schema=Output)
+        .add_node(add)
+        .add_edge(START, "add")
+        .compile()
     )
 
     def multiply(state):
         return {"result": state["a"] * state["b"]}
 
     multiply_subgraph = (
-        StateGraph(State, output=Output)
+        StateGraph(State, output_schema=Output)
         .add_node(multiply)
         .add_edge(START, "multiply")
         .compile()

libs/prebuilt/uv.lock

@@ -390,6 +390,7 @@ dev = [
     { name = "dataclasses-json" },
     { name = "mypy" },
     { name = "numpy" },
+    { name = "pandas" },
     { name = "pytest" },
     { name = "pytest-asyncio" },
     { name = "pytest-mock" },