Merge branch 'main' into cc/update_quickstart

.github/workflows/deploy_docs.yml

@@ -85,7 +85,8 @@ jobs:
           if [ "${{ github.event_name }}" == "schedule" ] || [ "${{ github.event_name }}" == "workflow_dispatch" ] || ([ "${{ github.event_name }}" == "push" ] && [ "${{ github.ref }}" == "refs/heads/main" ]); then
             echo "Running link check on all HTML files matching notebooks in docs directory..."
             poetry run pytest -v \
-              --check-links-ignore "https://(api|web|docs|academy)\.smith\.langchain\.com/.*" \
+              --check-links-ignore "https://(api|web|docs)\.smith\.langchain\.com/.*" \
+              --check-links-ignore "https://academy\.langchain\.com/.*" \
               --check-links-ignore "https://x.com/.*" \
               --check-links-ignore "https://github\.com/.*" \
               --check-links-ignore "http://localhost:8123/.*" \
@@ -106,7 +107,8 @@ jobs:
             if [ -n "${CHANGED_FILES}" ]; then
               echo "Running link check on HTML files matching changed notebook files..."
               poetry run pytest -v \
-                --check-links-ignore "https://(api|web|docs|academy)\.smith\.langchain\.com/.*" \
+                --check-links-ignore "https://(api|web|docs)\.smith\.langchain\.com/.*" \
+                --check-links-ignore "https://academy\.langchain\.com/.*" \
                 --check-links-ignore "http://localhost:8123/.*" \
                 --check-links-ignore "http://localhost:2024.*" \
                 --check-links-ignore "http://127.0.0.1:.*" \

docs/docs/cloud/reference/api/openapi_control_plane.json

@@ -561,6 +561,16 @@
           },
           "resource": {
             "$ref": "#/components/schemas/ResourceService"
+          },
+          "status": {
+            "type": "string",
+            "enum": [
+              "AWAITING_DATABASE",
+              "READY",
+              "AWAITING_DELETE",
+              "UNKNOWN"
+            ],
+            "description": "Deployment status of the project.\n\nNon-terminal statuses: `AWAITING_DATABASE`, `AWAITING_DELETE`. All other statuses are terminal."
           }
         }
       },

docs/docs/concepts/langgraph_cloud.md

@@ -6,21 +6,29 @@
 
 ## Overview
 
-LangGraph's Cloud SaaS is a managed service for deploying LangGraph APIs, regardless of its definition or dependencies. The service offers managed implementations of checkpointers and stores, allowing you to focus on building the right cognitive architecture for your use case. By handling scalable & secure infrastructure, LangGraph Cloud offers the fastest path to getting your LangGraph API deployed to production.
+LangGraph's Cloud SaaS is a managed service for deploying LangGraph Servers, regardless of its definition or dependencies. The service offers managed implementations of checkpointers and stores, allowing you to focus on building the right cognitive architecture for your use case. By handling scalable & secure infrastructure, LangGraph Cloud SaaS offers the fastest path to getting your LangGraph Server deployed to production.
 
 ## Deployment
 
-A **deployment** is an instance of a LangGraph API. A single deployment can have many [revisions](#revision). When a deployment is created, all the necessary infrastructure (e.g. database, containers, secrets store) are automatically provisioned. See the [architecture diagram](#architecture) below for more details.
+A **deployment** is an instance of a LangGraph Server. A single deployment can have many [revisions](#revision). When a deployment is created, all the necessary infrastructure (e.g. database, containers, secrets store) are automatically provisioned. See the [architecture diagram](#architecture) below for more details.
 
-See the [how-to guide](../cloud/deployment/cloud.md#create-new-deployment) for creating a new deployment.
-
-## Resource Allocation
+Resource Allocation:
 
 | **Deployment Type** | **CPU** | **Memory** | **Scaling**         |
 |---------------------|---------|------------|---------------------|
 | Development         | 1 CPU   | 1 GB       | Up to 1 container   |
 | Production          | 2 CPU   | 2 GB       | Up to 10 containers |
 
+See the [how-to guide](../cloud/deployment/cloud.md#create-new-deployment) for creating a new deployment.
+
+## Persistence
+
+A dedicated database is automatically created for each deployment. The database serves as the [persistence layer](../concepts/persistence.md) for the deployment.
+
+When defining a graph to be deployed to LangGraph Cloud SaaS, a [checkpointer](../concepts/persistence.md#checkpointer-libraries) should not be configured by the user. Instead, a checkpointer is automatically configured for the graph.
+
+There is no direct access to the database. All access to the database occurs through the LangGraph Server APIs.
+
 ## Autoscaling
 `Production` type deployments automatically scale up to 10 containers. Scaling is based on the current request load for a single container. Specifically, the autoscaling implementation scales the deployment so that each container is processing about 10 concurrent requests. For example... 
 

docs/mkdocs.yml

@@ -20,6 +20,7 @@ theme:
     - content.action.edit
     - content.tooltips
     - header.autohide
+    - navigation.indexes
     - navigation.expand
     - navigation.footer
     - navigation.instant

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

@@ -57,6 +57,9 @@ MIGRATIONS = [
     PRIMARY KEY (thread_id, checkpoint_ns, checkpoint_id, task_id, idx)
 );""",
     "ALTER TABLE checkpoint_blobs ALTER COLUMN blob DROP not null;",
+    # NOTE: this is a no-op migration to ensure that the versions in the migrations table are correct.
+    # This is necessary due to an empty migration previously added to the list.
+    "SELECT 1;",
     """
     CREATE INDEX CONCURRENTLY IF NOT EXISTS checkpoints_thread_id_idx ON checkpoints(thread_id);
     """,

libs/checkpoint-postgres/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "langgraph-checkpoint-postgres"
-version = "2.0.12"
+version = "2.0.13"
 description = "Library with a Postgres implementation of LangGraph checkpoint saver."
 authors = []
 license = "MIT"

libs/langgraph/langgraph/constants.py

@@ -75,11 +75,11 @@ CONFIG_KEY_CHECKPOINT_ID = sys.intern("checkpoint_id")
 CONFIG_KEY_CHECKPOINT_NS = sys.intern("checkpoint_ns")
 # holds the current checkpoint_ns, "" for root graph
 CONFIG_KEY_NODE_FINISHED = sys.intern("__pregel_node_finished")
-# holds the value that "answers" an interrupt() call
-CONFIG_KEY_WRITES = sys.intern("__pregel_writes")
-# read-only list of existing task writes
+# holds a callback to be called when a node is finished
 CONFIG_KEY_SCRATCHPAD = sys.intern("__pregel_scratchpad")
 # holds a mutable dict for temporary storage scoped to the current task
+CONFIG_KEY_END = sys.intern("__pregel_previous")
+# holds the previous return value from a stateful Pregel graph.
 
 # --- Other constants ---
 PUSH = sys.intern("__pregel_push")

libs/langgraph/langgraph/errors.py

@@ -107,18 +107,3 @@ class CheckpointNotLatest(Exception):
     """Raised when the checkpoint is not the latest version (for distributed mode)."""
 
     pass
-
-
-class MultipleSubgraphsError(Exception):
-    """Raised when multiple subgraphs are called inside the same node.
-
-    Troubleshooting guides:
-
-    - [MULTIPLE_SUBGRAPHS](https://python.langchain.com/docs/troubleshooting/errors/MULTIPLE_SUBGRAPHS)
-    """
-
-    pass
-
-
-_SEEN_CHECKPOINT_NS: set[str] = set()
-"""Used for subgraph detection."""

libs/langgraph/langgraph/func/__init__.py

@@ -110,23 +110,99 @@ def entrypoint(
     *,
     checkpointer: Optional[BaseCheckpointSaver] = None,
     store: Optional[BaseStore] = None,
+    config_schema: Optional[type[Any]] = None,
 ) -> Callable[[types.FunctionType], Pregel]:
     def _imp(func: types.FunctionType) -> Pregel:
+        """Convert a function into a Pregel graph.
+
+        Args:
+            func: The function to convert. Support both sync and async functions, as well
+                   as generator and async generator functions.
+
+        Returns:
+            A Pregel graph.
+        """
+        # wrap generators in a function that writes to StreamWriter
         if inspect.isgeneratorfunction(func):
+            original_sig = inspect.signature(func)
+            # Check if original signature has a writer argument with a matching type.
+            # If not, we'll inject it into the decorator, but not pass it
+            # to the wrapped function.
+            if "writer" in original_sig.parameters:
 
-            def gen_wrapper(*args: Any, writer: StreamWriter, **kwargs: Any) -> Any:
-                for chunk in func(*args, **kwargs):
-                    writer(chunk)
+                @functools.wraps(func)
+                def gen_wrapper(*args: Any, writer: StreamWriter, **kwargs: Any) -> Any:
+                    chunks = []
+                    for chunk in func(*args, writer=writer, **kwargs):
+                        writer(chunk)
+                        chunks.append(chunk)
+                    return chunks
+            else:
 
+                @functools.wraps(func)
+                def gen_wrapper(*args: Any, writer: StreamWriter, **kwargs: Any) -> Any:
+                    chunks = []
+                    # Do not pass the writer argument to the wrapped function
+                    # as it does not have a matching parameter
+                    for chunk in func(*args, **kwargs):
+                        writer(chunk)
+                        chunks.append(chunk)
+                    return chunks
+
+                # Create a new parameter for the writer argument
+                extra_param = inspect.Parameter(
+                    "writer",
+                    inspect.Parameter.KEYWORD_ONLY,
+                    # The extra argument is a keyword-only argument
+                    default=lambda _: None,
+                )
+                # Update the function's signature to include the extra argument
+                new_params = list(original_sig.parameters.values()) + [extra_param]
+                new_sig = original_sig.replace(parameters=new_params)
+                # Update the signature of the wrapper function
+                gen_wrapper.__signature__ = new_sig  # type: ignore
             bound = get_runnable_for_func(gen_wrapper)
             stream_mode: StreamMode = "custom"
         elif inspect.isasyncgenfunction(func):
+            original_sig = inspect.signature(func)
+            # Check if original signature has a writer argument with a matching type.
+            # If not, we'll inject it into the decorator, but not pass it
+            # to the wrapped function.
+            if "writer" in original_sig.parameters:
 
-            async def agen_wrapper(
-                *args: Any, writer: StreamWriter, **kwargs: Any
-            ) -> Any:
-                async for chunk in func(*args, **kwargs):
-                    writer(chunk)
+                @functools.wraps(func)
+                async def agen_wrapper(
+                    *args: Any, writer: StreamWriter, **kwargs: Any
+                ) -> Any:
+                    chunks = []
+                    async for chunk in func(*args, writer=writer, **kwargs):
+                        writer(chunk)
+                        chunks.append(chunk)
+                    return chunks
+            else:
+
+                @functools.wraps(func)
+                async def agen_wrapper(
+                    *args: Any, writer: StreamWriter, **kwargs: Any
+                ) -> Any:
+                    chunks = []
+                    async for chunk in func(*args, **kwargs):
+                        writer(chunk)
+                        chunks.append(chunk)
+                    return chunks
+
+                # Create a new parameter for the writer argument
+                extra_param = inspect.Parameter(
+                    "writer",
+                    inspect.Parameter.KEYWORD_ONLY,
+                    # The extra argument is a keyword-only argument
+                    default=lambda _: None,
+                )
+                # Update the function's signature to include the extra argument
+                new_params = list(original_sig.parameters.values()) + [extra_param]
+                new_sig = original_sig.replace(parameters=new_params)
+                # Update the signature of the wrapper function
+                agen_wrapper.__signature__ = new_sig  # type: ignore
 
             bound = get_runnable_for_func(agen_wrapper)
             stream_mode = "custom"
@@ -134,6 +210,23 @@ def entrypoint(
             bound = get_runnable_for_func(func)
             stream_mode = "updates"
 
+        # get input and output types
+        sig = inspect.signature(func)
+        first_parameter_name = next(iter(sig.parameters.keys()), None)
+        if not first_parameter_name:
+            raise ValueError("Entrypoint function must have at least one parameter")
+        input_type = (
+            sig.parameters[first_parameter_name].annotation
+            if sig.parameters[first_parameter_name].annotation
+            is not inspect.Signature.empty
+            else Any
+        )
+        output_type = (
+            sig.return_annotation
+            if sig.return_annotation is not inspect.Signature.empty
+            else Any
+        )
+
         return Pregel(
             nodes={
                 func.__name__: PregelNode(
@@ -143,13 +236,18 @@ def entrypoint(
                     writers=[ChannelWrite([ChannelWriteEntry(END)], tags=[TAG_HIDDEN])],
                 )
             },
-            channels={START: EphemeralValue(Any), END: LastValue(Any, END)},
+            channels={
+                START: EphemeralValue(input_type),
+                END: LastValue(output_type, END),
+            },
             input_channels=START,
             output_channels=END,
             stream_channels=END,
             stream_mode=stream_mode,
+            stream_eager=True,
             checkpointer=checkpointer,
             store=store,
+            config_type=config_schema,
         )
 
     return _imp

libs/langgraph/langgraph/graph/state.py

@@ -847,14 +847,10 @@ def _control_branch(value: Any) -> Sequence[Union[str, Send]]:
     commands: list[Command] = []
     if isinstance(value, Command):
         commands.append(value)
-    elif (
-        isinstance(value, (list, tuple))
-        and value
-        and all(isinstance(i, Command) for i in value)
-    ):
-        commands.extend(value)
-    else:
-        return EMPTY_SEQ
+    elif isinstance(value, (list, tuple)):
+        for cmd in value:
+            if isinstance(cmd, Command):
+                commands.append(cmd)
     rtn: list[Union[str, Send]] = []
     for command in commands:
         if command.graph == Command.PARENT:
@@ -874,14 +870,10 @@ async def _acontrol_branch(value: Any) -> Sequence[Union[str, Send]]:
     commands: list[Command] = []
     if isinstance(value, Command):
         commands.append(value)
-    elif (
-        isinstance(value, (list, tuple))
-        and value
-        and all(isinstance(i, Command) for i in value)
-    ):
-        commands.extend(value)
-    else:
-        return EMPTY_SEQ
+    elif isinstance(value, (list, tuple)):
+        for cmd in value:
+            if isinstance(cmd, Command):
+                commands.append(cmd)
     rtn: list[Union[str, Send]] = []
     for command in commands:
         if command.graph == Command.PARENT:

libs/langgraph/langgraph/prebuilt/interrupt.py

@@ -0,0 +1,94 @@
+from typing import (
+    Literal,
+    Optional,
+    Union,
+)
+
+from typing_extensions import TypedDict
+
+
+class HumanInterruptConfig(TypedDict):
+    """Configuration that defines what actions are allowed for a human interrupt.
+
+    This controls the available interaction options when the graph is paused for human input.
+
+    Attributes:
+        allow_ignore: Whether the human can choose to ignore/skip the current step
+        allow_respond: Whether the human can provide a text response/feedback
+        allow_edit: Whether the human can edit the provided content/state
+        allow_accept: Whether the human can accept/approve the current state
+    """
+
+    allow_ignore: bool
+    allow_respond: bool
+    allow_edit: bool
+    allow_accept: bool
+
+
+class ActionRequest(TypedDict):
+    """Represents a request for human action within the graph execution.
+
+    Contains the action type and any associated arguments needed for the action.
+
+    Attributes:
+        action: The type or name of action being requested (e.g., "Approve XYZ action")
+        args: Key-value pairs of arguments needed for the action
+    """
+
+    action: str
+    args: dict
+
+
+class HumanInterrupt(TypedDict):
+    """Represents an interrupt triggered by the graph that requires human intervention.
+
+    This is passed to the `interrupt` function when execution is paused for human input.
+
+    Attributes:
+        action_request: The specific action being requested from the human
+        config: Configuration defining what actions are allowed
+        description: Optional detailed description of what input is needed
+
+    Example:
+        ```python
+        # Extract a tool call from the state and create an interrupt request
+        request = HumanInterrupt(
+            action_request=ActionRequest(
+                action="run_command",  # The action being requested
+                args={"command": "ls", "args": ["-l"]}  # Arguments for the action
+            ),
+            config=HumanInterruptConfig(
+                allow_ignore=True,    # Allow skipping this step
+                allow_respond=True,   # Allow text feedback
+                allow_edit=False,     # Don't allow editing
+                allow_accept=True     # Allow direct acceptance
+            ),
+            description="Please review the command before execution"
+        )
+        # Send the interrupt request and get the response
+        response = interrupt([request])[0]
+        ```
+    """
+
+    action_request: ActionRequest
+    config: HumanInterruptConfig
+    description: Optional[str]
+
+
+class HumanResponse(TypedDict):
+    """The response provided by a human to an interrupt, which is returned when graph execution resumes.
+
+    Attributes:
+        type: The type of response:
+            - "accept": Approves the current state without changes
+            - "ignore": Skips/ignores the current step
+            - "response": Provides text feedback or instructions
+            - "edit": Modifies the current state/content
+        arg: The response payload:
+            - None: For ignore/accept actions
+            - str: For text responses
+            - ActionRequest: For edit actions with updated content
+    """
+
+    type: Literal["accept", "ignore", "response", "edit"]
+    args: Union[None, str, ActionRequest]

libs/langgraph/langgraph/pregel/__init__.py

@@ -115,6 +115,7 @@ from langgraph.utils.config import (
     patch_checkpoint_map,
     patch_config,
     patch_configurable,
+    recast_checkpoint_ns,
 )
 from langgraph.utils.fields import get_enhanced_type_hints
 from langgraph.utils.pydantic import create_model
@@ -203,6 +204,10 @@ class Pregel(PregelProtocol):
     stream_mode: StreamMode = "values"
     """Mode to stream output, defaults to 'values'."""
 
+    stream_eager: bool = False
+    """Whether to force emitting stream events eagerly, automatically turned on
+    for stream_mode "messages" and "custom"."""
+
     output_channels: Union[str, Sequence[str]]
 
     stream_channels: Optional[Union[str, Sequence[str]]] = None
@@ -242,6 +247,7 @@ class Pregel(PregelProtocol):
         channels: Optional[dict[str, Union[BaseChannel, ManagedValueSpec]]],
         auto_validate: bool = True,
         stream_mode: StreamMode = "values",
+        stream_eager: bool = False,
         output_channels: Union[str, Sequence[str]],
         stream_channels: Optional[Union[str, Sequence[str]]] = None,
         interrupt_after_nodes: Union[All, Sequence[str]] = (),
@@ -259,6 +265,7 @@ class Pregel(PregelProtocol):
         self.nodes = nodes
         self.channels = channels or {}
         self.stream_mode = stream_mode
+        self.stream_eager = stream_eager
         self.output_channels = output_channels
         self.stream_channels = stream_channels
         self.interrupt_after_nodes = interrupt_after_nodes
@@ -494,7 +501,9 @@ class Pregel(PregelProtocol):
                 saved.metadata.get("step", -1) + 1,
                 for_execution=True,
                 store=self.store,
-                checkpointer=self.checkpointer or None,
+                checkpointer=self.checkpointer
+                if isinstance(self.checkpointer, BaseCheckpointSaver)
+                else None,
                 manager=None,
             )
             # get the subgraphs
@@ -606,7 +615,9 @@ class Pregel(PregelProtocol):
                 saved.metadata.get("step", -1) + 1,
                 for_execution=True,
                 store=self.store,
-                checkpointer=self.checkpointer or None,
+                checkpointer=self.checkpointer
+                if isinstance(self.checkpointer, BaseCheckpointSaver)
+                else None,
                 manager=None,
             )
             # get the subgraphs
@@ -690,19 +701,15 @@ class Pregel(PregelProtocol):
             checkpoint_ns := config[CONF].get(CONFIG_KEY_CHECKPOINT_NS, "")
         ) and CONFIG_KEY_CHECKPOINTER not in config[CONF]:
             # remove task_ids from checkpoint_ns
-            recast_checkpoint_ns = NS_SEP.join(
-                part.split(NS_END)[0] for part in checkpoint_ns.split(NS_SEP)
-            )
+            recast = recast_checkpoint_ns(checkpoint_ns)
             # find the subgraph with the matching name
-            for _, pregel in self.get_subgraphs(
-                namespace=recast_checkpoint_ns, recurse=True
-            ):
+            for _, pregel in self.get_subgraphs(namespace=recast, recurse=True):
                 return pregel.get_state(
                     patch_configurable(config, {CONFIG_KEY_CHECKPOINTER: checkpointer}),
                     subgraphs=subgraphs,
                 )
             else:
-                raise ValueError(f"Subgraph {recast_checkpoint_ns} not found")
+                raise ValueError(f"Subgraph {recast} not found")
 
         config = merge_configs(self.config, config) if self.config else config
         saved = checkpointer.get_tuple(config)
@@ -727,19 +734,15 @@ class Pregel(PregelProtocol):
             checkpoint_ns := config[CONF].get(CONFIG_KEY_CHECKPOINT_NS, "")
         ) and CONFIG_KEY_CHECKPOINTER not in config[CONF]:
             # remove task_ids from checkpoint_ns
-            recast_checkpoint_ns = NS_SEP.join(
-                part.split(NS_END)[0] for part in checkpoint_ns.split(NS_SEP)
-            )
+            recast = recast_checkpoint_ns(checkpoint_ns)
             # find the subgraph with the matching name
-            async for _, pregel in self.aget_subgraphs(
-                namespace=recast_checkpoint_ns, recurse=True
-            ):
+            async for _, pregel in self.aget_subgraphs(namespace=recast, recurse=True):
                 return await pregel.aget_state(
                     patch_configurable(config, {CONFIG_KEY_CHECKPOINTER: checkpointer}),
                     subgraphs=subgraphs,
                 )
             else:
-                raise ValueError(f"Subgraph {recast_checkpoint_ns} not found")
+                raise ValueError(f"Subgraph {recast} not found")
 
         config = merge_configs(self.config, config) if self.config else config
         saved = await checkpointer.aget_tuple(config)
@@ -770,13 +773,9 @@ class Pregel(PregelProtocol):
             checkpoint_ns := config[CONF].get(CONFIG_KEY_CHECKPOINT_NS, "")
         ) and CONFIG_KEY_CHECKPOINTER not in config[CONF]:
             # remove task_ids from checkpoint_ns
-            recast_checkpoint_ns = NS_SEP.join(
-                part.split(NS_END)[0] for part in checkpoint_ns.split(NS_SEP)
-            )
+            recast = recast_checkpoint_ns(checkpoint_ns)
             # find the subgraph with the matching name
-            for _, pregel in self.get_subgraphs(
-                namespace=recast_checkpoint_ns, recurse=True
-            ):
+            for _, pregel in self.get_subgraphs(namespace=recast, recurse=True):
                 yield from pregel.get_state_history(
                     patch_configurable(config, {CONFIG_KEY_CHECKPOINTER: checkpointer}),
                     filter=filter,
@@ -785,7 +784,7 @@ class Pregel(PregelProtocol):
                 )
                 return
             else:
-                raise ValueError(f"Subgraph {recast_checkpoint_ns} not found")
+                raise ValueError(f"Subgraph {recast} not found")
 
         config = merge_configs(
             self.config,
@@ -820,13 +819,9 @@ class Pregel(PregelProtocol):
             checkpoint_ns := config[CONF].get(CONFIG_KEY_CHECKPOINT_NS, "")
         ) and CONFIG_KEY_CHECKPOINTER not in config[CONF]:
             # remove task_ids from checkpoint_ns
-            recast_checkpoint_ns = NS_SEP.join(
-                part.split(NS_END)[0] for part in checkpoint_ns.split(NS_SEP)
-            )
+            recast = recast_checkpoint_ns(checkpoint_ns)
             # find the subgraph with the matching name
-            async for _, pregel in self.aget_subgraphs(
-                namespace=recast_checkpoint_ns, recurse=True
-            ):
+            async for _, pregel in self.aget_subgraphs(namespace=recast, recurse=True):
                 async for state in pregel.aget_state_history(
                     patch_configurable(config, {CONFIG_KEY_CHECKPOINTER: checkpointer}),
                     filter=filter,
@@ -836,7 +831,7 @@ class Pregel(PregelProtocol):
                     yield state
                 return
             else:
-                raise ValueError(f"Subgraph {recast_checkpoint_ns} not found")
+                raise ValueError(f"Subgraph {recast} not found")
 
         config = merge_configs(
             self.config,
@@ -875,20 +870,16 @@ class Pregel(PregelProtocol):
             checkpoint_ns := config[CONF].get(CONFIG_KEY_CHECKPOINT_NS, "")
         ) and CONFIG_KEY_CHECKPOINTER not in config[CONF]:
             # remove task_ids from checkpoint_ns
-            recast_checkpoint_ns = NS_SEP.join(
-                part.split(NS_END)[0] for part in checkpoint_ns.split(NS_SEP)
-            )
+            recast = recast_checkpoint_ns(checkpoint_ns)
             # find the subgraph with the matching name
-            for _, pregel in self.get_subgraphs(
-                namespace=recast_checkpoint_ns, recurse=True
-            ):
+            for _, pregel in self.get_subgraphs(namespace=recast, recurse=True):
                 return pregel.update_state(
                     patch_configurable(config, {CONFIG_KEY_CHECKPOINTER: checkpointer}),
                     values,
                     as_node,
                 )
             else:
-                raise ValueError(f"Subgraph {recast_checkpoint_ns} not found")
+                raise ValueError(f"Subgraph {recast} not found")
 
         # get last checkpoint
         config = ensure_config(self.config, config)
@@ -926,7 +917,9 @@ class Pregel(PregelProtocol):
                         saved.metadata.get("step", -1) + 1,
                         for_execution=True,
                         store=self.store,
-                        checkpointer=self.checkpointer or None,
+                        checkpointer=self.checkpointer
+                        if isinstance(self.checkpointer, BaseCheckpointSaver)
+                        else None,
                         manager=None,
                     )
                     # apply null writes
@@ -1020,7 +1013,9 @@ class Pregel(PregelProtocol):
                     saved.metadata.get("step", -1) + 1,
                     for_execution=True,
                     store=self.store,
-                    checkpointer=self.checkpointer or None,
+                    checkpointer=self.checkpointer
+                    if isinstance(self.checkpointer, BaseCheckpointSaver)
+                    else None,
                     manager=None,
                 )
                 # apply null writes
@@ -1155,20 +1150,16 @@ class Pregel(PregelProtocol):
             checkpoint_ns := config[CONF].get(CONFIG_KEY_CHECKPOINT_NS, "")
         ) and CONFIG_KEY_CHECKPOINTER not in config[CONF]:
             # remove task_ids from checkpoint_ns
-            recast_checkpoint_ns = NS_SEP.join(
-                part.split(NS_END)[0] for part in checkpoint_ns.split(NS_SEP)
-            )
+            recast = recast_checkpoint_ns(checkpoint_ns)
             # find the subgraph with the matching name
-            async for _, pregel in self.aget_subgraphs(
-                namespace=recast_checkpoint_ns, recurse=True
-            ):
+            async for _, pregel in self.aget_subgraphs(namespace=recast, recurse=True):
                 return await pregel.aupdate_state(
                     patch_configurable(config, {CONFIG_KEY_CHECKPOINTER: checkpointer}),
                     values,
                     as_node,
                 )
             else:
-                raise ValueError(f"Subgraph {recast_checkpoint_ns} not found")
+                raise ValueError(f"Subgraph {recast} not found")
 
         # get last checkpoint
         config = ensure_config(self.config, config)
@@ -1209,7 +1200,9 @@ class Pregel(PregelProtocol):
                         saved.metadata.get("step", -1) + 1,
                         for_execution=True,
                         store=self.store,
-                        checkpointer=self.checkpointer or None,
+                        checkpointer=self.checkpointer
+                        if isinstance(self.checkpointer, BaseCheckpointSaver)
+                        else None,
                         manager=None,
                     )
                     # apply null writes
@@ -1303,7 +1296,9 @@ class Pregel(PregelProtocol):
                     saved.metadata.get("step", -1) + 1,
                     for_execution=True,
                     store=self.store,
-                    checkpointer=self.checkpointer or None,
+                    checkpointer=self.checkpointer
+                    if isinstance(self.checkpointer, BaseCheckpointSaver)
+                    else None,
                     manager=None,
                 )
                 # apply null writes
@@ -1455,6 +1450,8 @@ class Pregel(PregelProtocol):
             checkpointer: Optional[BaseCheckpointSaver] = None
         elif CONFIG_KEY_CHECKPOINTER in config.get(CONF, {}):
             checkpointer = config[CONF][CONFIG_KEY_CHECKPOINTER]
+        elif self.checkpointer is True:
+            raise RuntimeError("checkpointer=True cannot be used for root graphs.")
         else:
             checkpointer = self.checkpointer
         if checkpointer and not config.get(CONF):
@@ -1598,6 +1595,12 @@ class Pregel(PregelProtocol):
                 interrupt_after=interrupt_after,
                 debug=debug,
             )
+            # set up subgraph checkpointing
+            if self.checkpointer is True:
+                ns = cast(str, config[CONF][CONFIG_KEY_CHECKPOINT_NS])
+                config[CONF][CONFIG_KEY_CHECKPOINT_NS] = NS_SEP.join(
+                    part.split(NS_END)[0] for part in ns.split(NS_SEP)
+                )
             # set up messages stream mode
             if "messages" in stream_modes:
                 run_manager.inheritable_handlers.append(
@@ -1634,7 +1637,12 @@ class Pregel(PregelProtocol):
                 if subgraphs:
                     loop.config[CONF][CONFIG_KEY_STREAM] = loop.stream
                 # enable concurrent streaming
-                if subgraphs or "messages" in stream_modes or "custom" in stream_modes:
+                if (
+                    self.stream_eager
+                    or subgraphs
+                    or "messages" in stream_modes
+                    or "custom" in stream_modes
+                ):
                     # we are careful to have a single waiter live at any one time
                     # because on exit we increment semaphore count by exactly 1
                     waiter: Optional[concurrent.futures.Future] = None
@@ -1864,7 +1872,12 @@ class Pregel(PregelProtocol):
                         stream_put, stream_modes
                     )
                 # enable concurrent streaming
-                if subgraphs or "messages" in stream_modes or "custom" in stream_modes:
+                if (
+                    self.stream_eager
+                    or subgraphs
+                    or "messages" in stream_modes
+                    or "custom" in stream_modes
+                ):
 
                     def get_waiter() -> asyncio.Task[None]:
                         return aioloop.create_task(stream.wait())

libs/langgraph/langgraph/pregel/algo.py

@@ -37,15 +37,16 @@ from langgraph.constants import (
     CONFIG_KEY_CHECKPOINT_MAP,
     CONFIG_KEY_CHECKPOINT_NS,
     CONFIG_KEY_CHECKPOINTER,
+    CONFIG_KEY_END,
     CONFIG_KEY_READ,
     CONFIG_KEY_SCRATCHPAD,
     CONFIG_KEY_SEND,
     CONFIG_KEY_STORE,
     CONFIG_KEY_TASK_ID,
-    CONFIG_KEY_WRITES,
     EMPTY_SEQ,
     ERROR,
     INTERRUPT,
+    MISSING,
     NO_WRITES,
     NS_END,
     NS_SEP,
@@ -71,6 +72,7 @@ from langgraph.types import (
     All,
     LoopProtocol,
     PregelExecutableTask,
+    PregelScratchpad,
     PregelTask,
     RetryPolicy,
 )
@@ -502,13 +504,13 @@ def prepare_single_task(
                         },
                         CONFIG_KEY_CHECKPOINT_ID: None,
                         CONFIG_KEY_CHECKPOINT_NS: task_checkpoint_ns,
-                        CONFIG_KEY_WRITES: [
-                            w
-                            for w in pending_writes
-                            + configurable.get(CONFIG_KEY_WRITES, [])
-                            if w[0] in (NULL_TASK_ID, task_id)
-                        ],
-                        CONFIG_KEY_SCRATCHPAD: {},
+                        CONFIG_KEY_SCRATCHPAD: _scratchpad(
+                            pending_writes,
+                            task_id,
+                        ),
+                        CONFIG_KEY_END: checkpoint["channel_values"].get(
+                            "__end__", None
+                        ),
                     },
                 ),
                 triggers,
@@ -614,13 +616,13 @@ def prepare_single_task(
                             },
                             CONFIG_KEY_CHECKPOINT_ID: None,
                             CONFIG_KEY_CHECKPOINT_NS: task_checkpoint_ns,
-                            CONFIG_KEY_WRITES: [
-                                w
-                                for w in pending_writes
-                                + configurable.get(CONFIG_KEY_WRITES, [])
-                                if w[0] in (NULL_TASK_ID, task_id)
-                            ],
-                            CONFIG_KEY_SCRATCHPAD: {},
+                            CONFIG_KEY_SCRATCHPAD: _scratchpad(
+                                pending_writes,
+                                task_id,
+                            ),
+                            CONFIG_KEY_END: checkpoint["channel_values"].get(
+                                "__end__", None
+                            ),
                         },
                     ),
                     triggers,
@@ -685,7 +687,7 @@ def prepare_single_task(
                 "langgraph_checkpoint_ns": task_checkpoint_ns,
             }
             if task_id_checksum is not None:
-                assert task_id == task_id_checksum
+                assert task_id == task_id_checksum, f"{task_id} != {task_id_checksum}"
             if for_execution:
                 if node := proc.node:
                     if proc.metadata:
@@ -738,13 +740,13 @@ def prepare_single_task(
                                 },
                                 CONFIG_KEY_CHECKPOINT_ID: None,
                                 CONFIG_KEY_CHECKPOINT_NS: task_checkpoint_ns,
-                                CONFIG_KEY_WRITES: [
-                                    w
-                                    for w in pending_writes
-                                    + configurable.get(CONFIG_KEY_WRITES, [])
-                                    if w[0] in (NULL_TASK_ID, task_id)
-                                ],
-                                CONFIG_KEY_SCRATCHPAD: {},
+                                CONFIG_KEY_SCRATCHPAD: _scratchpad(
+                                    pending_writes,
+                                    task_id,
+                                ),
+                                CONFIG_KEY_END: checkpoint["channel_values"].get(
+                                    "__end__", None
+                                ),
                             },
                         ),
                         triggers,
@@ -758,6 +760,27 @@ def prepare_single_task(
                 return PregelTask(task_id, name, task_path[:3])
 
 
+def _scratchpad(
+    pending_writes: Sequence[PendingWrite],
+    task_id: str,
+) -> PregelScratchpad:
+    return PregelScratchpad(
+        # call
+        call_counter=0,
+        # interrupt
+        interrupt_counter=-1,
+        resume=next(
+            (w[2] for w in pending_writes if w[0] == task_id and w[1] == RESUME), []
+        ),
+        null_resume=next(
+            (w[2] for w in pending_writes if w[0] == NULL_TASK_ID and w[1] == RESUME),
+            MISSING,
+        ),
+        # subgraph
+        subgraph_counter=0,
+    )
+
+
 def _proc_input(
     proc: PregelNode,
     managed: ManagedValueMapping,

libs/langgraph/langgraph/pregel/loop.py

@@ -47,12 +47,14 @@ from langgraph.constants import (
     CONFIG_KEY_DELEGATE,
     CONFIG_KEY_ENSURE_LATEST,
     CONFIG_KEY_RESUMING,
+    CONFIG_KEY_SCRATCHPAD,
     CONFIG_KEY_STREAM,
     CONFIG_KEY_TASK_ID,
     EMPTY_SEQ,
     ERROR,
     INPUT,
     INTERRUPT,
+    MISSING,
     NS_SEP,
     NULL_TASK_ID,
     PUSH,
@@ -61,12 +63,10 @@ from langgraph.constants import (
     TAG_HIDDEN,
 )
 from langgraph.errors import (
-    _SEEN_CHECKPOINT_NS,
     CheckpointNotLatest,
     EmptyInputError,
     GraphDelegate,
     GraphInterrupt,
-    MultipleSubgraphsError,
 )
 from langgraph.managed.base import (
     ManagedValueMapping,
@@ -114,6 +114,7 @@ from langgraph.types import (
     Command,
     LoopProtocol,
     PregelExecutableTask,
+    PregelScratchpad,
     StreamChunk,
     StreamProtocol,
 )
@@ -201,7 +202,6 @@ class PregelLoop(LoopProtocol):
         interrupt_after: Union[All, Sequence[str]] = EMPTY_SEQ,
         interrupt_before: Union[All, Sequence[str]] = EMPTY_SEQ,
         manager: Union[None, AsyncParentRunManager, ParentRunManager] = None,
-        check_subgraphs: bool = True,
         debug: bool = False,
     ) -> None:
         super().__init__(
@@ -228,20 +228,26 @@ class PregelLoop(LoopProtocol):
         self.debug = debug
         if self.stream is not None and CONFIG_KEY_STREAM in config[CONF]:
             self.stream = DuplexStream(self.stream, config[CONF][CONFIG_KEY_STREAM])
+        scratchpad: Optional[PregelScratchpad] = config[CONF].get(CONFIG_KEY_SCRATCHPAD)
+        if not self.config[CONF].get(CONFIG_KEY_DELEGATE) and scratchpad is not None:
+            if scratchpad["subgraph_counter"]:
+                self.config = patch_configurable(
+                    self.config,
+                    {
+                        CONFIG_KEY_CHECKPOINT_NS: NS_SEP.join(
+                            (
+                                config[CONF][CONFIG_KEY_CHECKPOINT_NS],
+                                str(scratchpad["subgraph_counter"]),
+                            )
+                        )
+                    },
+                )
+            scratchpad["subgraph_counter"] += 1
         if not self.is_nested and config[CONF].get(CONFIG_KEY_CHECKPOINT_NS):
             self.config = patch_configurable(
                 self.config,
                 {CONFIG_KEY_CHECKPOINT_NS: "", CONFIG_KEY_CHECKPOINT_ID: None},
             )
-        if check_subgraphs and self.is_nested and self.checkpointer is not None:
-            if self.config[CONF][CONFIG_KEY_CHECKPOINT_NS] in _SEEN_CHECKPOINT_NS:
-                raise MultipleSubgraphsError(
-                    "Multiple subgraphs called inside the same node\n\n"
-                    "Troubleshooting URL: https://python.langchain.com/docs"
-                    "/troubleshooting/errors/MULTIPLE_SUBGRAPHS/"
-                )
-            else:
-                _SEEN_CHECKPOINT_NS.add(self.config[CONF][CONFIG_KEY_CHECKPOINT_NS])
         if (
             CONFIG_KEY_CHECKPOINT_MAP in self.config[CONF]
             and self.config[CONF].get(CONFIG_KEY_CHECKPOINT_NS)
@@ -339,11 +345,11 @@ class PregelLoop(LoopProtocol):
                 (PUSH, task.path, write_idx, task.id, call),
                 None,
                 checkpoint=self.checkpoint,
-                pending_writes=[(task.id, *w) for w in task.writes],
+                pending_writes=self.checkpoint_pending_writes,
                 processes=self.nodes,
                 channels=self.channels,
                 managed=self.managed,
-                config=self.config,
+                config=task.config,
                 step=self.step,
                 for_execution=True,
                 store=self.store,
@@ -556,8 +562,16 @@ class PregelLoop(LoopProtocol):
             )
         )
 
+        # take resume value from parent
+        if scratchpad := configurable.get(CONFIG_KEY_SCRATCHPAD):
+            if scratchpad["null_resume"] is not MISSING:
+                self.put_writes(NULL_TASK_ID, [(RESUME, scratchpad["null_resume"])])
         # map command to writes
         if isinstance(self.input, Command):
+            if self.input.resume is not None and not self.checkpointer:
+                raise RuntimeError(
+                    "Cannot use Command(resume=...) without checkpointer"
+                )
             writes: defaultdict[str, list[tuple[str, Any]]] = defaultdict(list)
             # group writes by task ID
             for tid, c, v in map_command(self.input, self.checkpoint_pending_writes):
@@ -807,7 +821,6 @@ class SyncPregelLoop(PregelLoop, ContextManager):
         interrupt_before: Union[All, Sequence[str]] = EMPTY_SEQ,
         output_keys: Union[str, Sequence[str]] = EMPTY_SEQ,
         stream_keys: Union[str, Sequence[str]] = EMPTY_SEQ,
-        check_subgraphs: bool = True,
         debug: bool = False,
     ) -> None:
         super().__init__(
@@ -822,7 +835,6 @@ class SyncPregelLoop(PregelLoop, ContextManager):
             stream_keys=stream_keys,
             interrupt_after=interrupt_after,
             interrupt_before=interrupt_before,
-            check_subgraphs=check_subgraphs,
             manager=manager,
             debug=debug,
         )
@@ -944,7 +956,6 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager):
         manager: Union[None, AsyncParentRunManager, ParentRunManager] = None,
         output_keys: Union[str, Sequence[str]] = EMPTY_SEQ,
         stream_keys: Union[str, Sequence[str]] = EMPTY_SEQ,
-        check_subgraphs: bool = True,
         debug: bool = False,
     ) -> None:
         super().__init__(
@@ -959,7 +970,6 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager):
             stream_keys=stream_keys,
             interrupt_after=interrupt_after,
             interrupt_before=interrupt_before,
-            check_subgraphs=check_subgraphs,
             manager=manager,
             debug=debug,
         )

libs/langgraph/langgraph/pregel/retry.py

@@ -12,7 +12,7 @@ from langgraph.constants import (
     CONFIG_KEY_RESUMING,
     NS_SEP,
 )
-from langgraph.errors import _SEEN_CHECKPOINT_NS, GraphBubbleUp, ParentCommand
+from langgraph.errors import GraphBubbleUp, ParentCommand
 from langgraph.types import Command, PregelExecutableTask, RetryPolicy
 from langgraph.utils.config import patch_configurable
 
@@ -48,7 +48,10 @@ def run_with_retry(
                 break
             elif cmd.graph == Command.PARENT:
                 # this command is for the parent graph, assign it to the parent
-                parent_ns = NS_SEP.join(ns.split(NS_SEP)[:-1])
+                parts = ns.split(NS_SEP)
+                if parts[-1].isdigit():
+                    parts.pop()
+                parent_ns = NS_SEP.join(parts[:-1])
                 exc.args = (replace(cmd, graph=parent_ns),)
             # bubble up
             raise
@@ -96,13 +99,6 @@ def run_with_retry(
             )
             # signal subgraphs to resume (if available)
             config = patch_configurable(config, {CONFIG_KEY_RESUMING: True})
-            # clear checkpoint_ns seen (for subgraph detection)
-            if checkpoint_ns := config[CONF].get(CONFIG_KEY_CHECKPOINT_NS):
-                _SEEN_CHECKPOINT_NS.discard(checkpoint_ns)
-        finally:
-            # clear checkpoint_ns seen (for subgraph detection)
-            if checkpoint_ns := config[CONF].get(CONFIG_KEY_CHECKPOINT_NS):
-                _SEEN_CHECKPOINT_NS.discard(checkpoint_ns)
 
 
 async def arun_with_retry(
@@ -140,7 +136,10 @@ async def arun_with_retry(
                 break
             elif cmd.graph == Command.PARENT:
                 # this command is for the parent graph, assign it to the parent
-                parent_ns = NS_SEP.join(ns.split(NS_SEP)[:-1])
+                parts = ns.split(NS_SEP)
+                if parts[-1].isdigit():
+                    parts.pop()
+                parent_ns = NS_SEP.join(parts[:-1])
                 exc.args = (replace(cmd, graph=parent_ns),)
             # bubble up
             raise
@@ -188,10 +187,3 @@ async def arun_with_retry(
             )
             # signal subgraphs to resume (if available)
             config = patch_configurable(config, {CONFIG_KEY_RESUMING: True})
-            # clear checkpoint_ns seen (for subgraph detection)
-            if checkpoint_ns := config[CONF].get(CONFIG_KEY_CHECKPOINT_NS):
-                _SEEN_CHECKPOINT_NS.discard(checkpoint_ns)
-        finally:
-            # clear checkpoint_ns seen (for subgraph detection)
-            if checkpoint_ns := config[CONF].get(CONFIG_KEY_CHECKPOINT_NS):
-                _SEEN_CHECKPOINT_NS.discard(checkpoint_ns)

libs/langgraph/langgraph/pregel/runner.py

@@ -107,12 +107,11 @@ class PregelRunner:
                     elif next_task.writes:
                         # if it already ran, return the result
                         fut = concurrent.futures.Future()
-                        if (
-                            val := next(
-                                (v for c, v in next_task.writes if c == RETURN), MISSING
-                            )
-                        ) and val is not MISSING:
-                            fut.set_result(val)
+                        ret = next(
+                            (v for c, v in next_task.writes if c == RETURN), MISSING
+                        )
+                        if ret is not MISSING:
+                            fut.set_result(ret)
                         elif exc := next(
                             (v for c, v in next_task.writes if c == ERROR), None
                         ):
@@ -295,12 +294,11 @@ class PregelRunner:
                     elif next_task.writes:
                         # if it already ran, return the result
                         fut = asyncio.Future()
-                        if (
-                            val := next(
-                                (v for c, v in next_task.writes if c == RETURN), MISSING
-                            )
-                        ) and val is not MISSING:
-                            fut.set_result(val)
+                        ret = next(
+                            (v for c, v in next_task.writes if c == RETURN), MISSING
+                        )
+                        if ret is not MISSING:
+                            fut.set_result(ret)
                         elif exc := next(
                             (v for c, v in next_task.writes if c == ERROR), None
                         ):

libs/langgraph/langgraph/types.py

@@ -21,11 +21,7 @@ from typing import (
 from langchain_core.runnables import Runnable, RunnableConfig
 from typing_extensions import Self, TypedDict
 
-from langgraph.checkpoint.base import (
-    BaseCheckpointSaver,
-    CheckpointMetadata,
-    PendingWrite,
-)
+from langgraph.checkpoint.base import BaseCheckpointSaver, CheckpointMetadata
 
 if TYPE_CHECKING:
     from langgraph.store.base import BaseStore
@@ -42,9 +38,11 @@ except ImportError:
 All = Literal["*"]
 """Special value to indicate that graph should interrupt on all nodes."""
 
-Checkpointer = Union[None, Literal[False], BaseCheckpointSaver]
-"""Type of the checkpointer to use for a subgraph. False disables checkpointing,
-even if the parent graph has a checkpointer. None inherits checkpointer."""
+Checkpointer = Union[None, bool, BaseCheckpointSaver]
+"""Type of the checkpointer to use for a subgraph.
+- True enables persistent checkpointing for this subgraph.
+- False disables checkpointing, even if the parent graph has a checkpointer.
+- None inherits checkpointer from the parent graph."""
 
 StreamMode = Literal["values", "updates", "debug", "messages", "custom"]
 """How the stream method should emit outputs.
@@ -341,13 +339,15 @@ class LoopProtocol:
         self.stop = stop
 
 
-class PregelScratchpad(TypedDict, total=False):
-    # interrupt
-    interrupt_counter: int
-    used_null_resume: bool
-    resume: list[Any]
+class PregelScratchpad(TypedDict):
     # call
     call_counter: int
+    # interrupt
+    interrupt_counter: int
+    resume: list[Any]
+    null_resume: Any
+    # subgraph
+    subgraph_counter: int
 
 
 def interrupt(value: Any) -> Any:
@@ -449,10 +449,8 @@ def interrupt(value: Any) -> Any:
         CONFIG_KEY_CHECKPOINT_NS,
         CONFIG_KEY_SCRATCHPAD,
         CONFIG_KEY_SEND,
-        CONFIG_KEY_TASK_ID,
-        CONFIG_KEY_WRITES,
+        MISSING,
         NS_SEP,
-        NULL_TASK_ID,
         RESUME,
     )
     from langgraph.errors import GraphInterrupt
@@ -461,29 +459,20 @@ def interrupt(value: Any) -> Any:
     conf = get_config()["configurable"]
     # track interrupt index
     scratchpad: PregelScratchpad = conf[CONFIG_KEY_SCRATCHPAD]
-    if "interrupt_counter" not in scratchpad:
-        scratchpad["interrupt_counter"] = 0
-    else:
-        scratchpad["interrupt_counter"] += 1
+    scratchpad["interrupt_counter"] += 1
     idx = scratchpad["interrupt_counter"]
     # find previous resume values
-    task_id = conf[CONFIG_KEY_TASK_ID]
-    writes: list[PendingWrite] = conf[CONFIG_KEY_WRITES]
-    scratchpad.setdefault(
-        "resume", next((w[2] for w in writes if w[0] == task_id and w[1] == RESUME), [])
-    )
     if scratchpad["resume"]:
         if idx < len(scratchpad["resume"]):
             return scratchpad["resume"][idx]
     # find current resume value
-    if not scratchpad.get("used_null_resume"):
-        scratchpad["used_null_resume"] = True
-        for tid, c, v in sorted(writes, key=lambda x: x[0], reverse=True):
-            if tid == NULL_TASK_ID and c == RESUME:
-                assert len(scratchpad["resume"]) == idx, (scratchpad["resume"], idx)
-                scratchpad["resume"].append(v)
-                conf[CONFIG_KEY_SEND]([(RESUME, scratchpad["resume"])])
-                return v
+    if scratchpad["null_resume"] is not MISSING:
+        assert len(scratchpad["resume"]) == idx, (scratchpad["resume"], idx)
+        v = scratchpad["null_resume"]
+        scratchpad["null_resume"] = MISSING
+        scratchpad["resume"].append(v)
+        conf[CONFIG_KEY_SEND]([(RESUME, scratchpad["resume"])])
+        return v
     # no resume value found
     raise GraphInterrupt(
         (

libs/langgraph/langgraph/utils/config.py

@@ -23,9 +23,25 @@ from langgraph.constants import (
     CONFIG_KEY_CHECKPOINT_ID,
     CONFIG_KEY_CHECKPOINT_MAP,
     CONFIG_KEY_CHECKPOINT_NS,
+    NS_END,
+    NS_SEP,
 )
 
 
+def recast_checkpoint_ns(ns: str) -> str:
+    """Remove task IDs from checkpoint namespace.
+
+    Args:
+        ns (str): The checkpoint namespace with task IDs.
+
+    Returns:
+        str: The checkpoint namespace without task IDs.
+    """
+    return NS_SEP.join(
+        part.split(NS_END)[0] for part in ns.split(NS_SEP) if not part.isdigit()
+    )
+
+
 def patch_configurable(
     config: Optional[RunnableConfig], patch: dict[str, Any]
 ) -> RunnableConfig:

libs/langgraph/langgraph/utils/runnable.py

@@ -34,7 +34,12 @@ from langchain_core.runnables.utils import Input
 from langchain_core.tracers._streaming import _StreamingCallbackHandler
 from typing_extensions import TypeGuard
 
-from langgraph.constants import CONF, CONFIG_KEY_STORE, CONFIG_KEY_STREAM_WRITER
+from langgraph.constants import (
+    CONF,
+    CONFIG_KEY_END,
+    CONFIG_KEY_STORE,
+    CONFIG_KEY_STREAM_WRITER,
+)
 from langgraph.store.base import BaseStore
 from langgraph.types import StreamWriter
 from langgraph.utils.config import (
@@ -58,6 +63,10 @@ class StrEnum(str, enum.Enum):
     """A string enum."""
 
 
+# Special type to denote any type is accepted
+ANY_TYPE = object()
+
+
 ASYNCIO_ACCEPTS_CONTEXT = sys.version_info >= (3, 11)
 
 KWARGS_CONFIG_KEYS: tuple[tuple[str, tuple[Any, ...], str, Any], ...] = (
@@ -73,6 +82,12 @@ KWARGS_CONFIG_KEYS: tuple[tuple[str, tuple[Any, ...], str, Any], ...] = (
         CONFIG_KEY_STORE,
         inspect.Parameter.empty,
     ),
+    (
+        sys.intern("previous"),
+        (ANY_TYPE,),
+        CONFIG_KEY_END,
+        inspect.Parameter.empty,
+    ),
 )
 """List of kwargs that can be passed to functions, and their corresponding
 config keys, default values and type annotations.
@@ -135,9 +150,12 @@ class RunnableCallable(Runnable):
         self.func_accepts: dict[str, bool] = {}
         for kw, typ, _, _ in KWARGS_CONFIG_KEYS:
             p = params.get(kw)
-            self.func_accepts[kw] = (
-                p is not None and p.annotation in typ and p.kind in VALID_KINDS
-            )
+            if typ == (ANY_TYPE,):
+                self.func_accepts[kw] = p is not None and p.kind in VALID_KINDS
+            else:
+                self.func_accepts[kw] = (
+                    p is not None and p.annotation in typ and p.kind in VALID_KINDS
+                )
 
     def __repr__(self) -> str:
         repr_args = {
@@ -162,16 +180,20 @@ class RunnableCallable(Runnable):
         if self.func_accepts_config:
             kwargs["config"] = config
         _conf = config[CONF]
-        for kw, _, ck, defv in KWARGS_CONFIG_KEYS:
+        for kw, _, config_key, default_value in KWARGS_CONFIG_KEYS:
             if not self.func_accepts[kw]:
                 continue
 
-            if defv is inspect.Parameter.empty and kw not in kwargs and ck not in _conf:
+            if (
+                default_value is inspect.Parameter.empty
+                and kw not in kwargs
+                and config_key not in _conf
+            ):
                 raise ValueError(
-                    f"Missing required config key '{ck}' for '{self.name}'."
+                    f"Missing required config key '{config_key}' for '{self.name}'."
                 )
             elif kwargs.get(kw) is None:
-                kwargs[kw] = _conf.get(ck, defv)
+                kwargs[kw] = _conf.get(config_key, default_value)
 
         context = copy_context()
         if self.trace:
@@ -210,16 +232,20 @@ class RunnableCallable(Runnable):
         if self.func_accepts_config:
             kwargs["config"] = config
         _conf = config[CONF]
-        for kw, _, ck, defv in KWARGS_CONFIG_KEYS:
+        for kw, _, config_key, default_value in KWARGS_CONFIG_KEYS:
             if not self.func_accepts[kw]:
                 continue
 
-            if defv is inspect.Parameter.empty and kw not in kwargs and ck not in _conf:
+            if (
+                default_value is inspect.Parameter.empty
+                and kw not in kwargs
+                and config_key not in _conf
+            ):
                 raise ValueError(
-                    f"Missing required config key '{ck}' for '{self.name}'."
+                    f"Missing required config key '{config_key}' for '{self.name}'."
                 )
             elif kwargs.get(kw) is None:
-                kwargs[kw] = _conf.get(ck, defv)
+                kwargs[kw] = _conf.get(config_key, default_value)
         context = copy_context()
         if self.trace:
             callback_manager = get_async_callback_manager_for_config(config, self.tags)

libs/langgraph/pyproject.toml

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

libs/langgraph/tests/test_pregel.py

@@ -51,7 +51,7 @@ from langgraph.checkpoint.base import (
 )
 from langgraph.checkpoint.memory import MemorySaver
 from langgraph.constants import CONFIG_KEY_NODE_FINISHED, ERROR, PULL, START
-from langgraph.errors import InvalidUpdateError, MultipleSubgraphsError
+from langgraph.errors import InvalidUpdateError
 from langgraph.func import entrypoint, task
 from langgraph.graph import END, Graph, StateGraph
 from langgraph.graph.message import MessageGraph, MessagesState, add_messages
@@ -1436,6 +1436,9 @@ def test_imp_task(request: pytest.FixtureRequest, checkpointer_name: str) -> Non
     checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}")
     mapper_calls = 0
 
+    class Config:
+        model: str
+
     @task()
     def mapper(input: int) -> str:
         nonlocal mapper_calls
@@ -1443,13 +1446,57 @@ def test_imp_task(request: pytest.FixtureRequest, checkpointer_name: str) -> Non
         time.sleep(input / 100)
         return str(input) * 2
 
-    @entrypoint(checkpointer=checkpointer)
+    @entrypoint(checkpointer=checkpointer, config_schema=Config)
     def graph(input: list[int]) -> list[str]:
         futures = [mapper(i) for i in input]
         mapped = [f.result() for f in futures]
         answer = interrupt("question")
         return [m + answer for m in mapped]
 
+    assert graph.get_input_jsonschema() == {
+        "type": "array",
+        "items": {"type": "integer"},
+        "title": "LangGraphInput",
+    }
+    assert graph.get_output_jsonschema() == {
+        "type": "array",
+        "items": {"type": "string"},
+        "title": "LangGraphOutput",
+    }
+    assert graph.get_config_jsonschema() == {
+        "$defs": {
+            "Configurable": {
+                "properties": {
+                    "model": {"default": None, "title": "Model", "type": "string"},
+                    "checkpoint_id": {
+                        "anyOf": [{"type": "string"}, {"type": "null"}],
+                        "default": None,
+                        "description": "Pass to fetch a past checkpoint. If None, fetches the latest checkpoint.",
+                        "title": "Checkpoint ID",
+                    },
+                    "checkpoint_ns": {
+                        "default": "",
+                        "description": 'Checkpoint namespace. Denotes the path to the subgraph node the checkpoint originates from, separated by `|` character, e.g. `"child|grandchild"`. Defaults to "" (root graph).',
+                        "title": "Checkpoint NS",
+                        "type": "string",
+                    },
+                    "thread_id": {
+                        "default": "",
+                        "title": "Thread ID",
+                        "type": "string",
+                    },
+                },
+                "title": "Configurable",
+                "type": "object",
+            }
+        },
+        "properties": {
+            "configurable": {"$ref": "#/$defs/Configurable", "default": None}
+        },
+        "title": "LangGraphConfig",
+        "type": "object",
+    }
+
     thread1 = {"configurable": {"thread_id": "1"}}
     assert [*graph.stream([0, 1], thread1)] == [
         {"mapper": "00"},
@@ -1745,9 +1792,8 @@ def test_invoke_join_then_call_other_pregel(
 
     # add checkpointer
     app.checkpointer = checkpointer
-    # subgraph is called twice in the same node, through .map(), so raises
-    with pytest.raises(MultipleSubgraphsError):
-        app.invoke([2, 3], {"configurable": {"thread_id": "1"}})
+    # subgraph is called twice in the same node, but that works
+    assert app.invoke([2, 3], {"configurable": {"thread_id": "1"}}) == 27
 
     # set inner graph checkpointer NeverCheckpoint
     inner_app.checkpointer = False
@@ -2167,10 +2213,10 @@ def test_in_one_fan_out_state_graph_waiting_edge(
 
     @workflow.add_node
     def rewrite_query(data: State) -> State:
-        return {"query": f'query: {data["query"]}'}
+        return {"query": f"query: {data['query']}"}
 
     def analyzer_one(data: State) -> State:
-        return {"query": f'analyzed: {data["query"]}'}
+        return {"query": f"analyzed: {data['query']}"}
 
     def retriever_one(data: State) -> State:
         return {"docs": ["doc1", "doc2"]}
@@ -2307,10 +2353,10 @@ def test_in_one_fan_out_state_graph_waiting_edge_via_branch(
         docs: Annotated[list[str], sorted_add]
 
     def rewrite_query(data: State) -> State:
-        return {"query": f'query: {data["query"]}'}
+        return {"query": f"query: {data['query']}"}
 
     def analyzer_one(data: State) -> State:
-        return {"query": f'analyzed: {data["query"]}'}
+        return {"query": f"analyzed: {data['query']}"}
 
     def retriever_one(data: State) -> State:
         return {"docs": ["doc1", "doc2"]}
@@ -2740,11 +2786,11 @@ def test_in_one_fan_out_state_graph_waiting_edge_plus_regular(
         docs: Annotated[list[str], sorted_add]
 
     def rewrite_query(data: State) -> State:
-        return {"query": f'query: {data["query"]}'}
+        return {"query": f"query: {data['query']}"}
 
     def analyzer_one(data: State) -> State:
         time.sleep(0.1)
-        return {"query": f'analyzed: {data["query"]}'}
+        return {"query": f"analyzed: {data['query']}"}
 
     def retriever_one(data: State) -> State:
         return {"docs": ["doc1", "doc2"]}
@@ -2830,10 +2876,10 @@ def test_in_one_fan_out_state_graph_waiting_edge_multiple() -> None:
         docs: Annotated[list[str], sorted_add]
 
     def rewrite_query(data: State) -> State:
-        return {"query": f'query: {data["query"]}'}
+        return {"query": f"query: {data['query']}"}
 
     def analyzer_one(data: State) -> State:
-        return {"query": f'analyzed: {data["query"]}'}
+        return {"query": f"analyzed: {data['query']}"}
 
     def retriever_one(data: State) -> State:
         return {"docs": ["doc1", "doc2"]}
@@ -2903,10 +2949,10 @@ def test_callable_in_conditional_edges_with_no_path_map() -> None:
         query: str
 
     def rewrite(data: State) -> State:
-        return {"query": f'query: {data["query"]}'}
+        return {"query": f"query: {data['query']}"}
 
     def analyze(data: State) -> State:
-        return {"query": f'analyzed: {data["query"]}'}
+        return {"query": f"analyzed: {data['query']}"}
 
     class ChooseAnalyzer:
         def __call__(self, data: State) -> str:
@@ -2929,10 +2975,10 @@ def test_function_in_conditional_edges_with_no_path_map() -> None:
         query: str
 
     def rewrite(data: State) -> State:
-        return {"query": f'query: {data["query"]}'}
+        return {"query": f"query: {data['query']}"}
 
     def analyze(data: State) -> State:
-        return {"query": f'analyzed: {data["query"]}'}
+        return {"query": f"analyzed: {data['query']}"}
 
     def choose_analyzer(data: State) -> str:
         return "analyzer"
@@ -2965,13 +3011,13 @@ def test_in_one_fan_out_state_graph_waiting_edge_multiple_cond_edge() -> None:
         docs: Annotated[list[str], sorted_add]
 
     def rewrite_query(data: State) -> State:
-        return {"query": f'query: {data["query"]}'}
+        return {"query": f"query: {data['query']}"}
 
     def retriever_picker(data: State) -> list[str]:
         return ["analyzer_one", "retriever_two"]
 
     def analyzer_one(data: State) -> State:
-        return {"query": f'analyzed: {data["query"]}'}
+        return {"query": f"analyzed: {data['query']}"}
 
     def retriever_one(data: State) -> State:
         return {"docs": ["doc1", "doc2"]}
@@ -3189,6 +3235,66 @@ def test_nested_graph(snapshot: SnapshotAssertion) -> None:
     ]
 
 
+@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC)
+def test_subgraph_checkpoint_true(
+    request: pytest.FixtureRequest, checkpointer_name: str
+) -> None:
+    checkpointer = request.getfixturevalue("checkpointer_" + checkpointer_name)
+
+    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
+
+    graph = StateGraph(State)
+    graph.add_node("inner", inner.compile(checkpointer=True))
+    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=checkpointer)
+
+    config = {"configurable": {"thread_id": "2"}}
+    assert [c for c in app.stream({"my_key": ""}, config, subgraphs=True)] == [
+        (("inner",), {"inner_1": {"my_key": " got here", "my_other_key": ""}}),
+        (("inner",), {"inner_2": {"my_key": " and there"}}),
+        ((), {"inner": {"my_key": " got here and there"}}),
+        (
+            ("inner",),
+            {
+                "inner_1": {
+                    "my_key": " got here",
+                    "my_other_key": " got here and there got here and there",
+                }
+            },
+        ),
+        (("inner",), {"inner_2": {"my_key": " and there"}}),
+        (
+            (),
+            {
+                "inner": {
+                    "my_key": " got here and there got here and there got here and there"
+                }
+            },
+        ),
+    ]
+
+
 @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC)
 def test_stream_subgraphs_during_execution(
     request: pytest.FixtureRequest, checkpointer_name: str
@@ -4843,6 +4949,69 @@ def test_interrupt_loop(request: pytest.FixtureRequest, checkpointer_name: str):
     ]
 
 
+@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC)
+def test_interrupt_functional(
+    request: pytest.FixtureRequest, checkpointer_name: str
+) -> None:
+    checkpointer: BaseCheckpointSaver = request.getfixturevalue(
+        f"checkpointer_{checkpointer_name}"
+    )
+
+    @task
+    def foo(state: dict) -> dict:
+        return {"a": state["a"] + "foo"}
+
+    @task
+    def bar(state: dict) -> dict:
+        return {"a": state["a"] + "bar", "b": state["b"]}
+
+    @entrypoint(checkpointer=checkpointer)
+    def graph(inputs: dict) -> dict:
+        fut_foo = foo(inputs)
+        value = interrupt("Provide value for bar:")
+        bar_input = {**fut_foo.result(), "b": value}
+        fut_bar = bar(bar_input)
+        return fut_bar.result()
+
+    config = {"configurable": {"thread_id": "1"}}
+    # First run, interrupted at bar
+    graph.invoke({"a": ""}, config)
+    # Resume with an answer
+    res = graph.invoke(Command(resume="bar"), config)
+    assert res == {"a": "foobar", "b": "bar"}
+
+
+@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC)
+def test_interrupt_task_functional(
+    request: pytest.FixtureRequest, checkpointer_name: str
+) -> None:
+    checkpointer: BaseCheckpointSaver = request.getfixturevalue(
+        f"checkpointer_{checkpointer_name}"
+    )
+
+    @task
+    def foo(state: dict) -> dict:
+        return {"a": state["a"] + "foo"}
+
+    @task
+    def bar(state: dict) -> dict:
+        value = interrupt("Provide value for bar:")
+        return {"a": state["a"] + value}
+
+    @entrypoint(checkpointer=checkpointer)
+    def graph(inputs: dict) -> dict:
+        fut_foo = foo(inputs)
+        fut_bar = bar(fut_foo.result())
+        return fut_bar.result()
+
+    config = {"configurable": {"thread_id": "1"}}
+    # First run, interrupted at bar
+    graph.invoke({"a": ""}, config)
+    # Resume with an answer
+    res = graph.invoke(Command(resume="bar"), config)
+    assert res == {"a": "foobar"}
+
+
 def test_root_mixed_return() -> None:
     def my_node(state: list[str]):
         return [Command(update=["a"]), ["b"]]
@@ -5262,9 +5431,10 @@ def test_multiple_updates() -> None:
     ]
 
 
-def test_falsy_return_from_task() -> None:
+@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC)
+def test_falsy_return_from_task(request: pytest.FixtureRequest, checkpointer_name: str):
     """Test with a falsy return from a task."""
-    checkpointer = MemorySaver()
+    checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}")
 
     @task
     def falsy_task() -> bool:
@@ -5276,17 +5446,18 @@ def test_falsy_return_from_task() -> None:
         falsy_task().result()
         interrupt("test")
 
-    configurable = {"configurable": {"thread_id": uuid.uuid4()}}
+    configurable = {"configurable": {"thread_id": str(uuid.uuid4())}}
     graph.invoke({"a": 5}, configurable)
     graph.invoke(Command(resume="123"), configurable)
 
 
-def test_multiple_interrupts_imperative() -> None:
+@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC)
+def test_multiple_interrupts_imperative(
+    request: pytest.FixtureRequest, checkpointer_name: str
+):
     """Test multiple interrupts with an imperative API."""
-    from langgraph.checkpoint.memory import MemorySaver
-    from langgraph.func import entrypoint, task
+    checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}")
 
-    checkpointer = MemorySaver()
     counter = 0
 
     @task
@@ -5307,7 +5478,7 @@ def test_multiple_interrupts_imperative() -> None:
 
         return {"values": values}
 
-    configurable = {"configurable": {"thread_id": uuid.uuid4()}}
+    configurable = {"configurable": {"thread_id": str(uuid.uuid4())}}
     graph.invoke({}, configurable)
     graph.invoke(Command(resume="a"), configurable)
     graph.invoke(Command(resume="b"), configurable)
@@ -5317,3 +5488,312 @@ def test_multiple_interrupts_imperative() -> None:
         "values": [2, "a", 4, "b", 6, "c"],
     }
     assert counter == 3
+
+
+@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC)
+def test_double_interrupt_subgraph(
+    request: pytest.FixtureRequest, checkpointer_name: str
+) -> None:
+    checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}")
+
+    class AgentState(TypedDict):
+        input: str
+
+    def node_1(state: AgentState):
+        result = interrupt("interrupt node 1")
+        return {"input": result}
+
+    def node_2(state: AgentState):
+        result = interrupt("interrupt node 2")
+        return {"input": result}
+
+    subgraph_builder = (
+        StateGraph(AgentState)
+        .add_node("node_1", node_1)
+        .add_node("node_2", node_2)
+        .add_edge(START, "node_1")
+        .add_edge("node_1", "node_2")
+        .add_edge("node_2", END)
+    )
+
+    # invoke the sub graph
+    subgraph = subgraph_builder.compile(checkpointer=checkpointer)
+    thread = {"configurable": {"thread_id": str(uuid.uuid4())}}
+    assert [c for c in subgraph.stream({"input": "test"}, thread)] == [
+        {
+            "__interrupt__": (
+                Interrupt(
+                    value="interrupt node 1",
+                    resumable=True,
+                    ns=[AnyStr("node_1:")],
+                    when="during",
+                ),
+            )
+        },
+    ]
+    # resume from the first interrupt
+    assert [c for c in subgraph.stream(Command(resume="123"), thread)] == [
+        {
+            "node_1": {"input": "123"},
+        },
+        {
+            "__interrupt__": (
+                Interrupt(
+                    value="interrupt node 2",
+                    resumable=True,
+                    ns=[AnyStr("node_2:")],
+                    when="during",
+                ),
+            )
+        },
+    ]
+    # resume from the second interrupt
+    assert [c for c in subgraph.stream(Command(resume="123"), thread)] == [
+        {
+            "node_2": {"input": "123"},
+        },
+    ]
+
+    subgraph = subgraph_builder.compile()
+
+    def invoke_sub_agent(state: AgentState):
+        return subgraph.invoke(state)
+
+    parent_agent = (
+        StateGraph(AgentState)
+        .add_node("invoke_sub_agent", invoke_sub_agent)
+        .add_edge(START, "invoke_sub_agent")
+        .add_edge("invoke_sub_agent", END)
+        .compile(checkpointer=checkpointer)
+    )
+
+    assert [c for c in parent_agent.stream({"input": "test"}, thread)] == [
+        {
+            "__interrupt__": (
+                Interrupt(
+                    value="interrupt node 1",
+                    resumable=True,
+                    ns=[AnyStr("invoke_sub_agent:"), AnyStr("node_1:")],
+                    when="during",
+                ),
+            )
+        },
+    ]
+
+    # resume from the first interrupt
+    assert [c for c in parent_agent.stream(Command(resume=True), thread)] == [
+        {
+            "__interrupt__": (
+                Interrupt(
+                    value="interrupt node 2",
+                    resumable=True,
+                    ns=[AnyStr("invoke_sub_agent:"), AnyStr("node_2:")],
+                    when="during",
+                ),
+            )
+        }
+    ]
+
+    # resume from 2nd interrupt
+    assert [c for c in parent_agent.stream(Command(resume=True), thread)] == [
+        {
+            "invoke_sub_agent": {"input": True},
+        },
+    ]
+
+
+def test_sync_streaming_with_functional_api() -> None:
+    """Test streaming with functional API.
+
+    This test verifies that we're able to stream results as they're being generated
+    rather than have all the results arrive at once after the graph has completed.
+
+    The time of arrival between the two updates corresponding to the two `slow` tasks
+    should be greater than the time delay between the two tasks.
+    """
+
+    time_delay = 0.01
+
+    @task()
+    def slow() -> dict:
+        time.sleep(time_delay)  # Simulate a delay of 10 ms
+        return {"tic": time.time()}
+
+    @entrypoint()
+    def graph(inputs: dict) -> list:
+        first = slow().result()
+        second = slow().result()
+        return [first, second]
+
+    arrival_times = []
+
+    for chunk in graph.stream({}):
+        if "slow" not in chunk:  # We'll just look at the updates from `slow`
+            continue
+        arrival_times.append(time.time())
+
+    assert len(arrival_times) == 2
+    delta = arrival_times[1] - arrival_times[0]
+    # Delta cannot be less than 10 ms if it is streaming as results are generated.
+    assert delta > time_delay
+
+
+def test_entrypoint_without_checkpointer() -> None:
+    """Test no checkpointer."""
+    states = []
+    config = {"configurable": {"thread_id": "1"}}
+
+    # Test without previous
+    @entrypoint()
+    def foo(inputs: Any) -> Any:
+        states.append(inputs)
+        return inputs
+
+    assert foo.invoke({"a": "1"}, config) == {"a": "1"}
+
+    @entrypoint()
+    def foo(inputs: Any, *, previous: Any) -> Any:
+        states.append(previous)
+        return {"previous": previous, "current": inputs}
+
+    assert foo.invoke({"a": "1"}, config) == {"current": {"a": "1"}, "previous": None}
+    assert foo.invoke({"a": "1"}, config) == {"current": {"a": "1"}, "previous": None}
+
+
+async def test_async_entrypoint_without_checkpointer() -> None:
+    """Test no checkpointer."""
+    states = []
+    config = {"configurable": {"thread_id": "1"}}
+
+    # Test without previous
+    @entrypoint()
+    async def foo(inputs: Any) -> Any:
+        states.append(inputs)
+        return inputs
+
+    assert (await foo.ainvoke({"a": "1"}, config)) == {"a": "1"}
+
+    @entrypoint()
+    async def foo(inputs: Any, *, previous: Any) -> Any:
+        states.append(previous)
+        return {"previous": previous, "current": inputs}
+
+    assert (await foo.ainvoke({"a": "1"}, config)) == {
+        "current": {"a": "1"},
+        "previous": None,
+    }
+    assert (await foo.ainvoke({"a": "1"}, config)) == {
+        "current": {"a": "1"},
+        "previous": None,
+    }
+
+
+def test_entrypoint_stateful() -> None:
+    """Test stateful entrypoint invoke."""
+
+    # Test invoke
+    states = []
+
+    @entrypoint(checkpointer=MemorySaver())
+    def foo(inputs, *, previous: Any) -> Any:
+        states.append(previous)
+        return {"previous": previous, "current": inputs}
+
+    config = {"configurable": {"thread_id": "1"}}
+
+    assert foo.invoke({"a": "1"}, config) == {"current": {"a": "1"}, "previous": None}
+    assert foo.invoke({"a": "2"}, config) == {
+        "current": {"a": "2"},
+        "previous": {"current": {"a": "1"}, "previous": None},
+    }
+    assert foo.invoke({"a": "3"}, config) == {
+        "current": {"a": "3"},
+        "previous": {
+            "current": {"a": "2"},
+            "previous": {"current": {"a": "1"}, "previous": None},
+        },
+    }
+    assert states == [
+        None,
+        {"current": {"a": "1"}, "previous": None},
+        {"current": {"a": "2"}, "previous": {"current": {"a": "1"}, "previous": None}},
+    ]
+
+    # Test stream
+    @entrypoint(checkpointer=MemorySaver())
+    def foo(inputs, *, previous: Any) -> Any:
+        return {"previous": previous, "current": inputs}
+
+    config = {"configurable": {"thread_id": "1"}}
+    items = [item for item in foo.stream({"a": "1"}, config)]
+    assert items == [{"foo": {"current": {"a": "1"}, "previous": None}}]
+
+
+def test_entrypoint_from_sync_generator() -> None:
+    """@entrypoint does not support sync generators."""
+    previous_return_values = []
+
+    @entrypoint(checkpointer=MemorySaver())
+    def foo(inputs, previous=None) -> Any:
+        previous_return_values.append(previous)
+        yield "a"
+        yield "b"
+
+    config = {"configurable": {"thread_id": "1"}}
+
+    assert foo.invoke({"a": "1"}, config) == ["a", "b"]
+    assert previous_return_values == [None]
+    assert foo.invoke({"a": "2"}, config) == ["a", "b"]
+    assert previous_return_values == [None, ["a", "b"]]
+
+
+def test_entrypoint_request_stream_writer() -> None:
+    """Test using a stream writer with an entrypoint."""
+
+    @entrypoint(checkpointer=MemorySaver())
+    def foo(inputs, writer: StreamWriter) -> Any:
+        writer("a")
+        yield "b"
+
+    config = {"configurable": {"thread_id": "1"}}
+
+    # Different invocations
+    # Are any of these confusing or unexpected?
+    assert list(foo.invoke({}, config)) == ["b"]
+    assert list(foo.stream({}, config)) == ["a", "b"]
+
+    # Stream modes
+    assert list(foo.stream({}, config, stream_mode=["updates"])) == [
+        ("updates", {"foo": ["b"]})
+    ]
+    assert list(foo.stream({}, config, stream_mode=["values"])) == [("values", ["b"])]
+    assert list(foo.stream({}, config, stream_mode=["custom"])) == [
+        (
+            "custom",
+            "a",
+        ),
+        (
+            "custom",
+            "b",
+        ),
+    ]
+
+
+async def test_entrypoint_from_async_generator() -> None:
+    """@entrypoint does not support sync generators."""
+    # Test invoke
+    previous_return_values = []
+
+    # In this version reducers do not work
+    @entrypoint(checkpointer=MemorySaver())
+    async def foo(inputs, previous=None) -> Any:
+        previous_return_values.append(previous)
+        yield "a"
+        yield "b"
+
+    config = {"configurable": {"thread_id": "1"}}
+
+    assert list(await foo.ainvoke({"a": "1"}, config)) == ["a", "b"]
+    assert previous_return_values == [None]
+    assert list(foo.invoke({"a": "2"}, config)) == ["a", "b"]
+    assert previous_return_values == [None, ["a", "b"]]

libs/langgraph/tests/test_pregel_async.py

@@ -48,7 +48,7 @@ from langgraph.checkpoint.base import (
 )
 from langgraph.checkpoint.memory import MemorySaver
 from langgraph.constants import CONFIG_KEY_NODE_FINISHED, ERROR, PULL, PUSH, START
-from langgraph.errors import InvalidUpdateError, MultipleSubgraphsError, NodeInterrupt
+from langgraph.errors import InvalidUpdateError, NodeInterrupt
 from langgraph.func import entrypoint, task
 from langgraph.graph import END, Graph, StateGraph
 from langgraph.graph.message import MessagesState, add_messages
@@ -89,6 +89,11 @@ logger = logging.getLogger(__name__)
 
 pytestmark = pytest.mark.anyio
 
+NEEDS_CONTEXTVARS = pytest.mark.skipif(
+    sys.version_info < (3, 11),
+    reason="Python 3.11+ is required for async contextvars support",
+)
+
 
 async def test_checkpoint_errors() -> None:
     class FaultyGetCheckpointer(MemorySaver):
@@ -501,10 +506,7 @@ async def test_node_cancellation_on_other_node_exception_two() -> None:
         await graph.ainvoke(1)
 
 
-@pytest.mark.skipif(
-    sys.version_info < (3, 11),
-    reason="Python 3.11+ is required for async contextvars support",
-)
+@NEEDS_CONTEXTVARS
 @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
 async def test_dynamic_interrupt(checkpointer_name: str) -> None:
     class State(TypedDict):
@@ -678,10 +680,7 @@ async def test_dynamic_interrupt(checkpointer_name: str) -> None:
         )
 
 
-@pytest.mark.skipif(
-    sys.version_info < (3, 11),
-    reason="Python 3.11+ is required for async contextvars support",
-)
+@NEEDS_CONTEXTVARS
 @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
 async def test_dynamic_interrupt_subgraph(checkpointer_name: str) -> None:
     class SubgraphState(TypedDict):
@@ -872,10 +871,7 @@ async def test_dynamic_interrupt_subgraph(checkpointer_name: str) -> None:
         )
 
 
-@pytest.mark.skipif(
-    sys.version_info < (3, 11),
-    reason="Python 3.11+ is required for async contextvars support",
-)
+@NEEDS_CONTEXTVARS
 @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
 async def test_copy_checkpoint(checkpointer_name: str) -> None:
     class State(TypedDict):
@@ -1079,10 +1075,7 @@ async def test_copy_checkpoint(checkpointer_name: str) -> None:
         )
 
 
-@pytest.mark.skipif(
-    sys.version_info < (3, 11),
-    reason="Python 3.11+ is required for async contextvars support",
-)
+@NEEDS_CONTEXTVARS
 @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
 async def test_node_not_cancelled_on_other_node_interrupted(
     checkpointer_name: str,
@@ -2442,10 +2435,7 @@ async def test_send_sequences(checkpointer_name: str) -> None:
         ]
 
 
-@pytest.mark.skipif(
-    sys.version_info < (3, 11),
-    reason="Python 3.11+ is required for async contextvars support",
-)
+@NEEDS_CONTEXTVARS
 @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
 async def test_imp_task(checkpointer_name: str) -> None:
     async with awith_checkpointer(checkpointer_name) as checkpointer:
@@ -2493,10 +2483,7 @@ async def test_imp_task(checkpointer_name: str) -> None:
         assert mapper_calls == 2
 
 
-@pytest.mark.skipif(
-    sys.version_info < (3, 11),
-    reason="Python 3.11+ is required for async contextvars support",
-)
+@NEEDS_CONTEXTVARS
 @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
 async def test_imp_task_cancel(checkpointer_name: str) -> None:
     async with awith_checkpointer(checkpointer_name) as checkpointer:
@@ -2547,10 +2534,7 @@ async def test_imp_task_cancel(checkpointer_name: str) -> None:
         assert mapper_cancels == 2
 
 
-@pytest.mark.skipif(
-    sys.version_info < (3, 11),
-    reason="Python 3.11+ is required for async contextvars support",
-)
+@NEEDS_CONTEXTVARS
 @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
 async def test_imp_sync_from_async(checkpointer_name: str) -> None:
     async with awith_checkpointer(checkpointer_name) as checkpointer:
@@ -2583,10 +2567,7 @@ async def test_imp_sync_from_async(checkpointer_name: str) -> None:
         ]
 
 
-@pytest.mark.skipif(
-    sys.version_info < (3, 11),
-    reason="Python 3.11+ is required for async contextvars support",
-)
+@NEEDS_CONTEXTVARS
 @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
 async def test_imp_stream_order(checkpointer_name: str) -> None:
     async with awith_checkpointer(checkpointer_name) as checkpointer:
@@ -4068,9 +4049,8 @@ async def test_invoke_join_then_call_other_pregel(
     async with awith_checkpointer(checkpointer_name) as checkpointer:
         # add checkpointer
         app.checkpointer = checkpointer
-        # subgraph is called twice in the same node, through .map(), so raises
-        with pytest.raises(MultipleSubgraphsError):
-            await app.ainvoke([2, 3], {"configurable": {"thread_id": "1"}})
+        # subgraph is called twice, and that works
+        assert await app.ainvoke([2, 3], {"configurable": {"thread_id": "1"}}) == 27
 
         # set inner graph checkpointer NeverCheckpoint
         inner_app.checkpointer = False
@@ -4294,10 +4274,10 @@ async def test_in_one_fan_out_state_graph_waiting_edge(checkpointer_name: str) -
         docs: Annotated[list[str], sorted_add]
 
     async def rewrite_query(data: State) -> State:
-        return {"query": f'query: {data["query"]}'}
+        return {"query": f"query: {data['query']}"}
 
     async def analyzer_one(data: State) -> State:
-        return {"query": f'analyzed: {data["query"]}'}
+        return {"query": f"analyzed: {data['query']}"}
 
     async def retriever_one(data: State) -> State:
         return {"docs": ["doc1", "doc2"]}
@@ -4384,10 +4364,10 @@ async def test_in_one_fan_out_state_graph_waiting_edge_via_branch(
         docs: Annotated[list[str], sorted_add]
 
     async def rewrite_query(data: State) -> State:
-        return {"query": f'query: {data["query"]}'}
+        return {"query": f"query: {data['query']}"}
 
     async def analyzer_one(data: State) -> State:
-        return {"query": f'analyzed: {data["query"]}'}
+        return {"query": f"analyzed: {data['query']}"}
 
     async def retriever_one(data: State) -> State:
         return {"docs": ["doc1", "doc2"]}
@@ -4801,11 +4781,11 @@ async def test_in_one_fan_out_state_graph_waiting_edge_plus_regular(
         docs: Annotated[list[str], sorted_add]
 
     async def rewrite_query(data: State) -> State:
-        return {"query": f'query: {data["query"]}'}
+        return {"query": f"query: {data['query']}"}
 
     async def analyzer_one(data: State) -> State:
         await asyncio.sleep(0.1)
-        return {"query": f'analyzed: {data["query"]}'}
+        return {"query": f"analyzed: {data['query']}"}
 
     async def retriever_one(data: State) -> State:
         return {"docs": ["doc1", "doc2"]}
@@ -4895,10 +4875,10 @@ async def test_in_one_fan_out_state_graph_waiting_edge_multiple() -> None:
         docs: Annotated[list[str], sorted_add]
 
     async def rewrite_query(data: State) -> State:
-        return {"query": f'query: {data["query"]}'}
+        return {"query": f"query: {data['query']}"}
 
     async def analyzer_one(data: State) -> State:
-        return {"query": f'analyzed: {data["query"]}'}
+        return {"query": f"analyzed: {data['query']}"}
 
     async def retriever_one(data: State) -> State:
         return {"docs": ["doc1", "doc2"]}
@@ -4979,13 +4959,13 @@ async def test_in_one_fan_out_state_graph_waiting_edge_multiple_cond_edge() -> N
         docs: Annotated[list[str], sorted_add]
 
     async def rewrite_query(data: State) -> State:
-        return {"query": f'query: {data["query"]}'}
+        return {"query": f"query: {data['query']}"}
 
     async def retriever_picker(data: State) -> list[str]:
         return ["analyzer_one", "retriever_two"]
 
     async def analyzer_one(data: State) -> State:
-        return {"query": f'analyzed: {data["query"]}'}
+        return {"query": f"analyzed: {data['query']}"}
 
     async def retriever_one(data: State) -> State:
         return {"docs": ["doc1", "doc2"]}
@@ -6117,10 +6097,7 @@ async def test_parent_command(checkpointer_name: str) -> None:
         )
 
 
-@pytest.mark.skipif(
-    sys.version_info < (3, 11),
-    reason="Python 3.11+ is required for async contextvars support",
-)
+@NEEDS_CONTEXTVARS
 @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
 async def test_interrupt_subgraph(checkpointer_name: str):
     class State(TypedDict):
@@ -6153,10 +6130,7 @@ async def test_interrupt_subgraph(checkpointer_name: str):
         assert await graph.ainvoke(Command(resume="bar"), thread1)
 
 
-@pytest.mark.skipif(
-    sys.version_info < (3, 11),
-    reason="Python 3.11+ is required for async contextvars support",
-)
+@NEEDS_CONTEXTVARS
 @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
 async def test_interrupt_multiple(checkpointer_name: str):
     class State(TypedDict):
@@ -6220,10 +6194,7 @@ async def test_interrupt_multiple(checkpointer_name: str):
         ]
 
 
-@pytest.mark.skipif(
-    sys.version_info < (3, 11),
-    reason="Python 3.11+ is required for async contextvars support",
-)
+@NEEDS_CONTEXTVARS
 @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
 async def test_interrupt_loop(checkpointer_name: str):
     class State(TypedDict):
@@ -6310,6 +6281,63 @@ async def test_interrupt_loop(checkpointer_name: str):
         ]
 
 
+@NEEDS_CONTEXTVARS
+@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
+async def test_interrupt_functional(checkpointer_name: str) -> None:
+    @task
+    async def foo(state: dict) -> dict:
+        return {"a": state["a"] + "foo"}
+
+    @task
+    async def bar(state: dict) -> dict:
+        return {"a": state["a"] + "bar", "b": state["b"]}
+
+    async with awith_checkpointer(checkpointer_name) as checkpointer:
+
+        @entrypoint(checkpointer=checkpointer)
+        async def graph(inputs: dict) -> dict:
+            foo_result = await foo(inputs)
+            value = interrupt("Provide value for bar:")
+            bar_input = {**foo_result, "b": value}
+            bar_result = await bar(bar_input)
+            return bar_result
+
+        config = {"configurable": {"thread_id": "1"}}
+        # First run, interrupted at bar
+        await graph.ainvoke({"a": ""}, config)
+        # Resume with an answer
+        res = await graph.ainvoke(Command(resume="bar"), config)
+        assert res == {"a": "foobar", "b": "bar"}
+
+
+@NEEDS_CONTEXTVARS
+@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
+async def test_interrupt_task_functional(checkpointer_name: str) -> None:
+    @task
+    async def foo(state: dict) -> dict:
+        return {"a": state["a"] + "foo"}
+
+    @task
+    async def bar(state: dict) -> dict:
+        value = interrupt("Provide value for bar:")
+        return {"a": state["a"] + value}
+
+    async with awith_checkpointer(checkpointer_name) as checkpointer:
+
+        @entrypoint(checkpointer=checkpointer)
+        async def graph(inputs: dict) -> dict:
+            foo_result = await foo(inputs)
+            bar_result = await bar(foo_result)
+            return bar_result
+
+        config = {"configurable": {"thread_id": "1"}}
+        # First run, interrupted at bar
+        await graph.ainvoke({"a": ""}, config)
+        # Resume with an answer
+        res = await graph.ainvoke(Command(resume="bar"), config)
+        assert res == {"a": "foobar"}
+
+
 @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
 async def test_command_with_static_breakpoints(checkpointer_name: str) -> None:
     """Test that we can use Command to resume and update with static breakpoints."""
@@ -6508,10 +6536,7 @@ async def test_parallel_node_execution():
     assert duration < 3.0
 
 
-@pytest.mark.skipif(
-    sys.version_info < (3, 11),
-    reason="Python 3.11+ is required for async contextvars support",
-)
+@NEEDS_CONTEXTVARS
 @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
 async def test_multiple_interrupt_state_persistence(checkpointer_name: str) -> None:
     """Test that state is preserved correctly across multiple interrupts."""
@@ -6692,33 +6717,29 @@ async def test_multiple_updates() -> None:
     ]
 
 
-@pytest.mark.skipif(
-    sys.version_info < (3, 11),
-    reason="Python 3.11+ is required for async contextvars support",
-)
-async def test_falsy_return_from_task() -> None:
+@NEEDS_CONTEXTVARS
+@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
+async def test_falsy_return_from_task(checkpointer_name: str) -> None:
     """Test with a falsy return from a task."""
-    checkpointer = MemorySaver()
 
     @task
     async def falsy_task() -> bool:
         return False
 
-    @entrypoint(checkpointer=checkpointer)
-    async def graph(state: dict) -> dict:
-        """React tool."""
-        await falsy_task()
-        interrupt("test")
+    async with awith_checkpointer(checkpointer_name) as checkpointer:
 
-    configurable = {"configurable": {"thread_id": uuid.uuid4()}}
-    await graph.ainvoke({"a": 5}, configurable)
-    await graph.ainvoke(Command(resume="123"), configurable)
+        @entrypoint(checkpointer=checkpointer)
+        async def graph(state: dict) -> dict:
+            """React tool."""
+            await falsy_task()
+            interrupt("test")
+
+        configurable = {"configurable": {"thread_id": str(uuid.uuid4())}}
+        await graph.ainvoke({"a": 5}, configurable)
+        await graph.ainvoke(Command(resume="123"), configurable)
 
 
-@pytest.mark.skipif(
-    sys.version_info < (3, 11),
-    reason="Python 3.11+ is required for async contextvars support",
-)
+@NEEDS_CONTEXTVARS
 @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
 async def test_multiple_interrupts_imperative(checkpointer_name: str) -> None:
     """Test multiple interrupts with an imperative API."""
@@ -6756,3 +6777,154 @@ async def test_multiple_interrupts_imperative(checkpointer_name: str) -> None:
             "values": [2, "a", 4, "b", 6, "c"],
         }
         assert counter == 3
+
+
+@NEEDS_CONTEXTVARS
+@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
+async def test_double_interrupt_subgraph(checkpointer_name: str) -> None:
+    class AgentState(TypedDict):
+        input: str
+
+    def node_1(state: AgentState):
+        result = interrupt("interrupt node 1")
+        return {"input": result}
+
+    def node_2(state: AgentState):
+        result = interrupt("interrupt node 2")
+        return {"input": result}
+
+    subgraph_builder = (
+        StateGraph(AgentState)
+        .add_node("node_1", node_1)
+        .add_node("node_2", node_2)
+        .add_edge(START, "node_1")
+        .add_edge("node_1", "node_2")
+        .add_edge("node_2", END)
+    )
+
+    async with awith_checkpointer(checkpointer_name) as checkpointer:
+        # invoke the sub graph
+        subgraph = subgraph_builder.compile(checkpointer=checkpointer)
+        thread = {"configurable": {"thread_id": str(uuid.uuid4())}}
+        assert [c async for c in subgraph.astream({"input": "test"}, thread)] == [
+            {
+                "__interrupt__": (
+                    Interrupt(
+                        value="interrupt node 1",
+                        resumable=True,
+                        ns=[AnyStr("node_1:")],
+                        when="during",
+                    ),
+                )
+            },
+        ]
+        # resume from the first interrupt
+        assert [c async for c in subgraph.astream(Command(resume="123"), thread)] == [
+            {
+                "node_1": {"input": "123"},
+            },
+            {
+                "__interrupt__": (
+                    Interrupt(
+                        value="interrupt node 2",
+                        resumable=True,
+                        ns=[AnyStr("node_2:")],
+                        when="during",
+                    ),
+                )
+            },
+        ]
+        # resume from the second interrupt
+        assert [c async for c in subgraph.astream(Command(resume="123"), thread)] == [
+            {
+                "node_2": {"input": "123"},
+            },
+        ]
+
+        subgraph = subgraph_builder.compile()
+
+        def invoke_sub_agent(state: AgentState):
+            return subgraph.invoke(state)
+
+        parent_agent = (
+            StateGraph(AgentState)
+            .add_node("invoke_sub_agent", invoke_sub_agent)
+            .add_edge(START, "invoke_sub_agent")
+            .add_edge("invoke_sub_agent", END)
+            .compile(checkpointer=checkpointer)
+        )
+
+        assert [c async for c in parent_agent.astream({"input": "test"}, thread)] == [
+            {
+                "__interrupt__": (
+                    Interrupt(
+                        value="interrupt node 1",
+                        resumable=True,
+                        ns=[AnyStr("invoke_sub_agent:"), AnyStr("node_1:")],
+                        when="during",
+                    ),
+                )
+            },
+        ]
+
+        # resume from the first interrupt
+        assert [
+            c async for c in parent_agent.astream(Command(resume=True), thread)
+        ] == [
+            {
+                "__interrupt__": (
+                    Interrupt(
+                        value="interrupt node 2",
+                        resumable=True,
+                        ns=[AnyStr("invoke_sub_agent:"), AnyStr("node_2:")],
+                        when="during",
+                    ),
+                )
+            }
+        ]
+
+        # resume from 2nd interrupt
+        assert [
+            c async for c in parent_agent.astream(Command(resume=True), thread)
+        ] == [
+            {
+                "invoke_sub_agent": {"input": True},
+            },
+        ]
+
+
+@NEEDS_CONTEXTVARS
+async def test_async_streaming_with_functional_api() -> None:
+    """Test streaming with functional API.
+
+    This test verifies that we're able to stream results as they're being generated
+    rather than have all the results arrive at once after the graph has completed.
+
+    The time of arrival between the two updates corresponding to the two `slow` tasks
+    should be greater than the time delay between the two tasks.
+    """
+
+    time_delay = 0.01
+
+    @task()
+    async def slow() -> dict:
+        await asyncio.sleep(time_delay)  # Simulate a delay of 10 ms
+        return {"tic": asyncio.get_running_loop().time()}
+
+    @entrypoint()
+    async def graph(inputs: dict) -> list:
+        first = await slow()
+        second = await slow()
+        return [first, second]
+
+    arrival_times = []
+
+    async for chunk in graph.astream({}):
+        if "slow" not in chunk:  # We'll just look at the updates from `slow`
+            continue
+        arrival_times.append(asyncio.get_running_loop().time())
+
+    assert len(arrival_times) == 2
+    delta = arrival_times[1] - arrival_times[0]
+    # Delta cannot be less than 10 ms if it is streaming as results are generated.
+    assert delta > time_delay

libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py

@@ -1,5 +1,6 @@
 import asyncio
 import concurrent.futures
+from collections.abc import Sequence
 from contextlib import (
     AbstractAsyncContextManager,
     AbstractContextManager,
@@ -7,7 +8,7 @@ from contextlib import (
     ExitStack,
 )
 from functools import partial
-from typing import Any, Optional, Sequence
+from typing import Any, Optional
 from uuid import UUID
 
 import orjson
@@ -15,7 +16,7 @@ from langchain_core.runnables import RunnableConfig
 from typing_extensions import Self
 
 import langgraph.scheduler.kafka.serde as serde
-from langgraph.constants import CONFIG_KEY_DELEGATE, ERROR, NS_END, NS_SEP
+from langgraph.constants import CONFIG_KEY_DELEGATE, ERROR
 from langgraph.errors import CheckpointNotLatest, GraphDelegate, TaskNotFound
 from langgraph.pregel import Pregel
 from langgraph.pregel.algo import prepare_single_task
@@ -39,7 +40,7 @@ from langgraph.scheduler.kafka.types import (
     Topics,
 )
 from langgraph.types import LoopProtocol, PregelExecutableTask, RetryPolicy
-from langgraph.utils.config import patch_configurable
+from langgraph.utils.config import patch_configurable, recast_checkpoint_ns
 
 
 class AsyncKafkaExecutor(AbstractAsyncContextManager):
@@ -165,14 +166,12 @@ class AsyncKafkaExecutor(AbstractAsyncContextManager):
         # find graph
         if checkpoint_ns := msg["config"]["configurable"].get("checkpoint_ns"):
             # remove task_ids from checkpoint_ns
-            recast_checkpoint_ns = NS_SEP.join(
-                part.split(NS_END)[0] for part in checkpoint_ns.split(NS_SEP)
-            )
+            recast = recast_checkpoint_ns(checkpoint_ns)
             # find the subgraph with the matching name
-            if recast_checkpoint_ns in self.subgraphs:
-                graph = self.subgraphs[recast_checkpoint_ns]
+            if recast in self.subgraphs:
+                graph = self.subgraphs[recast]
             else:
-                raise ValueError(f"Subgraph {recast_checkpoint_ns} not found")
+                raise ValueError(f"Subgraph {recast} not found")
         else:
             graph = self.graph
         # process message
@@ -183,16 +182,19 @@ class AsyncKafkaExecutor(AbstractAsyncContextManager):
             raise RuntimeError("Checkpoint not found")
         if saved.checkpoint["id"] != msg["config"]["configurable"]["checkpoint_id"]:
             raise CheckpointNotLatest()
-        async with AsyncChannelsManager(
-            graph.channels,
-            saved.checkpoint,
-            LoopProtocol(
-                config=msg["config"],
-                store=self.graph.store,
-                step=saved.metadata["step"] + 1,
-                stop=saved.metadata["step"] + 2,
-            ),
-        ) as (channels, managed), AsyncBackgroundExecutor(msg["config"]) as submit:
+        async with (
+            AsyncChannelsManager(
+                graph.channels,
+                saved.checkpoint,
+                LoopProtocol(
+                    config=msg["config"],
+                    store=self.graph.store,
+                    step=saved.metadata["step"] + 1,
+                    stop=saved.metadata["step"] + 2,
+                ),
+            ) as (channels, managed),
+            AsyncBackgroundExecutor(msg["config"]) as submit,
+        ):
             if task := await asyncio.to_thread(
                 prepare_single_task,
                 msg["task"]["path"],
@@ -378,14 +380,12 @@ class KafkaExecutor(AbstractContextManager):
         # find graph
         if checkpoint_ns := msg["config"]["configurable"].get("checkpoint_ns"):
             # remove task_ids from checkpoint_ns
-            recast_checkpoint_ns = NS_SEP.join(
-                part.split(NS_END)[0] for part in checkpoint_ns.split(NS_SEP)
-            )
+            recast = recast_checkpoint_ns(checkpoint_ns)
             # find the subgraph with the matching name
-            if recast_checkpoint_ns in self.subgraphs:
-                graph = self.subgraphs[recast_checkpoint_ns]
+            if recast in self.subgraphs:
+                graph = self.subgraphs[recast]
             else:
-                raise ValueError(f"Subgraph {recast_checkpoint_ns} not found")
+                raise ValueError(f"Subgraph {recast} not found")
         else:
             graph = self.graph
         # process message
@@ -396,16 +396,19 @@ class KafkaExecutor(AbstractContextManager):
             raise RuntimeError("Checkpoint not found")
         if saved.checkpoint["id"] != msg["config"]["configurable"]["checkpoint_id"]:
             raise CheckpointNotLatest()
-        with ChannelsManager(
-            graph.channels,
-            saved.checkpoint,
-            LoopProtocol(
-                config=msg["config"],
-                store=self.graph.store,
-                step=saved.metadata["step"] + 1,
-                stop=saved.metadata["step"] + 2,
-            ),
-        ) as (channels, managed), BackgroundExecutor({}) as submit:
+        with (
+            ChannelsManager(
+                graph.channels,
+                saved.checkpoint,
+                LoopProtocol(
+                    config=msg["config"],
+                    store=self.graph.store,
+                    step=saved.metadata["step"] + 1,
+                    stop=saved.metadata["step"] + 2,
+                ),
+            ) as (channels, managed),
+            BackgroundExecutor({}) as submit,
+        ):
             if task := prepare_single_task(
                 msg["task"]["path"],
                 msg["task"]["id"],

libs/scheduler-kafka/langgraph/scheduler/kafka/orchestrator.py

@@ -13,11 +13,11 @@ from typing_extensions import Self
 
 import langgraph.scheduler.kafka.serde as serde
 from langgraph.constants import (
+    CONF,
     CONFIG_KEY_DEDUPE_TASKS,
     CONFIG_KEY_ENSURE_LATEST,
+    CONFIG_KEY_SCRATCHPAD,
     INTERRUPT,
-    NS_END,
-    NS_SEP,
     SCHEDULED,
 )
 from langgraph.errors import CheckpointNotLatest, GraphInterrupt
@@ -37,7 +37,7 @@ from langgraph.scheduler.kafka.types import (
     Topics,
 )
 from langgraph.types import RetryPolicy
-from langgraph.utils.config import patch_configurable
+from langgraph.utils.config import patch_configurable, recast_checkpoint_ns
 
 
 class AsyncKafkaOrchestrator(AbstractAsyncContextManager):
@@ -140,14 +140,12 @@ class AsyncKafkaOrchestrator(AbstractAsyncContextManager):
         # find graph
         if checkpoint_ns := msg["config"]["configurable"].get("checkpoint_ns"):
             # remove task_ids from checkpoint_ns
-            recast_checkpoint_ns = NS_SEP.join(
-                part.split(NS_END)[0] for part in checkpoint_ns.split(NS_SEP)
-            )
+            recast = recast_checkpoint_ns(checkpoint_ns)
             # find the subgraph with the matching name
-            if recast_checkpoint_ns in self.subgraphs:
-                graph = self.subgraphs[recast_checkpoint_ns]
+            if recast in self.subgraphs:
+                graph = self.subgraphs[recast]
             else:
-                raise ValueError(f"Subgraph {recast_checkpoint_ns} not found")
+                raise ValueError(f"Subgraph {recast} not found")
         else:
             graph = self.graph
         # process message
@@ -163,7 +161,6 @@ class AsyncKafkaOrchestrator(AbstractAsyncContextManager):
             stream_keys=graph.stream_channels,
             interrupt_after=graph.interrupt_after_nodes,
             interrupt_before=graph.interrupt_before_nodes,
-            check_subgraphs=False,
         ) as loop:
             if loop.tick(input_keys=graph.input_channels):
                 # wait for checkpoint to be saved
@@ -173,6 +170,16 @@ class AsyncKafkaOrchestrator(AbstractAsyncContextManager):
                 if new_tasks := [
                     t for t in loop.tasks.values() if not t.scheduled and not t.writes
                 ]:
+                    config = patch_configurable(
+                        loop.config,
+                        {
+                            **loop.checkpoint_config["configurable"],
+                            CONFIG_KEY_DEDUPE_TASKS: True,
+                            CONFIG_KEY_ENSURE_LATEST: True,
+                        },
+                    )
+                    if CONFIG_KEY_SCRATCHPAD in config[CONF]:
+                        config[CONF][CONFIG_KEY_SCRATCHPAD]["subgraph_counter"] = 0
                     # send messages to executor
                     futures = await asyncio.gather(
                         *(
@@ -180,16 +187,7 @@ class AsyncKafkaOrchestrator(AbstractAsyncContextManager):
                                 self.topics.executor,
                                 value=serde.dumps(
                                     MessageToExecutor(
-                                        config=patch_configurable(
-                                            loop.config,
-                                            {
-                                                **loop.checkpoint_config[
-                                                    "configurable"
-                                                ],
-                                                CONFIG_KEY_DEDUPE_TASKS: True,
-                                                CONFIG_KEY_ENSURE_LATEST: True,
-                                            },
-                                        ),
+                                        config=config,
                                         task=ExecutorTask(id=task.id, path=task.path),
                                         finally_send=msg.get("finally_send"),
                                     )
@@ -330,14 +328,12 @@ class KafkaOrchestrator(AbstractContextManager):
         # find graph
         if checkpoint_ns := msg["config"]["configurable"].get("checkpoint_ns"):
             # remove task_ids from checkpoint_ns
-            recast_checkpoint_ns = NS_SEP.join(
-                part.split(NS_END)[0] for part in checkpoint_ns.split(NS_SEP)
-            )
+            recast = recast_checkpoint_ns(checkpoint_ns)
             # find the subgraph with the matching name
-            if recast_checkpoint_ns in self.subgraphs:
-                graph = self.subgraphs[recast_checkpoint_ns]
+            if recast in self.subgraphs:
+                graph = self.subgraphs[recast]
             else:
-                raise ValueError(f"Subgraph {recast_checkpoint_ns} not found")
+                raise ValueError(f"Subgraph {recast} not found")
         else:
             graph = self.graph
         # process message
@@ -353,7 +349,6 @@ class KafkaOrchestrator(AbstractContextManager):
             stream_keys=graph.stream_channels,
             interrupt_after=graph.interrupt_after_nodes,
             interrupt_before=graph.interrupt_before_nodes,
-            check_subgraphs=False,
         ) as loop:
             if loop.tick(input_keys=graph.input_channels):
                 # wait for checkpoint to be saved
@@ -363,20 +358,23 @@ class KafkaOrchestrator(AbstractContextManager):
                 if new_tasks := [
                     t for t in loop.tasks.values() if not t.scheduled and not t.writes
                 ]:
+                    config = patch_configurable(
+                        loop.config,
+                        {
+                            **loop.checkpoint_config["configurable"],
+                            CONFIG_KEY_DEDUPE_TASKS: True,
+                            CONFIG_KEY_ENSURE_LATEST: True,
+                        },
+                    )
+                    if CONFIG_KEY_SCRATCHPAD in config[CONF]:
+                        config[CONF][CONFIG_KEY_SCRATCHPAD]["subgraph_counter"] = 0
                     # send messages to executor
                     futures = [
                         self.producer.send(
                             self.topics.executor,
                             value=serde.dumps(
                                 MessageToExecutor(
-                                    config=patch_configurable(
-                                        loop.config,
-                                        {
-                                            **loop.checkpoint_config["configurable"],
-                                            CONFIG_KEY_DEDUPE_TASKS: True,
-                                            CONFIG_KEY_ENSURE_LATEST: True,
-                                        },
-                                    ),
+                                    config=config,
                                     task=ExecutorTask(id=task.id, path=task.path),
                                     finally_send=msg.get("finally_send"),
                                 )

libs/scheduler-kafka/tests/any.py

@@ -53,3 +53,11 @@ class AnyList(list):
                 return False
         else:
             return True
+
+
+class AnyInt(int):
+    def __init__(self) -> None:
+        super().__init__()
+
+    def __eq__(self, other: object) -> bool:
+        return isinstance(other, int)

libs/scheduler-kafka/tests/test_subgraph.py

@@ -15,7 +15,7 @@ from langgraph.graph.state import StateGraph
 from langgraph.pregel import Pregel
 from langgraph.scheduler.kafka import serde
 from langgraph.scheduler.kafka.types import MessageToOrchestrator, Topics
-from tests.any import AnyDict, AnyList
+from tests.any import AnyDict, AnyInt
 from tests.drain import drain_topics_async
 from tests.messages import _AnyIdAIMessage, _AnyIdHumanMessage
 
@@ -195,10 +195,16 @@ async def test_subgraph_w_interrupt(
                             "__pregel_ensure_latest": True,
                             "__pregel_dedupe_tasks": True,
                             "__pregel_resuming": False,
+                            "__pregel_previous": None,
                             "__pregel_store": None,
                             "__pregel_task_id": history[0].tasks[0].id,
-                            "__pregel_scratchpad": {},
-                            "__pregel_writes": AnyList(),
+                            "__pregel_scratchpad": {
+                                "subgraph_counter": AnyInt(),
+                                "call_counter": 0,
+                                "interrupt_counter": -1,
+                                "null_resume": None,
+                                "resume": [],
+                            },
                             "checkpoint_id": None,
                             "checkpoint_map": {
                                 "": history[0].config["configurable"]["checkpoint_id"]
@@ -262,10 +268,16 @@ async def test_subgraph_w_interrupt(
                             "__pregel_ensure_latest": True,
                             "__pregel_dedupe_tasks": True,
                             "__pregel_resuming": False,
+                            "__pregel_previous": None,
                             "__pregel_store": None,
                             "__pregel_task_id": history[0].tasks[0].id,
-                            "__pregel_scratchpad": {},
-                            "__pregel_writes": AnyList(),
+                            "__pregel_scratchpad": {
+                                "subgraph_counter": AnyInt(),
+                                "call_counter": 0,
+                                "interrupt_counter": -1,
+                                "null_resume": None,
+                                "resume": [],
+                            },
                             "checkpoint_id": c.config["configurable"]["checkpoint_id"],
                             "checkpoint_map": {
                                 "": history[0].config["configurable"]["checkpoint_id"]
@@ -359,10 +371,16 @@ async def test_subgraph_w_interrupt(
                             "__pregel_ensure_latest": True,
                             "__pregel_dedupe_tasks": True,
                             "__pregel_resuming": False,
+                            "__pregel_previous": None,
                             "__pregel_store": None,
                             "__pregel_task_id": history[0].tasks[0].id,
-                            "__pregel_scratchpad": {},
-                            "__pregel_writes": AnyList(),
+                            "__pregel_scratchpad": {
+                                "subgraph_counter": AnyInt(),
+                                "call_counter": 0,
+                                "interrupt_counter": -1,
+                                "null_resume": None,
+                                "resume": [],
+                            },
                             "checkpoint_id": c.config["configurable"]["checkpoint_id"],
                             "checkpoint_map": {
                                 "": history[0].config["configurable"]["checkpoint_id"]
@@ -466,10 +484,16 @@ async def test_subgraph_w_interrupt(
                             "__pregel_ensure_latest": True,
                             "__pregel_dedupe_tasks": True,
                             "__pregel_resuming": True,
+                            "__pregel_previous": None,
                             "__pregel_store": None,
                             "__pregel_task_id": history[1].tasks[0].id,
-                            "__pregel_scratchpad": {},
-                            "__pregel_writes": AnyList(),
+                            "__pregel_scratchpad": {
+                                "subgraph_counter": AnyInt(),
+                                "call_counter": 0,
+                                "interrupt_counter": -1,
+                                "null_resume": None,
+                                "resume": [],
+                            },
                             "checkpoint_id": None,
                             "checkpoint_map": {
                                 "": history[1].config["configurable"]["checkpoint_id"]
@@ -528,10 +552,16 @@ async def test_subgraph_w_interrupt(
                             "__pregel_ensure_latest": True,
                             "__pregel_dedupe_tasks": True,
                             "__pregel_resuming": True,
+                            "__pregel_previous": None,
                             "__pregel_store": None,
                             "__pregel_task_id": history[1].tasks[0].id,
-                            "__pregel_scratchpad": {},
-                            "__pregel_writes": AnyList(),
+                            "__pregel_scratchpad": {
+                                "subgraph_counter": AnyInt(),
+                                "call_counter": 0,
+                                "interrupt_counter": -1,
+                                "null_resume": None,
+                                "resume": [],
+                            },
                             "checkpoint_id": c.config["configurable"]["checkpoint_id"],
                             "checkpoint_map": {
                                 "": history[1].config["configurable"]["checkpoint_id"]
@@ -646,10 +676,16 @@ async def test_subgraph_w_interrupt(
                             "__pregel_ensure_latest": True,
                             "__pregel_dedupe_tasks": True,
                             "__pregel_resuming": True,
+                            "__pregel_previous": None,
                             "__pregel_store": None,
                             "__pregel_task_id": history[1].tasks[0].id,
-                            "__pregel_scratchpad": {},
-                            "__pregel_writes": AnyList(),
+                            "__pregel_scratchpad": {
+                                "subgraph_counter": AnyInt(),
+                                "call_counter": 0,
+                                "interrupt_counter": -1,
+                                "null_resume": None,
+                                "resume": [],
+                            },
                             "checkpoint_id": c.config["configurable"]["checkpoint_id"],
                             "checkpoint_map": {
                                 "": history[1].config["configurable"]["checkpoint_id"]

libs/scheduler-kafka/tests/test_subgraph_sync.py

@@ -15,7 +15,7 @@ from langgraph.pregel import Pregel
 from langgraph.scheduler.kafka import serde
 from langgraph.scheduler.kafka.default_sync import DefaultProducer
 from langgraph.scheduler.kafka.types import MessageToOrchestrator, Topics
-from tests.any import AnyDict, AnyList
+from tests.any import AnyDict, AnyInt
 from tests.drain import drain_topics
 from tests.messages import _AnyIdAIMessage, _AnyIdHumanMessage
 
@@ -194,10 +194,16 @@ def test_subgraph_w_interrupt(
                             "__pregel_ensure_latest": True,
                             "__pregel_dedupe_tasks": True,
                             "__pregel_resuming": False,
+                            "__pregel_previous": None,
                             "__pregel_store": None,
                             "__pregel_task_id": history[0].tasks[0].id,
-                            "__pregel_scratchpad": {},
-                            "__pregel_writes": AnyList(),
+                            "__pregel_scratchpad": {
+                                "subgraph_counter": AnyInt(),
+                                "call_counter": 0,
+                                "interrupt_counter": -1,
+                                "null_resume": None,
+                                "resume": [],
+                            },
                             "checkpoint_id": None,
                             "checkpoint_map": {
                                 "": history[0].config["configurable"]["checkpoint_id"]
@@ -262,9 +268,15 @@ def test_subgraph_w_interrupt(
                             "__pregel_store": None,
                             "__pregel_dedupe_tasks": True,
                             "__pregel_resuming": False,
+                            "__pregel_previous": None,
                             "__pregel_task_id": history[0].tasks[0].id,
-                            "__pregel_scratchpad": {},
-                            "__pregel_writes": AnyList(),
+                            "__pregel_scratchpad": {
+                                "subgraph_counter": AnyInt(),
+                                "call_counter": 0,
+                                "interrupt_counter": -1,
+                                "null_resume": None,
+                                "resume": [],
+                            },
                             "checkpoint_id": c.config["configurable"]["checkpoint_id"],
                             "checkpoint_map": {
                                 "": history[0].config["configurable"]["checkpoint_id"]
@@ -360,8 +372,14 @@ def test_subgraph_w_interrupt(
                             "__pregel_store": None,
                             "__pregel_resuming": False,
                             "__pregel_task_id": history[0].tasks[0].id,
-                            "__pregel_scratchpad": {},
-                            "__pregel_writes": AnyList(),
+                            "__pregel_previous": None,
+                            "__pregel_scratchpad": {
+                                "subgraph_counter": AnyInt(),
+                                "call_counter": 0,
+                                "interrupt_counter": -1,
+                                "null_resume": None,
+                                "resume": [],
+                            },
                             "checkpoint_id": c.config["configurable"]["checkpoint_id"],
                             "checkpoint_map": {
                                 "": history[0].config["configurable"]["checkpoint_id"]
@@ -465,9 +483,15 @@ def test_subgraph_w_interrupt(
                             "__pregel_dedupe_tasks": True,
                             "__pregel_store": None,
                             "__pregel_resuming": True,
+                            "__pregel_previous": None,
                             "__pregel_task_id": history[1].tasks[0].id,
-                            "__pregel_scratchpad": {},
-                            "__pregel_writes": AnyList(),
+                            "__pregel_scratchpad": {
+                                "subgraph_counter": AnyInt(),
+                                "call_counter": 0,
+                                "interrupt_counter": -1,
+                                "null_resume": None,
+                                "resume": [],
+                            },
                             "checkpoint_id": None,
                             "checkpoint_map": {
                                 "": history[1].config["configurable"]["checkpoint_id"]
@@ -527,9 +551,15 @@ def test_subgraph_w_interrupt(
                             "__pregel_dedupe_tasks": True,
                             "__pregel_store": None,
                             "__pregel_resuming": True,
+                            "__pregel_previous": None,
                             "__pregel_task_id": history[1].tasks[0].id,
-                            "__pregel_scratchpad": {},
-                            "__pregel_writes": AnyList(),
+                            "__pregel_scratchpad": {
+                                "subgraph_counter": AnyInt(),
+                                "call_counter": 0,
+                                "interrupt_counter": -1,
+                                "null_resume": None,
+                                "resume": [],
+                            },
                             "checkpoint_id": c.config["configurable"]["checkpoint_id"],
                             "checkpoint_map": {
                                 "": history[1].config["configurable"]["checkpoint_id"]
@@ -644,10 +674,16 @@ def test_subgraph_w_interrupt(
                             "__pregel_ensure_latest": True,
                             "__pregel_dedupe_tasks": True,
                             "__pregel_resuming": True,
+                            "__pregel_previous": None,
                             "__pregel_store": None,
                             "__pregel_task_id": history[1].tasks[0].id,
-                            "__pregel_scratchpad": {},
-                            "__pregel_writes": AnyList(),
+                            "__pregel_scratchpad": {
+                                "subgraph_counter": AnyInt(),
+                                "call_counter": 0,
+                                "interrupt_counter": -1,
+                                "null_resume": None,
+                                "resume": [],
+                            },
                             "checkpoint_id": c.config["configurable"]["checkpoint_id"],
                             "checkpoint_map": {
                                 "": history[1].config["configurable"]["checkpoint_id"]