langgraph-checkpoint-postgres 2.0.19 (#3945)

Leverage information about which channels were updated in the previous
step to determine which tasks should be triggered. This can result in
significant speed up in prepare_next_tasks in some situations.

.github/workflows/run_notebooks.yml

@@ -57,13 +57,13 @@ jobs:
         env:
           # these won't actually be used because of the VCR cassettes
           # but need to set them to avoid triggering getpass()
-          OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
-          ANTHROPIC_API_KEY: ${{ secrets.ANTHROPIC_API_KEY }}
-          TAVILY_API_KEY: ${{ secrets.TAVILY_API_KEY }}
-          LANGSMITH_API_KEY: ${{ secrets.LANGSMITH_API_KEY }}
-          NOMIC_API_KEY: ${{ secrets.NOMIC_API_KEY }}
-          COHERE_API_KEY: ${{ secrets.COHERE_API_KEY }}
-          FIREWORKS_API_KEY: ${{ secrets.FIREWORKS_API_KEY }}
+          OPENAI_API_KEY: "very-secret-key"
+          ANTHROPIC_API_KEY: "very-secret-key"
+          TAVILY_API_KEY: "very-secret-key"
+          LANGSMITH_API_KEY: "very-secret-key"
+          NOMIC_API_KEY: "very-secret-key"
+          COHERE_API_KEY: "very-secret-key"
+          FIREWORKS_API_KEY: "very-secret-key"
         run: |
           if [ "${{ github.event_name }}" = "workflow_dispatch" ] || [ "${{ github.event_name }}" = "schedule" ]; then
             echo "Running all notebooks"

docs/README.md

@@ -14,6 +14,8 @@ To run the documentation server locally you can run:
 make serve-docs
 ```
 
+This will start the documentation server on [http://127.0.0.1:8000/langgraph/](http://127.0.0.1:8000/langgraph/).
+
 ## Execute notebooks
 
 If you would like to automatically execute all of the notebooks, to mimic the "Run notebooks" GHA, you can run:

docs/_scripts/third_party_page/packages.yml

@@ -30,6 +30,9 @@ packages:
   - name: "langgraph-bigtool"
     repo: "langchain-ai/langgraph-bigtool"
     description: "Build LangGraph agents with large numbers of tools."
+  - name: "ai-data-science-team"
+    repo: "business-science/ai-data-science-team"
+    description: "An AI-powered data science team of agents to help you perform common data science tasks 10X faster."
   - name: "langgraph-reflection"
     repo: "langchain-ai/langgraph-reflection"
-    description: "LangGraph agent that runs a reflection step."
+    description: "LangGraph agent that runs a reflection step."

docs/docs/cloud/deployment/setup_pyproject.md

@@ -64,7 +64,7 @@ license = "MIT"
 readme = "README.md"
 
 [tool.poetry.dependencies]
-python = ">=3.9.0,<3.13"
+python = ">=3.9"
 langgraph = "^0.2.0"
 langchain-fireworks = "^0.1.3"
 

docs/docs/cloud/how-tos/human_in_the_loop_user_input.md

@@ -63,7 +63,7 @@ Now, let's invoke our graph by interrupting before `ask_human` node:
         "messages": [
             {
                 "role": "user",
-                "content": "Use the search tool to ask the user where they are, then look up the weather there",
+                "content": "Ask the user where they are, then look up the weather there",
             }
         ]
     }
@@ -85,8 +85,7 @@ Now, let's invoke our graph by interrupting before `ask_human` node:
       messages: [
         {
           role: "human",
-          content: "Use the search tool to ask the user where they are, then look up the weather there"
-        }
+          content: "Ask the user where they are, then look up the weather there" }
       ]
     };
 
@@ -115,7 +114,7 @@ Now, let's invoke our graph by interrupting before `ask_human` node:
      --header 'Content-Type: application/json' \
      --data "{
        \"assistant_id\": \"agent\",
-       \"input\": {\"messages\": [{\"role\": \"human\", \"content\": \"Use the search tool to ask the user where they are, then look up the weather there\"}]},
+       \"input\": {\"messages\": [{\"role\": \"human\", \"content\": \"Ask the user where they are, then look up the weather there\"}]},
        \"interrupt_before\": [\"ask_human\"],
        \"stream_mode\": [
          \"updates\"

docs/docs/cloud/how-tos/iterate_graph_studio.md

@@ -1,6 +1,133 @@
 # Prompt Engineering in LangGraph Studio
 
-In LangGraph Studio you can iterate on the prompts used within your graph by utilizing the LangSmith Playground. To do so:
+## Overview
+
+A central aspect of agent development is prompt engineering. LangGraph Studio makes it easy to iterate on the prompts used within your graph directly within the UI.
+
+## Setup
+
+The first step is to define your [configuration](https://langchain-ai.github.io/langgraph/how-tos/configuration/) such that LangGraph Studio is aware of the prompts you want to iterate on and which nodes they are associated with.
+
+### Reference
+
+When defining your configuration, you can use special metadata keys to instruct LangGraph Studio how to handle different fields. Here's a reference for the available configuration options:
+
+#### `langgraph_nodes`
+
+- **Description**: Specifies which graph nodes a configuration field is associated with.
+- **Value Type**: Array of strings, where each string is the name of a node in your graph.
+- **Usage Context**: Include in the `json_schema_extra` dictionary for Pydantic models or the `metadata["json_schema_extra"]` dictionary for dataclasses.
+- **Required**: No, but necessary if you want a field to be editable for specific nodes in the UI.
+- **Example**:
+  ```python
+  system_prompt: str = Field(
+      default="You are a helpful AI assistant.",
+      json_schema_extra={"langgraph_nodes": ["call_model", "other_node"]},
+  )
+  ```
+
+#### `langgraph_type`
+
+- **Description**: Specifies the type of configuration field, which determines how it's handled in the UI.
+- **Value Type**: String
+- **Supported Values**:
+  - `"prompt"`: Indicates the field contains prompt text that should be treated specially in the UI.
+- **Usage Context**: Include in the `json_schema_extra` dictionary for Pydantic models or the `metadata["json_schema_extra"]` dictionary for dataclasses.
+- **Required**: No, but helpful for prompt fields to enable special handling.
+- **Example**:
+  ```python
+  system_prompt: str = Field(
+      default="You are a helpful AI assistant.",
+      json_schema_extra={
+          "langgraph_nodes": ["call_model"],
+          "langgraph_type": "prompt",
+      },
+  )
+  ```
+
+### Example
+
+For example, if you have a node called `call_model` whose system prompt you want to iterate on, you can define a configuration like the following.
+
+```python
+## Using Pydantic
+from pydantic import BaseModel, Field
+from typing import Annotated, Literal
+
+class Configuration(BaseModel):
+    """The configuration for the agent."""
+
+    system_prompt: str = Field(
+        default="You are a helpful AI assistant.",
+        description="The system prompt to use for the agent's interactions. "
+        "This prompt sets the context and behavior for the agent.",
+        json_schema_extra={
+            "langgraph_nodes": ["call_model"],
+            "langgraph_type": "prompt",
+        },
+    )
+
+    model: Annotated[
+        Literal[
+            "anthropic/claude-3-7-sonnet-latest",
+            "anthropic/claude-3-5-haiku-latest",
+            "openai/o1",
+            "openai/gpt-4o-mini",
+            "openai/o1-mini",
+            "openai/o3-mini",
+        ],
+        {"__template_metadata__": {"kind": "llm"}},
+    ] = Field(
+        default="openai/gpt-4o-mini",
+        description="The name of the language model to use for the agent's main interactions. "
+        "Should be in the form: provider/model-name.",
+        json_schema_extra={"langgraph_nodes": ["call_model"]},
+    )
+
+## Using Dataclasses
+from dataclasses import dataclass, field
+
+@dataclass(kw_only=True)
+class Configuration:
+    """The configuration for the agent."""
+
+    system_prompt: str = field(
+        default="You are a helpful AI assistant.",
+        metadata={
+            "description": "The system prompt to use for the agent's interactions. "
+            "This prompt sets the context and behavior for the agent.",
+            "json_schema_extra": {"langgraph_nodes": ["call_model"]},
+        },
+    )
+
+    model: Annotated[str, {"__template_metadata__": {"kind": "llm"}}] = field(
+        default="anthropic/claude-3-5-sonnet-20240620",
+        metadata={
+            "description": "The name of the language model to use for the agent's main interactions. "
+            "Should be in the form: provider/model-name.",
+            "json_schema_extra": {"langgraph_nodes": ["call_model"]},
+        },
+    )
+
+```
+
+## Iterating on prompts
+
+### Node Configuration
+
+With this set up, running your graph and viewing in LangGraph Studio will result in the graph rendering like such.
+
+**Note the configuration icon in the top right corner of the `call_model` node**:
+
+![Graph in Studio](../img/studio_graph_with_configuration.png){width=1200}
+
+Clicking this icon will open a modal where you can edit the configuration for all of the fields associated with the `call_model` node. From here, you can save your changes and apply them to the graph. Note that these values reflect the currently active assistant, and saving will update the assistant with the new values.
+
+![Configuration modal](../img/studio_node_configuration.png){width=1200}
+
+### Playground
+
+LangGraph Studio also supports prompt engineering through an integration with the LangSmith Playground. To do so:
 
 1. Open an existing thread or create a new one.
 2. Within the thread log, any nodes that have made an LLM call will have a "View LLM Runs" button. Clicking this will open a popover with the LLM runs for that node.
@@ -8,8 +135,6 @@ In LangGraph Studio you can iterate on the prompts used within your graph by uti
 
 ![Playground in Studio](../img/studio_playground.png){width=1200}
 
-
-
 From here you can edit the prompt, test different model configurations and re-run just this LLM call without having to re-run the entire graph. When you are happy with your changes, you can copy the updated prompt back into your graph.
 
 For more information on how to use the LangSmith Playground, see the [LangSmith Playground documentation](https://docs.smith.langchain.com/prompt_engineering/how_to_guides#playground).

docs/docs/concepts/agentic_concepts.md

@@ -14,7 +14,7 @@ As a result, there are many different types of [agent architectures](https://blo
 
 ## Router
 
-A router allows an LLM to select a single step from a specified set of options. This is an agent architecture that exhibits a relatively limited level of control because the LLM usually focuses on making a single decision and produces a specific output from limited set of pre-defined options. Routers typically employ a few different concepts to achieve this.
+A router allows an LLM to select a single step from a specified set of options. This is an agent architecture that exhibits a relatively limited level of control because the LLM usually focuses on making a single decision and produces a specific output from a limited set of pre-defined options. Routers typically employ a few different concepts to achieve this.
 
 ### Structured Output
 

docs/docs/concepts/multi_agent.md

@@ -89,7 +89,7 @@ def transfer_to_bob(state):
     )
 ```
 
-This is a special case of updating the graph state from tools where in addition the state update, the control flow is included as well.
+This is a special case of updating the graph state from tools where, in addition to the state update, the control flow is included as well.
 
 !!! important
 
@@ -235,7 +235,7 @@ supervisor = create_react_agent(model, tools)
 
 ### Hierarchical
 
-As you add more agents to your system, it might become too hard for the supervisor to manage all of them. The supervisor might start making poor decisions about which agent to call next, the context might become too complex for a single supervisor to keep track of. In other words, you end up with the same problems that motivated the multi-agent architecture in the first place.
+As you add more agents to your system, it might become too hard for the supervisor to manage all of them. The supervisor might start making poor decisions about which agent to call next, or the context might become too complex for a single supervisor to keep track of. In other words, you end up with the same problems that motivated the multi-agent architecture in the first place.
 
 To address this, you can design your system _hierarchically_. For example, you can create separate, specialized teams of agents managed by individual supervisors, and a top-level supervisor to manage the teams.
 
@@ -339,9 +339,9 @@ builder.add_edge("agent_1", "agent_2")
 
 ## Communication between agents
 
-The most important thing when building multi-agent systems is figuring out how the agents communicate. There are few different considerations:
+The most important thing when building multi-agent systems is figuring out how the agents communicate. There are a few different considerations:
 
-- Do agents communicate via [**via graph state or via tool calls**](#graph-state-vs-tool-calls)?
+- Do agents communicate [**via graph state or via tool calls**](#graph-state-vs-tool-calls)?
 - What if two agents have [**different state schemas**](#different-state-schemas)?
 - How to communicate over a [**shared message list**](#shared-message-list)?
 

docs/docs/how-tos/command.ipynb

@@ -33,7 +33,7 @@
     "    )\n",
     "```\n",
     "\n",
-    "If you are using [subgraphs](#subgraphs), you might want to navigate from a node a subgraph to a different subgraph (i.e. a different node in the parent graph). To do so, you can specify `graph=Command.PARENT` in `Command`:\n",
+    "If you are using [subgraphs](#subgraphs), you might want to navigate from a node within a subgraph to a different subgraph (i.e. a different node in the parent graph). To do so, you can specify `graph=Command.PARENT` in `Command`:\n",
     "\n",
     "```python\n",
     "def my_node(state: State) -> Command[Literal[\"my_other_node\"]]:\n",

docs/docs/how-tos/human_in_the_loop/wait-user-input.ipynb

@@ -397,7 +397,8 @@
     "# We define a fake node to ask the human\n",
     "def ask_human(state):\n",
     "    tool_call_id = state[\"messages\"][-1].tool_calls[0][\"id\"]\n",
-    "    location = interrupt(\"Please provide your location:\")\n",
+    "    ask = AskHuman.model_validate(state[\"messages\"][-1].tool_calls[0][\"args\"])\n",
+    "    location = interrupt(ask.question)\n",
     "    tool_message = [{\"tool_call_id\": tool_call_id, \"type\": \"tool\", \"content\": location}]\n",
     "    return {\"messages\": tool_message}\n",
     "\n",
@@ -491,7 +492,7 @@
     "        \"messages\": [\n",
     "            (\n",
     "                \"user\",\n",
-    "                \"Use the search tool to ask the user where they are, then look up the weather there\",\n",
+    "                \"Ask the user where they are, then look up the weather there\",\n",
     "            )\n",
     "        ]\n",
     "    },\n",

docs/docs/how-tos/memory/add-summary-conversation-history.ipynb

@@ -99,7 +99,7 @@
     "from typing import Literal\n",
     "\n",
     "from langchain_anthropic import ChatAnthropic\n",
-    "from langchain_core.messages import SystemMessage, RemoveMessage\n",
+    "from langchain_core.messages import SystemMessage, RemoveMessage, HumanMessage\n",
     "from langgraph.checkpoint.memory import MemorySaver\n",
     "from langgraph.graph import MessagesState, StateGraph, START, END\n",
     "\n",

docs/docs/how-tos/memory/manage-conversation-history.ipynb

@@ -7,7 +7,7 @@
    "source": [
     "# How to manage conversation history\n",
     "\n",
-    "One of the most common use cases for persistence is to use it to keep track of conversation history. This is great - it makes it easy to continue conversations. As conversations get longer and longer, however, this conversation history can build up and take up more and more of the context window. This can often be undesirable as it leads to more expensive and longer calls to the LLM, and potentially ones that error. In order to prevent this from happening, you need to probably manage the conversation history.\n",
+    "One of the most common use cases for persistence is to use it to keep track of conversation history. This is great - it makes it easy to continue conversations. As conversations get longer and longer, however, this conversation history can build up and take up more and more of the context window. This can often be undesirable as it leads to more expensive and longer calls to the LLM, and potentially ones that error. In order to prevent this from happening, you need to properly manage the conversation history.\n",
     "\n",
     "Note: this guide focuses on how to do this in LangGraph, where you can fully customize how this is done. If you want a more off-the-shelf solution, you can look into functionality provided in LangChain:\n",
     "\n",

docs/docs/how-tos/pass-run-time-values-to-tools.ipynb

@@ -38,7 +38,7 @@
     "    </p>\n",
     "</div> \n",
     "\n",
-    "The core technique the examples below is to **annotate** a parameter as \"injected\", meaning it will be injected by your program and should not be seen or populated by the LLM. Let the following codesnippet serve as a tl;dr:\n",
+    "The core technique in the examples below is to **annotate** a parameter as \"injected\", meaning it will be injected by your program and should not be seen or populated by the LLM. Let the following codesnippet serve as a tl;dr:\n",
     "\n",
     "```python\n",
     "from typing import Annotated\n",

docs/docs/how-tos/react-agent-structured-output.ipynb

@@ -65,7 +65,7 @@
     "\n",
     "**Pros and Cons**\n",
     "\n",
-    "The benefit to this format is that you only need one LLM, and can save money and latency because of this. The downside to this option is that it isn't guaranteed that the single LLM will call the correct tool when you want it to. We can help the LLM by setting `tool_choice` to `any` when we use `bind_tools` which forces the LLM to select at least one tool at every turn, but this is far from a fool proof strategy. In addition, another downside is that the agent might call *multiple* tools, so we need to check for this explicitly in our routing function (or if we are using OpenAI we an set `parallell_tool_calling=False` to ensure only one tool is called at a time).\n",
+    "The benefit to this format is that you only need one LLM, and can save money and latency because of this. The downside to this option is that it isn't guaranteed that the single LLM will call the correct tool when you want it to. We can help the LLM by setting `tool_choice` to `any` when we use `bind_tools` which forces the LLM to select at least one tool at every turn, but this is far from a foolproof strategy. In addition, another downside is that the agent might call *multiple* tools, so we need to check for this explicitly in our routing function (or if we are using OpenAI we can set `parallell_tool_calling=False` to ensure only one tool is called at a time).\n",
     "\n",
     "**Option 2**\n",
     "\n",

docs/docs/how-tos/state-model.ipynb

@@ -266,6 +266,235 @@
     "    print(\"An exception was raised because bad_node sets `a` to an integer.\")\n",
     "    print(e)"
    ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "2270bc3c",
+   "metadata": {},
+   "source": [
+    "## Multiple Nodes\n",
+    "\n",
+    "Run-time validation will also work in a multi-node graph. In the example below `bad_node` updates `a` to an integer. \n",
+    "\n",
+    "Because run-time validation occurs on **inputs**, the validation error will occur when `ok_node` is called (not when `bad_node` returns an update to the state which is inconsistent with the schema)."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d832cdcc",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from langgraph.graph import StateGraph, START, END\n",
+    "from typing_extensions import TypedDict\n",
+    "\n",
+    "from pydantic import BaseModel\n",
+    "\n",
+    "\n",
+    "# The overall state of the graph (this is the public state shared across nodes)\n",
+    "class OverallState(BaseModel):\n",
+    "    a: str\n",
+    "\n",
+    "\n",
+    "def bad_node(state: OverallState):\n",
+    "    return {\n",
+    "        \"a\": 123  # Invalid\n",
+    "    }\n",
+    "\n",
+    "\n",
+    "def ok_node(state: OverallState):\n",
+    "    return {\"a\": \"goodbye\"}\n",
+    "\n",
+    "\n",
+    "# Build the state graph\n",
+    "builder = StateGraph(OverallState)\n",
+    "builder.add_node(bad_node)\n",
+    "builder.add_node(ok_node)\n",
+    "builder.add_edge(START, \"bad_node\")\n",
+    "builder.add_edge(\"bad_node\", \"ok_node\")\n",
+    "builder.add_edge(\"ok_node\", END)\n",
+    "graph = builder.compile()\n",
+    "\n",
+    "# Test the graph with a valid input\n",
+    "try:\n",
+    "    graph.invoke({\"a\": \"hello\"})\n",
+    "except Exception as e:\n",
+    "    print(\"An exception was raised because bad_node sets `a` to an integer.\")\n",
+    "    print(e)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "456b1f77",
+   "metadata": {},
+   "source": [
+    "## Advanced Pydantic Model Usage\n",
+    "\n",
+    "This section covers more advanced topics when using Pydantic models with LangGraph.\n",
+    "\n",
+    "### Serialization Behavior\n",
+    "\n",
+    "When using Pydantic models as state schemas, it's important to understand how serialization works, especially when:\n",
+    "- Passing Pydantic objects as inputs\n",
+    "- Receiving outputs from the graph\n",
+    "- Working with nested Pydantic models\n",
+    "\n",
+    "Let's see these behaviors in action:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "0e919cdc",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from langgraph.graph import StateGraph, START, END\n",
+    "from pydantic import BaseModel\n",
+    "\n",
+    "\n",
+    "class NestedModel(BaseModel):\n",
+    "    value: str\n",
+    "\n",
+    "\n",
+    "class ComplexState(BaseModel):\n",
+    "    text: str\n",
+    "    count: int\n",
+    "    nested: NestedModel\n",
+    "\n",
+    "\n",
+    "def process_node(state: ComplexState):\n",
+    "    # Node receives a validated Pydantic object\n",
+    "    print(f\"Input state type: {type(state)}\")\n",
+    "    print(f\"Nested type: {type(state.nested)}\")\n",
+    "\n",
+    "    # Return a dictionary update\n",
+    "    return {\"text\": state.text + \" processed\", \"count\": state.count + 1}\n",
+    "\n",
+    "\n",
+    "# Build the graph\n",
+    "builder = StateGraph(ComplexState)\n",
+    "builder.add_node(\"process\", process_node)\n",
+    "builder.add_edge(START, \"process\")\n",
+    "builder.add_edge(\"process\", END)\n",
+    "graph = builder.compile()\n",
+    "\n",
+    "# Create a Pydantic instance for input\n",
+    "input_state = ComplexState(text=\"hello\", count=0, nested=NestedModel(value=\"test\"))\n",
+    "print(f\"Input object type: {type(input_state)}\")\n",
+    "\n",
+    "# Invoke graph with a Pydantic instance\n",
+    "result = graph.invoke(input_state)\n",
+    "print(f\"Output type: {type(result)}\")\n",
+    "print(f\"Output content: {result}\")\n",
+    "\n",
+    "# Convert back to Pydantic model if needed\n",
+    "output_model = ComplexState(**result)\n",
+    "print(f\"Converted back to Pydantic: {type(output_model)}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "f13f28ce",
+   "metadata": {},
+   "source": [
+    "### Runtime Type Coercion\n",
+    "\n",
+    "Pydantic performs runtime type coercion for certain data types. This can be helpful but also lead to unexpected behavior if you're not aware of it."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "faf59316",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from langgraph.graph import StateGraph, START, END\n",
+    "from pydantic import BaseModel\n",
+    "\n",
+    "\n",
+    "class CoercionExample(BaseModel):\n",
+    "    # Pydantic will coerce string numbers to integers\n",
+    "    number: int\n",
+    "    # Pydantic will parse string booleans to bool\n",
+    "    flag: bool\n",
+    "\n",
+    "\n",
+    "def inspect_node(state: CoercionExample):\n",
+    "    print(f\"number: {state.number} (type: {type(state.number)})\")\n",
+    "    print(f\"flag: {state.flag} (type: {type(state.flag)})\")\n",
+    "    return {}\n",
+    "\n",
+    "\n",
+    "builder = StateGraph(CoercionExample)\n",
+    "builder.add_node(\"inspect\", inspect_node)\n",
+    "builder.add_edge(START, \"inspect\")\n",
+    "builder.add_edge(\"inspect\", END)\n",
+    "graph = builder.compile()\n",
+    "\n",
+    "# Demonstrate coercion with string inputs that will be converted\n",
+    "result = graph.invoke({\"number\": \"42\", \"flag\": \"true\"})\n",
+    "\n",
+    "# This would fail with a validation error\n",
+    "try:\n",
+    "    graph.invoke({\"number\": \"not-a-number\", \"flag\": \"true\"})\n",
+    "except Exception as e:\n",
+    "    print(f\"\\nExpected validation error: {e}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "2844475b",
+   "metadata": {},
+   "source": [
+    "### Working with Message Models\n",
+    "\n",
+    "When working with LangChain message types in your state schema, there are important considerations for serialization. You should use `AnyMessage` (rather than `BaseMessage`) for proper serialization/deserialization when using message objects over the wire:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "bd0734b0",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from langgraph.graph import StateGraph, START, END\n",
+    "from pydantic import BaseModel\n",
+    "from langchain_core.messages import HumanMessage, AIMessage, AnyMessage\n",
+    "from typing import List\n",
+    "\n",
+    "\n",
+    "class ChatState(BaseModel):\n",
+    "    messages: List[AnyMessage]\n",
+    "    context: str\n",
+    "\n",
+    "\n",
+    "def add_message(state: ChatState):\n",
+    "    return {\"messages\": state.messages + [AIMessage(content=\"Hello there!\")]}\n",
+    "\n",
+    "\n",
+    "builder = StateGraph(ChatState)\n",
+    "builder.add_node(\"add_message\", add_message)\n",
+    "builder.add_edge(START, \"add_message\")\n",
+    "builder.add_edge(\"add_message\", END)\n",
+    "graph = builder.compile()\n",
+    "\n",
+    "# Create input with a message\n",
+    "initial_state = ChatState(\n",
+    "    messages=[HumanMessage(content=\"Hi\")], context=\"Customer support chat\"\n",
+    ")\n",
+    "\n",
+    "result = graph.invoke(initial_state)\n",
+    "print(f\"Output: {result}\")\n",
+    "\n",
+    "# Convert back to Pydantic model to see message types\n",
+    "output_model = ChatState(**result)\n",
+    "for i, msg in enumerate(output_model.messages):\n",
+    "    print(f\"Message {i}: {type(msg).__name__} - {msg.content}\")"
+   ]
   }
  ],
  "metadata": {

docs/docs/how-tos/update-state-from-tools.ipynb

@@ -210,7 +210,7 @@
    "id": "cbb06aea-6654-4245-91f8-af6e8f2b5377",
    "metadata": {},
    "source": [
-    "Let's now add personalization: we'll respond differently to the user based on the state values AFTER the state has been updated from the tool. To achieve this, let's define a function that will dynamically construct the system prompt based on the graph state. It will be called ever time the LLM is called and the function output will be passed to the LLM:"
+    "Let's now add personalization: we'll respond differently to the user based on the state values AFTER the state has been updated from the tool. To achieve this, let's define a function that will dynamically construct the system prompt based on the graph state. It will be called every time the LLM is called and the function output will be passed to the LLM:"
    ]
   },
   {

libs/checkpoint-postgres/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "langgraph-checkpoint-postgres"
-version = "2.0.18"
+version = "2.0.19"
 description = "Library with a Postgres implementation of LangGraph checkpoint saver."
 authors = []
 license = "MIT"
@@ -10,7 +10,7 @@ packages = [{ include = "langgraph" }]
 
 [tool.poetry.dependencies]
 python = "^3.9.0,<4.0"
-langgraph-checkpoint = "^2.0.15"
+langgraph-checkpoint = "^2.0.21"
 orjson = ">=3.10.1"
 psycopg = "^3.2.0"
 psycopg-pool = "^3.2.0"

libs/checkpoint-postgres/tests/test_async_store.py

@@ -67,6 +67,10 @@ async def store(request) -> AsyncIterator[AsyncPostgresStore]:
                 for mig in store.MIGRATIONS
             ]
             await store.setup()
+            async with store._cursor() as cur:
+                # drop the migration index
+                await cur.execute("DROP TABLE IF EXISTS store_migrations")
+            await store.setup()  # Will fail if migrations aren't idempotent
 
         if request.param == "pipe":
             async with AsyncPostgresStore.from_conn_string(

libs/checkpoint-postgres/tests/test_store.py

@@ -413,6 +413,10 @@ def _create_vector_store(
             ttl={"default_ttl": 2, "refresh_on_read": True} if enable_ttl else None,
         ) as store:
             store.setup()
+            with store._cursor() as cur:
+                # drop the migration index
+                cur.execute("DROP TABLE IF EXISTS store_migrations")
+            store.setup()  # Will fail if migrations aren't idempotent
             yield store
     finally:
         with Connection.connect(admin_conn_string, autocommit=True) as conn:

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

@@ -56,7 +56,10 @@ class SqliteSaver(BaseCheckpointSaver[str]):
         >>> builder.add_node("add_one", lambda x: x + 1)
         >>> builder.set_entry_point("add_one")
         >>> builder.set_finish_point("add_one")
-        >>> conn = sqlite3.connect("checkpoints.sqlite")
+        >>> # Create a new SqliteSaver instance
+        >>> # Note: check_same_thread=False is OK as the implementation uses a lock
+        >>> # to ensure thread safety.
+        >>> conn = sqlite3.connect("checkpoints.sqlite", check_same_thread=False)
         >>> memory = SqliteSaver(conn)
         >>> graph = builder.compile(checkpointer=memory)
         >>> config = {"configurable": {"thread_id": "1"}}

libs/checkpoint/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "langgraph-checkpoint"
-version = "2.0.20"
+version = "2.0.21"
 description = "Library with base interfaces for LangGraph checkpoint savers."
 authors = []
 license = "MIT"

libs/cli/README.md

@@ -79,7 +79,7 @@ The CLI uses a `langgraph.json` configuration file with these key settings:
 }
 ```
 
-See the [full documentation](https://langchain-ai.github.io/langgraph/docs/cloud/reference/cli.html) for detailed configuration options.
+See the [full documentation](https://langchain-ai.github.io/langgraph/cloud/reference/cli/) for detailed configuration options.
 
 ## Development
 

libs/langgraph/Makefile

@@ -58,9 +58,11 @@ WORKERS ?= auto
 XDIST_ARGS := $(if $(WORKERS),-n $(WORKERS) --dist worksteal,)
 MAXFAIL ?=
 MAXFAIL_ARGS := $(if $(MAXFAIL),--maxfail $(MAXFAIL),)
+# Add an '-x' if xdist is enabled
+XDIST_ARGS := $(if $(WORKERS),-x $(XDIST_ARGS),)
 
 test_watch:
-	make start-postgres && poetry run ptw . -- --ff -vv -x $(XDIST_ARGS) $(MAXFAIL_ARGS) --snapshot-update --tb short $(TEST); \
+	make start-postgres && poetry run ptw . -- --ff -vv $(XDIST_ARGS) $(MAXFAIL_ARGS) --snapshot-update --tb short $(TEST); \
 	EXIT_CODE=$$?; \
 	make stop-postgres; \
 	exit $$EXIT_CODE

libs/langgraph/bench/__main__.py

@@ -6,7 +6,9 @@ from pyperf._runner import Runner
 from uvloop import new_event_loop
 
 from bench.fanout_to_subgraph import fanout_to_subgraph, fanout_to_subgraph_sync
+from bench.pydantic_state import pydantic_state
 from bench.react_agent import react_agent
+from bench.sequential import create_sequential
 from bench.wide_state import wide_state
 from langgraph.checkpoint.memory import MemorySaver
 from langgraph.pregel import Pregel
@@ -203,6 +205,126 @@ benchmarks = (
             ]
         },
     ),
+    (
+        "sequential_20",
+        create_sequential(20).compile(),
+        create_sequential(20).compile(),
+        {"messages": []},  # Empty list of messages
+    ),
+    (
+        "sequential_50",
+        create_sequential(50).compile(),
+        create_sequential(50).compile(),
+        {"messages": []},  # Empty list of messages
+    ),
+    (
+        "sequential_100",
+        create_sequential(100).compile(),
+        create_sequential(100).compile(),
+        {"messages": []},  # Empty list of messages
+    ),
+    (
+        "sequential_200",
+        create_sequential(200).compile(),
+        create_sequential(200).compile(),
+        {"messages": []},  # Empty list of messages
+    ),
+    (
+        "pydantic_state_25x300",
+        pydantic_state(300).compile(checkpointer=None),
+        pydantic_state(300).compile(checkpointer=None),
+        {
+            "messages": [
+                {
+                    str(i) * 10: {
+                        str(j) * 10: ["hi?" * 10, True, 1, 6327816386138, None] * 5
+                        for j in range(5)
+                    }
+                    for i in range(5)
+                }
+            ]
+        },
+    ),
+    (
+        "pydantic_state_25x300_checkpoint",
+        pydantic_state(300).compile(checkpointer=MemorySaver()),
+        pydantic_state(300).compile(checkpointer=MemorySaver()),
+        {
+            "messages": [
+                {
+                    str(i) * 10: {
+                        str(j) * 10: ["hi?" * 10, True, 1, 6327816386138, None] * 5
+                        for j in range(5)
+                    }
+                    for i in range(5)
+                }
+            ]
+        },
+    ),
+    (
+        "pydantic_state_15x600",
+        pydantic_state(600).compile(checkpointer=None),
+        pydantic_state(600).compile(checkpointer=None),
+        {
+            "messages": [
+                {
+                    str(i) * 10: {
+                        str(j) * 10: ["hi?" * 10, True, 1, 6327816386138, None] * 5
+                        for j in range(5)
+                    }
+                    for i in range(3)
+                }
+            ]
+        },
+    ),
+    (
+        "pydantic_state_15x600_checkpoint",
+        pydantic_state(600).compile(checkpointer=MemorySaver()),
+        pydantic_state(600).compile(checkpointer=MemorySaver()),
+        {
+            "messages": [
+                {
+                    str(i) * 10: {
+                        str(j) * 10: ["hi?" * 10, True, 1, 6327816386138, None] * 5
+                        for j in range(5)
+                    }
+                    for i in range(3)
+                }
+            ]
+        },
+    ),
+    (
+        "pydantic_state_9x1200",
+        pydantic_state(1200).compile(checkpointer=None),
+        pydantic_state(1200).compile(checkpointer=None),
+        {
+            "messages": [
+                {
+                    str(i) * 10: {
+                        str(j) * 10: ["hi?" * 10, True, 1, 6327816386138, None] * 5
+                        for j in range(3)
+                    }
+                    for i in range(3)
+                }
+            ]
+        },
+    ),
+    (
+        "pydantic_state_9x1200_checkpoint",
+        pydantic_state(1200).compile(checkpointer=MemorySaver()),
+        pydantic_state(1200).compile(checkpointer=MemorySaver()),
+        {
+            "messages": [
+                {
+                    str(i) * 10: {
+                        str(j) * 10: ["hi?" * 10, True, 1, 6327816386138, None] * 5
+                        for j in range(3)
+                    }
+                    for i in range(3)
+                }
+            ]
+        },
+    ),
 )
 
 

libs/langgraph/bench/pydantic_state.py

@@ -0,0 +1,327 @@
+import operator
+from functools import partial
+from random import choice
+from typing import Annotated, Optional, Sequence
+
+from pydantic import BaseModel, Field, field_validator
+
+from langgraph.constants import END, START
+from langgraph.graph.state import StateGraph
+
+
+def pydantic_state(n: int) -> StateGraph:
+    class State(BaseModel):
+        messages: Annotated[list, operator.add] = Field(default_factory=list)
+
+        @field_validator("messages", mode="after")
+        @classmethod
+        def validate_messages(cls, v):
+            if not isinstance(v, list):
+                raise TypeError("messages must be a list")
+            for msg in v:
+                if not isinstance(msg, dict):
+                    raise TypeError("messages must be a list of dicts")
+                if not all(isinstance(k, str) for k in msg.keys()):
+                    raise TypeError("messages must be a list of dicts with str keys")
+            return v
+
+        trigger_events: Annotated[list, operator.add] = Field(default_factory=list)
+        """The external events that are converted by the graph."""
+
+        @field_validator("trigger_events", mode="after")
+        @classmethod
+        def validate_trigger_events(cls, v):
+            if not isinstance(v, list):
+                raise TypeError("trigger_events must be a list")
+            for event in v:
+                if not isinstance(event, dict):
+                    raise TypeError("trigger_events must be a list of dicts")
+                if not all(isinstance(k, str) for k in event.keys()):
+                    raise TypeError(
+                        "trigger_events must be a list of dicts with str keys"
+                    )
+            return v
+
+        primary_issue_medium: Annotated[str, lambda x, y: y or x] = Field(
+            default="email"
+        )
+        """The primary issue medium for the current conversation."""
+
+        @field_validator("primary_issue_medium", mode="after")
+        @classmethod
+        def validate_primary_issue_medium(cls, v):
+            if not isinstance(v, str):
+                raise TypeError("primary_issue_medium must be a string")
+            return v
+
+        autoresponse: Annotated[Optional[dict], lambda _, y: y] = Field(
+            default=None
+        )  # Always overwrite
+
+        @field_validator("autoresponse", mode="after")
+        @classmethod
+        def validate_autoresponse(cls, v):
+            if v is not None and not isinstance(v, dict):
+                raise TypeError("autoresponse must be a dict or None")
+            return v
+
+        issue: Annotated[dict | None, lambda x, y: y if y else x] = Field(default=None)
+
+        @field_validator("issue", mode="after")
+        @classmethod
+        def validate_issue(cls, v):
+            if v is not None and not isinstance(v, dict):
+                raise TypeError("issue must be a dict or None")
+            return v
+
+        relevant_rules: Optional[list[dict]] = Field(default=None)
+        """SOPs fetched from the rulebook that are relevant to the current conversation."""
+
+        @field_validator("relevant_rules", mode="after")
+        @classmethod
+        def validate_relevant_rules(cls, v):
+            if v is None:
+                return v
+            if not isinstance(v, list):
+                raise TypeError("relevant_rules must be a list or None")
+            for rule in v:
+                if not isinstance(rule, dict):
+                    raise TypeError("relevant_rules must be a list of dicts")
+                if not all(isinstance(k, str) for k in rule.keys()):
+                    raise TypeError(
+                        "relevant_rules must be a list of dicts with str keys"
+                    )
+            return v
+
+        memory_docs: Optional[list[dict]] = Field(default=None)
+        """Memory docs fetched from the memory service that are relevant to the current conversation."""
+
+        @field_validator("memory_docs", mode="after")
+        @classmethod
+        def validate_memory_docs(cls, v):
+            if v is None:
+                return v
+            if not isinstance(v, list):
+                raise TypeError("memory_docs must be a list or None")
+            for doc in v:
+                if not isinstance(doc, dict):
+                    raise TypeError("memory_docs must be a list of dicts")
+                if not all(isinstance(k, str) for k in doc.keys()):
+                    raise TypeError("memory_docs must be a list of dicts with str keys")
+            return v
+
+        categorizations: Annotated[list[dict], operator.add] = Field(
+            default_factory=list
+        )
+        """The issue categorizations auto-generated by the AI."""
+
+        @field_validator("categorizations", mode="after")
+        @classmethod
+        def validate_categorizations(cls, v):
+            if not isinstance(v, list):
+                raise TypeError("categorizations must be a list")
+            for categorization in v:
+                if not isinstance(categorization, dict):
+                    raise TypeError("categorizations must be a list of dicts")
+                if not all(isinstance(k, str) for k in categorization.keys()):
+                    raise TypeError(
+                        "categorizations must be a list of dicts with str keys"
+                    )
+            return v
+
+        responses: Annotated[list[dict], operator.add] = Field(default_factory=list)
+        """The draft responses recommended by the AI."""
+
+        @field_validator("responses", mode="after")
+        @classmethod
+        def validate_responses(cls, v):
+            if not isinstance(v, list):
+                raise TypeError("responses must be a list")
+            for response in v:
+                if not isinstance(response, dict):
+                    raise TypeError("responses must be a list of dicts")
+                if not all(isinstance(k, str) for k in response.keys()):
+                    raise TypeError("responses must be a list of dicts with str keys")
+            return v
+
+        user_info: Annotated[Optional[dict], lambda x, y: y if y is not None else x] = (
+            Field(default=None)
+        )
+        """The current user state (by email)."""
+
+        @field_validator("user_info", mode="after")
+        @classmethod
+        def validate_user_info(cls, v):
+            if v is not None and not isinstance(v, dict):
+                raise TypeError("user_info must be a dict or None")
+            return v
+
+        crm_info: Annotated[Optional[dict], lambda x, y: y if y is not None else x] = (
+            Field(default=None)
+        )
+        """The CRM information for organization the current user is from."""
+
+        @field_validator("crm_info", mode="after")
+        @classmethod
+        def validate_crm_info(cls, v):
+            if v is not None and not isinstance(v, dict):
+                raise TypeError("crm_info must be a dict or None")
+            return v
+
+        email_thread_id: Annotated[
+            Optional[str], lambda x, y: y if y is not None else x
+        ] = Field(default=None)
+        """The current email thread ID."""
+
+        @field_validator("email_thread_id", mode="after")
+        @classmethod
+        def validate_email_thread_id(cls, v):
+            if v is not None and not isinstance(v, str):
+                raise TypeError("email_thread_id must be a string or None")
+            return v
+
+        slack_participants: Annotated[dict, operator.or_] = Field(default_factory=dict)
+        """The growing list of current slack participants."""
+
+        @field_validator("slack_participants", mode="after")
+        @classmethod
+        def validate_slack_participants(cls, v):
+            if not isinstance(v, dict):
+                raise TypeError("slack_participants must be a dict")
+            for participant in v:
+                if not isinstance(participant, str):
+                    raise TypeError("slack_participants must be a dict with str keys")
+            return v
+
+        bot_id: Optional[str] = Field(default=None)
+        """The ID of the bot user in the slack channel."""
+
+        @field_validator("bot_id", mode="after")
+        @classmethod
+        def validate_bot_id(cls, v):
+            if v is not None and not isinstance(v, str):
+                raise TypeError("bot_id must be a string or None")
+            return v
+
+        notified_assignees: Annotated[dict, operator.or_] = Field(default_factory=dict)
+
+        @field_validator("notified_assignees", mode="after")
+        def validate_notified_assignees(cls, v):
+            if not isinstance(v, dict):
+                raise TypeError("notified_assignees must be a dict")
+            for assignee in v:
+                if not isinstance(assignee, str):
+                    raise TypeError("notified_assignees must be a dict with str keys")
+            return v
+
+    list_fields = {
+        "messages",
+        "trigger_events",
+        "categorizations",
+        "responses",
+        "memory_docs",
+        "relevant_rules",
+    }
+    dict_fields = {
+        "user_info",
+        "crm_info",
+        "slack_participants",
+        "notified_assignees",
+        "autoresponse",
+        "issue",
+    }
+
+    def read_write(read: str, write: Sequence[str], input: State) -> dict:
+        val = getattr(input, read)
+        val = {val: val} if isinstance(val, str) else val
+        val_single = val[-1] if isinstance(val, list) else val
+        val_list = val if isinstance(val, list) else [val]
+        return {
+            k: val_list
+            if k in list_fields
+            else val_single
+            if k in dict_fields
+            else "".join(choice("abcdefghijklmnopqrstuvwxyz") for _ in range(n))
+            for k in write
+        }
+
+    builder = StateGraph(State)
+    builder.add_edge(START, "one")
+    builder.add_node(
+        "one",
+        partial(read_write, "messages", ["trigger_events", "primary_issue_medium"]),
+    )
+    builder.add_edge("one", "two")
+    builder.add_node(
+        "two",
+        partial(read_write, "trigger_events", ["autoresponse", "issue"]),
+    )
+    builder.add_edge("two", "three")
+    builder.add_edge("two", "four")
+    builder.add_node(
+        "three",
+        partial(read_write, "autoresponse", ["relevant_rules"]),
+    )
+    builder.add_node(
+        "four",
+        partial(
+            read_write,
+            "trigger_events",
+            ["categorizations", "responses", "memory_docs"],
+        ),
+    )
+    builder.add_node(
+        "five",
+        partial(
+            read_write,
+            "categorizations",
+            [
+                "user_info",
+                "crm_info",
+                "email_thread_id",
+                "slack_participants",
+                "bot_id",
+                "notified_assignees",
+            ],
+        ),
+    )
+    builder.add_edge(["three", "four"], "five")
+    builder.add_edge("five", "six")
+    builder.add_node(
+        "six",
+        partial(read_write, "responses", ["messages"]),
+    )
+    builder.add_conditional_edges(
+        "six", lambda state: END if len(state.messages) > n else "one"
+    )
+
+    return builder
+
+
+if __name__ == "__main__":
+    import asyncio
+
+    import uvloop
+
+    from langgraph.checkpoint.memory import MemorySaver
+
+    graph = pydantic_state(1000).compile(checkpointer=MemorySaver())
+    input = {
+        "messages": [
+            {
+                str(i) * 10: {
+                    str(j) * 10: ["hi?" * 10, True, 1, 6327816386138, None] * 5
+                    for j in range(5)
+                }
+                for i in range(5)
+            }
+        ]
+    }
+    config = {"configurable": {"thread_id": "1"}, "recursion_limit": 20000000000}
+
+    async def run():
+        async for c in graph.astream(input, config=config):
+            print(c.keys())
+
+    uvloop.install()
+    asyncio.run(run())

libs/langgraph/bench/sequential.py

@@ -0,0 +1,48 @@
+"""Create a sequential no-op graph consisting of a few hundred nodes."""
+
+from langgraph.graph import MessagesState, StateGraph
+from langgraph.utils.runnable import RunnableCallable
+
+
+def create_sequential(number_nodes) -> StateGraph:
+    """Create a sequential no-op graph consisting of a few hundred nodes."""
+    builder = StateGraph(MessagesState)
+
+    def noop(state: MessagesState) -> None:
+        """No-op function."""
+        pass
+
+    async def anoop(state: MessagesState) -> None:
+        """No-op function."""
+        pass
+
+    prev_node = "__start__"
+
+    for i in range(number_nodes):
+        name = f"node_{i}"
+        builder.add_node(name, RunnableCallable(noop, anoop))
+        builder.add_edge(prev_node, name)
+        prev_node = name
+
+    builder.add_edge(prev_node, "__end__")
+    return builder
+
+
+if __name__ == "__main__":
+    import asyncio
+    import time
+
+    import uvloop
+
+    graph = create_sequential(2000).compile()
+    input = {"messages": []}  # Empty list of messages
+    config = {"recursion_limit": 20000000000}
+
+    async def run():
+        len([c async for c in graph.astream(input, config=config)])
+
+    uvloop.install()
+    start = time.time()
+    asyncio.run(run())
+    end = time.time()
+    print(f"Time taken: {end - start:.4f} seconds")

libs/langgraph/langgraph/channels/any_value.py

@@ -1,8 +1,9 @@
-from typing import Generic, Optional, Sequence, Type
+from typing import Any, Generic, Optional, Sequence, Type
 
 from typing_extensions import Self
 
 from langgraph.channels.base import BaseChannel, Value
+from langgraph.constants import MISSING
 from langgraph.errors import EmptyChannelError
 
 
@@ -12,6 +13,10 @@ class AnyValue(Generic[Value], BaseChannel[Value, Value, Value]):
 
     __slots__ = ("typ", "value")
 
+    def __init__(self, typ: Any, key: str = "") -> None:
+        super().__init__(typ, key)
+        self.value = MISSING
+
     def __eq__(self, value: object) -> bool:
         return isinstance(value, AnyValue)
 
@@ -34,17 +39,19 @@ class AnyValue(Generic[Value], BaseChannel[Value, Value, Value]):
 
     def update(self, values: Sequence[Value]) -> bool:
         if len(values) == 0:
-            try:
-                del self.value
-                return True
-            except AttributeError:
+            if self.value is MISSING:
                 return False
+            else:
+                self.value = MISSING
+                return True
 
         self.value = values[-1]
         return True
 
     def get(self) -> Value:
-        try:
-            return self.value
-        except AttributeError:
+        if self.value is MISSING:
             raise EmptyChannelError()
+        return self.value
+
+    def is_available(self) -> bool:
+        return self.value is not MISSING

libs/langgraph/langgraph/channels/base.py

@@ -64,6 +64,17 @@ class BaseChannel(Generic[Value, Update, C], ABC):
         """
         return False
 
+    def is_available(self) -> bool:
+        """Return True if the channel is available (not empty), False otherwise.
+        Subclasses should override this method to provide a more efficient
+        implementation than calling get() and catching EmptyChannelError.
+        """
+        try:
+            self.get()
+            return True
+        except EmptyChannelError:
+            return False
+
 
 __all__ = [
     "BaseChannel",

libs/langgraph/langgraph/channels/binop.py

@@ -10,6 +10,7 @@ from typing import (
 from typing_extensions import NotRequired, Required, Self
 
 from langgraph.channels.base import BaseChannel, Value
+from langgraph.constants import MISSING
 from langgraph.errors import EmptyChannelError
 
 
@@ -51,7 +52,7 @@ class BinaryOperatorAggregate(Generic[Value], BaseChannel[Value, Value, Value]):
         try:
             self.value = typ()
         except Exception:
-            pass
+            self.value = MISSING
 
     def __eq__(self, value: object) -> bool:
         return isinstance(value, BinaryOperatorAggregate) and (
@@ -81,7 +82,7 @@ class BinaryOperatorAggregate(Generic[Value], BaseChannel[Value, Value, Value]):
     def update(self, values: Sequence[Value]) -> bool:
         if not values:
             return False
-        if not hasattr(self, "value"):
+        if self.value is MISSING:
             self.value = values[0]
             values = values[1:]
         for value in values:
@@ -89,7 +90,9 @@ class BinaryOperatorAggregate(Generic[Value], BaseChannel[Value, Value, Value]):
         return True
 
     def get(self) -> Value:
-        try:
-            return self.value
-        except AttributeError:
+        if self.value is MISSING:
             raise EmptyChannelError()
+        return self.value
+
+    def is_available(self) -> bool:
+        return self.value is not MISSING

libs/langgraph/langgraph/channels/dynamic_barrier_value.py

@@ -85,6 +85,9 @@ class DynamicBarrierValue(
             raise EmptyChannelError()
         return None
 
+    def is_available(self) -> bool:
+        return self.seen == self.names
+
     def consume(self) -> bool:
         if self.seen == self.names:
             self.seen = set()

libs/langgraph/langgraph/channels/ephemeral_value.py

@@ -3,6 +3,7 @@ from typing import Any, Generic, Optional, Sequence, Type
 from typing_extensions import Self
 
 from langgraph.channels.base import BaseChannel, Value
+from langgraph.constants import MISSING
 from langgraph.errors import EmptyChannelError, InvalidUpdateError
 
 
@@ -14,6 +15,7 @@ class EphemeralValue(Generic[Value], BaseChannel[Value, Value, Value]):
     def __init__(self, typ: Any, guard: bool = True) -> None:
         super().__init__(typ)
         self.guard = guard
+        self.value = MISSING
 
     def __eq__(self, value: object) -> bool:
         return isinstance(value, EphemeralValue) and value.guard == self.guard
@@ -37,10 +39,10 @@ class EphemeralValue(Generic[Value], BaseChannel[Value, Value, Value]):
 
     def update(self, values: Sequence[Value]) -> bool:
         if len(values) == 0:
-            try:
-                del self.value
+            if self.value is not MISSING:
+                self.value = MISSING
                 return True
-            except AttributeError:
+            else:
                 return False
         if len(values) != 1 and self.guard:
             raise InvalidUpdateError(
@@ -51,7 +53,9 @@ class EphemeralValue(Generic[Value], BaseChannel[Value, Value, Value]):
         return True
 
     def get(self) -> Value:
-        try:
-            return self.value
-        except AttributeError:
+        if self.value is MISSING:
             raise EmptyChannelError()
+        return self.value
+
+    def is_available(self) -> bool:
+        return self.value is not MISSING

libs/langgraph/langgraph/channels/last_value.py

@@ -1,8 +1,9 @@
-from typing import Generic, Optional, Sequence, Type
+from typing import Any, Generic, Optional, Sequence, Type
 
 from typing_extensions import Self
 
 from langgraph.channels.base import BaseChannel, Value
+from langgraph.constants import MISSING
 from langgraph.errors import (
     EmptyChannelError,
     ErrorCode,
@@ -16,6 +17,10 @@ class LastValue(Generic[Value], BaseChannel[Value, Value, Value]):
 
     __slots__ = ("value",)
 
+    def __init__(self, typ: Any, key: str = "") -> None:
+        super().__init__(typ, key)
+        self.value = MISSING
+
     def __eq__(self, value: object) -> bool:
         return isinstance(value, LastValue)
 
@@ -50,7 +55,9 @@ class LastValue(Generic[Value], BaseChannel[Value, Value, Value]):
         return True
 
     def get(self) -> Value:
-        try:
-            return self.value
-        except AttributeError:
+        if self.value is MISSING:
             raise EmptyChannelError()
+        return self.value
+
+    def is_available(self) -> bool:
+        return self.value is not MISSING

libs/langgraph/langgraph/channels/named_barrier_value.py

@@ -60,6 +60,9 @@ class NamedBarrierValue(Generic[Value], BaseChannel[Value, Value, set[Value]]):
             raise EmptyChannelError()
         return None
 
+    def is_available(self) -> bool:
+        return self.seen == self.names
+
     def consume(self) -> bool:
         if self.seen == self.names:
             self.seen = set()

libs/langgraph/langgraph/channels/topic.py

@@ -75,3 +75,6 @@ class Topic(
             return list(self.values)
         else:
             raise EmptyChannelError
+
+    def is_available(self) -> bool:
+        return bool(self.values)

libs/langgraph/langgraph/channels/untracked_value.py

@@ -3,6 +3,7 @@ from typing import Generic, Optional, Sequence, Type
 from typing_extensions import Self
 
 from langgraph.channels.base import BaseChannel, Value
+from langgraph.constants import MISSING
 from langgraph.errors import EmptyChannelError, InvalidUpdateError
 
 
@@ -14,6 +15,7 @@ class UntrackedValue(Generic[Value], BaseChannel[Value, Value, Value]):
     def __init__(self, typ: Type[Value], guard: bool = True) -> None:
         super().__init__(typ)
         self.guard = guard
+        self.value = MISSING
 
     def __eq__(self, value: object) -> bool:
         return isinstance(value, UntrackedValue) and value.guard == self.guard
@@ -48,7 +50,9 @@ class UntrackedValue(Generic[Value], BaseChannel[Value, Value, Value]):
         return True
 
     def get(self) -> Value:
-        try:
-            return self.value
-        except AttributeError:
+        if self.value is MISSING:
             raise EmptyChannelError()
+        return self.value
+
+    def is_available(self) -> bool:
+        return self.value is not MISSING

libs/langgraph/langgraph/graph/graph.py

@@ -1,3 +1,4 @@
+import asyncio
 import logging
 from collections import defaultdict
 from typing import (
@@ -31,6 +32,7 @@ from langgraph.constants import (
 )
 from langgraph.graph.branch import Branch
 from langgraph.pregel import Channel, Pregel
+from langgraph.pregel.protocol import PregelProtocol
 from langgraph.pregel.read import PregelNode
 from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry
 from langgraph.types import All, Checkpointer
@@ -418,7 +420,38 @@ class CompiledGraph(Pregel):
         *,
         xray: Union[int, bool] = False,
     ) -> DrawableGraph:
-        return self.get_graph(config, xray=xray)
+        """Returns a drawable representation of the computation graph."""
+        from langgraph.pregel.remote import RemoteGraph
+
+        # gather subgraphs
+        if xray:
+            subpregels: dict[str, PregelProtocol] = {
+                k: v
+                async for k, v in self.aget_subgraphs()
+                if isinstance(v, (CompiledGraph, RemoteGraph))
+            }
+            subgraphs = {
+                k: v
+                for k, v in zip(
+                    subpregels,
+                    await asyncio.gather(
+                        *(
+                            p.aget_graph(
+                                config,
+                                xray=xray
+                                if isinstance(xray, bool) or xray <= 0
+                                else xray - 1,
+                            )
+                            for p in subpregels.values()
+                        )
+                    ),
+                )
+            }
+        else:
+            subgraphs = {}
+
+        # draw the graph
+        return self._draw_graph(config, subgraphs=subgraphs)
 
     def get_graph(
         self,
@@ -427,17 +460,36 @@ class CompiledGraph(Pregel):
         xray: Union[int, bool] = False,
     ) -> DrawableGraph:
         """Returns a drawable representation of the computation graph."""
+        from langgraph.pregel.remote import RemoteGraph
+
+        # gather subgraphs
+        if xray:
+            subgraphs = {
+                k: v.get_graph(
+                    config,
+                    xray=xray if isinstance(xray, bool) or xray <= 0 else xray - 1,
+                )
+                for k, v in self.get_subgraphs()
+                if isinstance(v, (CompiledGraph, RemoteGraph))
+            }
+        else:
+            subgraphs = {}
+
+        # draw the graph
+        return self._draw_graph(config, subgraphs=subgraphs)
+
+    def _draw_graph(
+        self,
+        config: Optional[RunnableConfig] = None,
+        *,
+        subgraphs: dict[str, DrawableGraph] = {},
+    ) -> DrawableGraph:
+        # create the graph
         graph = DrawableGraph()
         start_nodes: dict[str, DrawableNode] = {
             START: graph.add_node(self.get_input_schema(config), START)
         }
         end_nodes: dict[str, DrawableNode] = {}
-        if xray:
-            subgraphs = {
-                k: v for k, v in self.get_subgraphs() if isinstance(v, CompiledGraph)
-            }
-        else:
-            subgraphs = {}
 
         def add_edge(
             start: str,
@@ -463,13 +515,8 @@ class CompiledGraph(Pregel):
                 metadata["__interrupt"] = "before"
             elif key in self.interrupt_after_nodes:
                 metadata["__interrupt"] = "after"
-            if xray and key in subgraphs:
-                subgraph = subgraphs[key].get_graph(
-                    config=config,
-                    xray=xray - 1
-                    if isinstance(xray, int) and not isinstance(xray, bool) and xray > 0
-                    else xray,
-                )
+            if key in subgraphs:
+                subgraph = subgraphs[key]
                 subgraph.trim_first_node()
                 subgraph.trim_last_node()
                 if len(subgraph.nodes) >= 1:

libs/langgraph/langgraph/graph/schema_utils.py

@@ -9,6 +9,7 @@ from typing import (
     Union,
     get_args,
     get_origin,
+    get_type_hints,
 )
 
 from pydantic import BaseModel
@@ -18,34 +19,56 @@ from typing_extensions import Annotated
 logger = logging.getLogger(__name__)
 
 
-class SchemaCoercionMapper:
-    _cache: weakref.WeakKeyDictionary[Type[Any], dict[int, "SchemaCoercionMapper"]] = (
-        weakref.WeakKeyDictionary()
-    )
+_cache: weakref.WeakKeyDictionary[Type[Any], dict[int, "SchemaCoercionMapper"]] = (
+    weakref.WeakKeyDictionary()
+)
 
-    def __new__(cls, schema: Type[Any], max_depth: int = 5) -> "SchemaCoercionMapper":
-        if schema not in cls._cache:
-            cls._cache[schema] = {}
-        if max_depth in cls._cache[schema]:
-            return cls._cache[schema][max_depth]
+
+class SchemaCoercionMapper:
+    def __new__(
+        cls,
+        schema: Type[Any],
+        type_hints: Optional[dict[str, Any]] = None,
+        max_depth: int = 12,
+    ) -> "SchemaCoercionMapper":
+        if schema not in _cache:
+            _cache[schema] = {}
+        if max_depth in _cache[schema]:
+            return _cache[schema][max_depth]
 
         inst = super().__new__(cls)
-        cls._cache[schema][max_depth] = inst
+        _cache[schema][max_depth] = inst
         return inst
 
-    def __init__(self, schema: Type[Any], max_depth: int = 5):
+    def __init__(
+        self,
+        schema: Type[Any],
+        type_hints: Optional[dict[str, Any]] = None,
+        max_depth: int = 12,
+    ):
         if hasattr(self, "_inited"):
             return
         self._inited = True
         self.schema = schema
+        self.type_hints = (
+            type_hints
+            if type_hints is not None
+            else get_type_hints(schema, localns={schema.__name__: schema})
+        )
         self.max_depth = max_depth
-        if hasattr(schema, "model_fields") and hasattr(schema, "model_construct"):
-            self._fields = {n: f.annotation for n, f in schema.model_fields.items()}
-            self._construct = schema.model_construct
-        elif hasattr(schema, "__fields__") and callable(
-            getattr(schema, "construct", None)
-        ):
-            self._fields = {n: f.annotation for n, f in schema.__fields__.items()}
+
+        if issubclass(schema, BaseModel):
+            self._fields = {
+                n: self.type_hints.get(n, f.annotation)
+                for n, f in schema.model_fields.items()
+            }
+            self._construct: Callable[..., Any] = schema.model_construct
+
+        elif issubclass(schema, BaseModelV1):
+            self._fields = {
+                n: self.type_hints.get(n, f.annotation)
+                for n, f in schema.__fields__.items()
+            }
             self._construct = schema.construct
         else:
             raise TypeError("Schema is neither valid Pydantic v1 nor v2 model.")
@@ -62,18 +85,23 @@ class SchemaCoercionMapper:
         processed = {}
         if self._field_coercers is None:
             self._field_coercers = {
-                n: self._build_coercer(t) for n, t in self._fields.items()
+                n: self._build_coercer(t, depth - 1) for n, t in self._fields.items()
             }
         for k, v in input_data.items():
             fn = self._field_coercers.get(k)
             processed[k] = fn(v, depth - 1) if fn else v
         return self._construct(**processed)
 
-    def _build_coercer(self, field_type: Any) -> Callable[[Any, Any], Any]:
+    def _build_coercer(
+        self, field_type: Any, depth: int, throw: bool = False
+    ) -> Callable[[Any, Any], Any]:
+        if depth == 0:
+            return self._passthrough
         origin = get_origin(field_type)
+
         if origin is Annotated:
             real_type, *_ = get_args(field_type)
-            sub = self._build_coercer(real_type)
+            sub = self._build_coercer(real_type, depth - 1)
             return lambda v, d: sub(v, d)
         if isclass(field_type):
             is_class_ = True
@@ -84,39 +112,54 @@ class SchemaCoercionMapper:
                 is_base_model = False
 
             if is_base_model:
-                mapper = SchemaCoercionMapper(field_type, self.max_depth)
+                mapper = SchemaCoercionMapper(field_type, max_depth=depth - 1)
                 return lambda v, d: mapper.coerce(v, d) if isinstance(v, dict) else v
             if is_class_ and issubclass(field_type, BaseModelV1):
-                mapper = SchemaCoercionMapper(field_type, self.max_depth)
+                mapper = SchemaCoercionMapper(field_type, max_depth=depth - 1)
                 return lambda v, d: mapper.coerce(v, d) if isinstance(v, dict) else v
         if origin is list or field_type is list:
             args = get_args(field_type)
             if len(args) != 1:
                 return lambda v, d: v
-            sub = self._build_coercer(args[0])
+            sub = self._build_coercer(args[0], depth - 1)
 
             def list_coercer(v: Any, d: Any) -> Any:
                 if not isinstance(v, (list, tuple)):
-                    raise TypeError(f"Expected list, got {type(v).__name__}")
+                    return v
                 return [sub(x, d - 1) for x in v]
 
             return list_coercer
+        if origin is set or field_type is set:
+            args = get_args(field_type)
+            if len(args) != 1:
+                return lambda v, d: v
+            sub = self._build_coercer(args[0], depth - 1)
+
+            def set_coercer(v: Any, d: Any) -> Any:
+                if not isinstance(v, (list, tuple, set)):
+                    return v
+                return {sub(x, d - 1) for x in v}
+
+            return set_coercer
         if origin is dict or field_type is dict:
             args = get_args(field_type)
             if len(args) != 2:
 
-                def plain_dict_coercer(v: Any, d: Any) -> Any:
+                def dict_coercer(v: Any, d: Any) -> Any:
                     if not isinstance(v, dict):
-                        raise TypeError(f"Expected dict, got {type(v).__name__}")
+                        if throw:
+                            raise TypeError("Expected dict, got %s" % type(v))
                     return v
 
-                return plain_dict_coercer
-            k_sub = self._build_coercer(args[0])
-            v_sub = self._build_coercer(args[1])
+                return dict_coercer
+            k_sub = self._build_coercer(args[0], depth - 1)
+            v_sub = self._build_coercer(args[1], depth - 1)
 
             def dict_coercer(v: Any, d: Any) -> Any:
                 if not isinstance(v, dict):
-                    raise TypeError(f"Expected dict, got {type(v).__name__}")
+                    if throw:
+                        raise TypeError("Expected dict, got %s" % type(v))
+                    return v
                 return {k_sub(k, d - 1): v_sub(val, d - 1) for k, val in v.items()}
 
             return dict_coercer
@@ -125,11 +168,11 @@ class SchemaCoercionMapper:
             targs = get_args(field_type)
             if not targs:
                 return lambda v, d: v
-            subs = [self._build_coercer(a) for a in targs]
+            subs = [self._build_coercer(a, depth - 1) for a in targs]
 
             def tuple_coercer(v: Any, d: Any) -> Any:
                 if not isinstance(v, (list, tuple)):
-                    raise TypeError(f"Expected tuple-like, got {type(v).__name__}")
+                    return v
                 out = []
                 for i, sp in enumerate(subs):
                     out.append(sp(v[i] if i < len(v) else None, d - 1))
@@ -139,11 +182,13 @@ class SchemaCoercionMapper:
         if origin is Union:
             uargs = get_args(field_type)
             subs, none_in_union = [], False
-            for arg in uargs:
+            for ix, arg in enumerate(uargs):
                 if arg is type(None):
                     none_in_union = True
                 else:
-                    subs.append(self._build_coercer(arg))
+                    subs.append(
+                        self._build_coercer(arg, depth - 1, throw=ix < len(uargs) - 1)
+                    )
 
             def union_coercer(v: Any, d: Any) -> Any:
                 if v is None and none_in_union:
@@ -152,11 +197,14 @@ class SchemaCoercionMapper:
                 for sp in subs:
                     try:
                         return sp(v, d - 1)
-                    except Exception as e:
+                    except TypeError as e:
                         err = e
                 if err:
                     raise err
                 return v
 
             return union_coercer
-        return lambda v, d: v
+        return self._passthrough
+
+    def _passthrough(self, v: Any, d: Any) -> Any:
+        return v

libs/langgraph/langgraph/graph/state.py

@@ -185,6 +185,7 @@ class StateGraph(Graph):
         self.schemas = {}
         self.channels = {}
         self.managed = {}
+        self.type_hints: dict[Type[Any], dict[str, Any]] = {}
         self.schema = state_schema
         self.input = input
         self.output = output
@@ -203,7 +204,7 @@ class StateGraph(Graph):
     def _add_schema(self, schema: Type[Any], /, allow_managed: bool = True) -> None:
         if schema not in self.schemas:
             _warn_invalid_state_schema(schema)
-            channels, managed = _get_channels(schema)
+            channels, managed, type_hints = _get_channels(schema)
             if managed and not allow_managed:
                 names = ", ".join(managed)
                 schema_name = getattr(schema, "__name__", "")
@@ -212,6 +213,7 @@ class StateGraph(Graph):
                     " Managed channels are not permitted in Input/Output schema."
                 )
             self.schemas[schema] = {**channels, **managed}
+            self.type_hints[schema] = type_hints
             for key, channel in channels.items():
                 if key in self.channels:
                     if self.channels[key] != channel:
@@ -416,7 +418,7 @@ class StateGraph(Graph):
                         and (vals := get_args(rargs[0]))
                     ):
                         ends = vals
-        except (TypeError, StopIteration):
+        except (NameError, TypeError, StopIteration):
             pass
 
         if destinations is not None:
@@ -821,13 +823,19 @@ class CompiledStateGraph(CompiledGraph):
             input_values = {k: k for k in self.builder.schemas[input_schema]}
             is_single_input = len(input_values) == 1 and "__root__" in input_values
 
+            branch_channel = f"branch:to:{key}"
             self.channels[key] = EphemeralValue(Any, guard=False)
+            self.channels[branch_channel] = EphemeralValue(Any, guard=False)
             self.nodes[key] = PregelNode(
-                triggers=[],
+                triggers=[branch_channel],
                 # read state keys and managed values
                 channels=(list(input_values) if is_single_input else input_values),
                 # coerce state dict to schema class (eg. pydantic model)
-                mapper=_pick_mapper(list(input_values), input_schema),
+                mapper=_pick_mapper(
+                    list(input_values),
+                    input_schema,
+                    self.builder.type_hints[input_schema],
+                ),
                 writers=[
                     # publish to this channel and state keys
                     ChannelWrite(
@@ -878,7 +886,7 @@ class CompiledStateGraph(CompiledGraph):
             if filtered := [p for p in packets if p != END]:
                 writes = [
                     (
-                        ChannelWriteEntry(f"branch:{start}:{name}:{p}", start)
+                        ChannelWriteEntry(f"branch:to:{p}", start)
                         if not isinstance(p, Send)
                         else p
                     )
@@ -914,11 +922,6 @@ class CompiledStateGraph(CompiledGraph):
             if branch.ends
             else [node for node in self.builder.nodes if node != branch.then]
         )
-        for end in ends:
-            if end != END:
-                channel_name = f"branch:{start}:{name}:{end}"
-                self.channels[channel_name] = EphemeralValue(Any, guard=False)
-                self.nodes[end].triggers.append(channel_name)
 
         # attach then subscriber
         if branch.then and branch.then != END:
@@ -942,12 +945,12 @@ def _get_state_reader(
         select=select[0] if select == ["__root__"] else select,
         fresh=True,
         # coerce state dict to schema class (eg. pydantic model)
-        mapper=_pick_mapper(state_keys, schema),
+        mapper=_pick_mapper(state_keys, schema, builder.type_hints[schema]),
     )
 
 
 def _pick_mapper(
-    state_keys: Sequence[str], schema: Type[Any]
+    state_keys: Sequence[str], schema: Type[Any], type_hints: Optional[dict[str, Any]]
 ) -> Optional[Callable[[Any], Any]]:
     if state_keys == ["__root__"]:
         return None
@@ -955,7 +958,7 @@ def _pick_mapper(
         if issubclass(schema, dict):
             return None
         if issubclass(schema, (BaseModel, BaseModelV1)):
-            return SchemaCoercionMapper(schema)
+            return SchemaCoercionMapper(schema, type_hints)
     return partial(_coerce_state, schema)
 
 
@@ -1017,18 +1020,24 @@ CONTROL_BRANCH = Branch(CONTROL_BRANCH_PATH, None)
 
 def _get_channels(
     schema: Type[dict],
-) -> tuple[dict[str, BaseChannel], dict[str, ManagedValueSpec]]:
+) -> tuple[dict[str, BaseChannel], dict[str, ManagedValueSpec], dict[str, Any]]:
     if not hasattr(schema, "__annotations__"):
-        return {"__root__": _get_channel("__root__", schema, allow_managed=False)}, {}
+        return (
+            {"__root__": _get_channel("__root__", schema, allow_managed=False)},
+            {},
+            {},
+        )
 
+    type_hints = get_type_hints(schema, include_extras=True)
     all_keys = {
         name: _get_channel(name, typ)
-        for name, typ in get_type_hints(schema, include_extras=True).items()
+        for name, typ in type_hints.items()
         if name != "__slots__"
     }
     return (
         {k: v for k, v in all_keys.items() if isinstance(v, BaseChannel)},
         {k: v for k, v in all_keys.items() if is_managed_value(v)},
+        type_hints,
     )
 
 

libs/langgraph/langgraph/pregel/algo.py

@@ -1,4 +1,4 @@
-import functools
+import binascii
 import itertools
 import sys
 from collections import defaultdict, deque
@@ -19,7 +19,6 @@ from typing import (
     cast,
     overload,
 )
-from uuid import UUID
 
 from langchain_core.callbacks import Callbacks
 from langchain_core.callbacks.manager import AsyncParentRunManager, ParentRunManager
@@ -28,6 +27,7 @@ from langchain_core.runnables.config import RunnableConfig
 from langgraph.channels.base import BaseChannel
 from langgraph.checkpoint.base import (
     BaseCheckpointSaver,
+    ChannelVersions,
     Checkpoint,
     PendingWrite,
     V,
@@ -233,10 +233,21 @@ def apply_writes(
     channels: Mapping[str, BaseChannel],
     tasks: Iterable[WritesProtocol],
     get_next_version: Optional[GetNextVersion],
-) -> dict[str, list[Any]]:
+) -> tuple[dict[str, list[Any]], set[str]]:
     """Apply writes from a set of tasks (usually the tasks from a Pregel step)
     to the checkpoint and channels, and return managed values writes to be applied
-    externally."""
+    externally.
+
+    Args:
+        checkpoint: The checkpoint to update.
+        channels: The channels to update.
+        tasks: The tasks to apply writes from.
+        get_next_version: Optional function to determine the next version of a channel.
+
+    Returns:
+        A tuple containing the managed values writes to be applied externally, and
+        the set of channels that were updated in this step.
+    """
     # sort tasks on path, to ensure deterministic order for update application
     # any path parts after the 3rd are ignored for sorting
     # (we use them for eg. task ids which aren't good for sorting)
@@ -318,9 +329,8 @@ def apply_writes(
                         max_version,
                         channels[chan],
                     )
-
     # Return managed values writes to be applied externally
-    return pending_writes_by_managed
+    return pending_writes_by_managed, updated_channels
 
 
 @overload
@@ -337,6 +347,8 @@ def prepare_next_tasks(
     store: Literal[None] = None,
     checkpointer: Literal[None] = None,
     manager: Literal[None] = None,
+    trigger_to_nodes: Optional[Mapping[str, list[str]]] = None,
+    updated_channels: Optional[set[str]] = None,
 ) -> dict[str, PregelTask]: ...
 
 
@@ -354,6 +366,8 @@ def prepare_next_tasks(
     store: Optional[BaseStore],
     checkpointer: Optional[BaseCheckpointSaver],
     manager: Union[None, ParentRunManager, AsyncParentRunManager],
+    trigger_to_nodes: Optional[Mapping[str, list[str]]] = None,
+    updated_channels: Optional[set[str]] = None,
 ) -> dict[str, PregelExecutableTask]: ...
 
 
@@ -370,10 +384,37 @@ def prepare_next_tasks(
     store: Optional[BaseStore] = None,
     checkpointer: Optional[BaseCheckpointSaver] = None,
     manager: Union[None, ParentRunManager, AsyncParentRunManager] = None,
+    trigger_to_nodes: Optional[Mapping[str, list[str]]] = None,
+    updated_channels: Optional[set[str]] = None,
 ) -> Union[dict[str, PregelTask], dict[str, PregelExecutableTask]]:
     """Prepare the set of tasks that will make up the next Pregel step.
-    This is the union of all PUSH tasks (Sends) and PULL tasks (nodes triggered
-    by edges)."""
+
+    Args:
+        checkpoint: The current checkpoint.
+        pending_writes: The list of pending writes.
+        processes: The mapping of process names to PregelNode instances.
+        channels: The mapping of channel names to BaseChannel instances.
+        managed: The mapping of managed value names to functions.
+        config: The runnable configuration.
+        step: The current step.
+        for_execution: Whether the tasks are being prepared for execution.
+        store: An instance of BaseStore to make it available for usage within tasks.
+        checkpointer: Checkpointer instance used for saving checkpoints.
+        manager: The parent run manager to use for the tasks.
+        trigger_to_nodes: Optional: Mapping of channel names to the set of nodes
+            that are can be triggered by that channel.
+        updated_channels: Optional. Set of channel names that have been updated during
+            the previous step. Using in conjunction with trigger_to_nodes to speed
+            up the process of determining which nodes should be triggered in the next
+            step.
+
+    Returns:
+        A dictionary of tasks to be executed. The keys are the task ids and the values
+        are the tasks themselves. This is the union of all PUSH tasks (Sends)
+        and PULL tasks (nodes triggered by edges).
+    """
+    checkpoint_id_bytes = binascii.unhexlify(checkpoint["id"].replace("-", ""))
+    null_version = checkpoint_null_version(checkpoint)
     tasks: list[Union[PregelTask, PregelExecutableTask]] = []
     # Consume pending_sends from previous step
     for idx, _ in enumerate(checkpoint["pending_sends"]):
@@ -381,6 +422,8 @@ def prepare_next_tasks(
             (PUSH, idx),
             None,
             checkpoint=checkpoint,
+            checkpoint_id_bytes=checkpoint_id_bytes,
+            checkpoint_null_version=null_version,
             pending_writes=pending_writes,
             processes=processes,
             channels=channels,
@@ -393,13 +436,34 @@ def prepare_next_tasks(
             manager=manager,
         ):
             tasks.append(task)
+
+    # This section is an optimization that allows which nodes will be active
+    # during the next step.
+    # When there's information about:
+    # 1. Which channels were updated in the previous step
+    # 2. Which nodes are triggered by which channels
+    # Then we can determine which nodes should be triggered in the next step
+    # without having to cycle through all nodes.
+    if updated_channels and trigger_to_nodes:
+        triggered_nodes: set[str] = set()
+        # Get all nodes that have triggers associated with an updated channel
+        for channel in updated_channels:
+            if node_ids := trigger_to_nodes.get(channel):
+                triggered_nodes.update(node_ids)
+        # Sort the nodes to ensure deterministic order
+        candidate_nodes: Iterable[str] = sorted(triggered_nodes)
+    else:
+        candidate_nodes = processes.keys()
+
     # Check if any processes should be run in next step
     # If so, prepare the values to be passed to them
-    for name in processes:
+    for name in candidate_nodes:
         if task := prepare_single_task(
             (PULL, name),
             None,
             checkpoint=checkpoint,
+            checkpoint_id_bytes=checkpoint_id_bytes,
+            checkpoint_null_version=null_version,
             pending_writes=pending_writes,
             processes=processes,
             channels=channels,
@@ -415,11 +479,16 @@ def prepare_next_tasks(
     return {t.id: t for t in tasks}
 
 
+PUSH_TRIGGER = (PUSH,)
+
+
 def prepare_single_task(
     task_path: tuple[Any, ...],
     task_id_checksum: Optional[str],
     *,
     checkpoint: Checkpoint,
+    checkpoint_id_bytes: bytes,
+    checkpoint_null_version: Optional[V],
     pending_writes: list[PendingWrite],
     processes: Mapping[str, PregelNode],
     channels: Mapping[str, BaseChannel],
@@ -433,7 +502,6 @@ def prepare_single_task(
 ) -> Union[None, PregelTask, PregelExecutableTask]:
     """Prepares a single task for the next Pregel step, given a task path, which
     uniquely identifies a PUSH or PULL task within the graph."""
-    checkpoint_id = UUID(checkpoint["id"]).bytes
     configurable = config.get(CONF, {})
     parent_ns = configurable.get(CONFIG_KEY_CHECKPOINT_NS, "")
 
@@ -446,10 +514,10 @@ def prepare_single_task(
         if name is None:
             raise ValueError("`call` functions must have a `__name__` attribute")
         # create task id
-        triggers = [PUSH]
+        triggers: Sequence[str] = PUSH_TRIGGER
         checkpoint_ns = f"{parent_ns}{NS_SEP}{name}" if parent_ns else name
         task_id = _uuid5_str(
-            checkpoint_id,
+            checkpoint_id_bytes,
             checkpoint_ns,
             str(step),
             name,
@@ -507,6 +575,7 @@ def prepare_single_task(
                         CONFIG_KEY_CHECKPOINT_ID: None,
                         CONFIG_KEY_CHECKPOINT_NS: task_checkpoint_ns,
                         CONFIG_KEY_SCRATCHPAD: _scratchpad(
+                            config[CONF].get(CONFIG_KEY_SCRATCHPAD),
                             pending_writes,
                             task_id,
                         ),
@@ -539,12 +608,12 @@ def prepare_single_task(
                 )
                 return
             # create task id
-            triggers = [PUSH]
+            triggers = PUSH_TRIGGER
             checkpoint_ns = (
                 f"{parent_ns}{NS_SEP}{packet.node}" if parent_ns else packet.node
             )
             task_id = _uuid5_str(
-                checkpoint_id,
+                checkpoint_id_bytes,
                 checkpoint_ns,
                 str(step),
                 packet.node,
@@ -616,6 +685,7 @@ def prepare_single_task(
                             CONFIG_KEY_CHECKPOINT_ID: None,
                             CONFIG_KEY_CHECKPOINT_NS: task_checkpoint_ns,
                             CONFIG_KEY_SCRATCHPAD: _scratchpad(
+                                config[CONF].get(CONFIG_KEY_SCRATCHPAD),
                                 pending_writes,
                                 task_id,
                             ),
@@ -640,21 +710,17 @@ def prepare_single_task(
         if name not in processes:
             return
         proc = processes[name]
-        version_type = type(next(iter(checkpoint["channel_versions"].values()), None))
-        null_version = version_type()  # type: ignore[misc]
-        if null_version is None:
+        if checkpoint_null_version is None:
             return
-        seen = checkpoint["versions_seen"].get(name, {})
         # If any of the channels read by this process were updated
-        if triggers := sorted(
-            chan
-            for chan in proc.triggers
-            if not isinstance(
-                read_channel(channels, chan, return_exception=True), EmptyChannelError
-            )
-            and checkpoint["channel_versions"].get(chan, null_version)  # type: ignore[operator]
-            > seen.get(chan, null_version)
+        if _triggers(
+            channels,
+            checkpoint["channel_versions"],
+            checkpoint["versions_seen"].get(name),
+            checkpoint_null_version,
+            proc,
         ):
+            triggers = tuple(sorted(proc.triggers))
             try:
                 val = next(
                     _proc_input(proc, managed, channels, for_execution=for_execution)
@@ -671,7 +737,7 @@ def prepare_single_task(
             # create task id
             checkpoint_ns = f"{parent_ns}{NS_SEP}{name}" if parent_ns else name
             task_id = _uuid5_str(
-                checkpoint_id,
+                checkpoint_id_bytes,
                 checkpoint_ns,
                 str(step),
                 name,
@@ -714,7 +780,7 @@ def prepare_single_task(
                                 CONFIG_KEY_SEND: partial(
                                     local_write,
                                     writes.extend,
-                                    processes.keys(),
+                                    tuple(processes.keys()),
                                 ),
                                 CONFIG_KEY_READ: partial(
                                     local_read,
@@ -723,7 +789,10 @@ def prepare_single_task(
                                     channels,
                                     managed,
                                     PregelTaskWrites(
-                                        task_path[:3], name, writes, triggers
+                                        task_path[:3],
+                                        name,
+                                        writes,
+                                        triggers,
                                     ),
                                     config,
                                 ),
@@ -741,6 +810,7 @@ def prepare_single_task(
                                 CONFIG_KEY_CHECKPOINT_ID: None,
                                 CONFIG_KEY_CHECKPOINT_NS: task_checkpoint_ns,
                                 CONFIG_KEY_SCRATCHPAD: _scratchpad(
+                                    config[CONF].get(CONFIG_KEY_SCRATCHPAD),
                                     pending_writes,
                                     task_id,
                                 ),
@@ -761,13 +831,59 @@ def prepare_single_task(
                 return PregelTask(task_id, name, task_path[:3])
 
 
+def checkpoint_null_version(
+    checkpoint: Checkpoint,
+) -> Optional[V]:
+    """Get the null version for the checkpoint, if available."""
+    for version in checkpoint["channel_versions"].values():
+        return type(version)()
+    return None
+
+
+def _triggers(
+    channels: Mapping[str, BaseChannel],
+    versions: ChannelVersions,
+    seen: Optional[ChannelVersions],
+    null_version: V,
+    proc: PregelNode,
+) -> Sequence[str]:
+    if seen is None:
+        for chan in proc.triggers:
+            if channels[chan].is_available():
+                return (chan,)
+    else:
+        for chan in proc.triggers:
+            if channels[chan].is_available() and versions.get(  # type: ignore[operator]
+                chan, null_version
+            ) > seen.get(chan, null_version):
+                return (chan,)
+    return EMPTY_SEQ
+
+
 def _scratchpad(
+    parent_scratchpad: Optional[PregelScratchpad],
     pending_writes: list[PendingWrite],
     task_id: str,
 ) -> PregelScratchpad:
+    # None cannot be used as a resume value, because it would be difficult to
+    # distinguish from missing when used over http
     null_resume_write = next(
         (w for w in pending_writes if w[0] == NULL_TASK_ID and w[1] == RESUME), None
     )
+
+    def get_null_resume(consume: bool = False) -> Any:
+        if null_resume_write is None:
+            if parent_scratchpad is not None:
+                return parent_scratchpad.get_null_resume(consume)
+            return None
+        if consume:
+            try:
+                pending_writes.remove(null_resume_write)
+                return null_resume_write[2]
+            except ValueError:
+                return None
+        return null_resume_write[2]
+
     # using itertools.count as an atomic counter (+= 1 is not thread-safe)
     return PregelScratchpad(
         # call
@@ -777,10 +893,7 @@ def _scratchpad(
         resume=next(
             (w[2] for w in pending_writes if w[0] == task_id and w[1] == RESUME), []
         ),
-        null_resume=null_resume_write[2] if null_resume_write is not None else None,
-        _consume_null_resume=functools.partial(pending_writes.remove, null_resume_write)
-        if null_resume_write is not None
-        else lambda: None,
+        get_null_resume=get_null_resume,
         # subgraph
         subgraph_counter=itertools.count(0).__next__,
     )

libs/langgraph/langgraph/pregel/debug.py

@@ -137,7 +137,12 @@ def map_debug_task_results(
             "result": [
                 w for w in writes if w[0] in stream_channels_list or w[0] == RETURN
             ],
-            "interrupts": [asdict(w[1]) for w in writes if w[0] == INTERRUPT],
+            "interrupts": [
+                asdict(v)
+                for w in writes
+                if w[0] == INTERRUPT
+                for v in (w[1] if isinstance(w[1], Sequence) else [w[1]])
+            ],
         },
     }
 
@@ -293,8 +298,9 @@ def tasks_w_writes(
                 ),
                 tuple(
                     v
-                    for tid, n, v in pending_writes
+                    for tid, n, vv in pending_writes
                     if tid == task.id and n == INTERRUPT
+                    for v in (vv if isinstance(vv, Sequence) else [vv])
                 ),
                 states.get(task.id) if states else None,
                 (

libs/langgraph/langgraph/pregel/io.py

@@ -14,7 +14,6 @@ from langgraph.constants import (
     NULL_TASK_ID,
     RESUME,
     RETURN,
-    SELF,
     START,
     TAG_HIDDEN,
     TASKS,
@@ -28,7 +27,7 @@ def is_task_id(task_id: str) -> bool:
     """Check if a string is a valid task id."""
     try:
         UUID(task_id)
-    except ValueError:
+    except Exception:
         return False
     return True
 
@@ -38,14 +37,11 @@ def read_channel(
     chan: str,
     *,
     catch: bool = True,
-    return_exception: bool = False,
 ) -> Any:
     try:
         return channels[chan].get()
-    except EmptyChannelError as exc:
-        if return_exception:
-            return exc
-        elif catch:
+    except EmptyChannelError:
+        if catch:
             return None
         else:
             raise
@@ -84,7 +80,7 @@ def map_command(
             if isinstance(send, Send):
                 yield (NULL_TASK_ID, TASKS, send)
             elif isinstance(send, str):
-                yield (NULL_TASK_ID, f"branch:{START}:{SELF}:{send}", START)
+                yield (NULL_TASK_ID, f"branch:to:{send}", START)
             else:
                 raise TypeError(
                     f"In Command.goto, expected Send/str, got {type(send).__name__}"

libs/langgraph/langgraph/pregel/loop.py

@@ -1,5 +1,7 @@
 import asyncio
+import binascii
 import concurrent.futures
+import dataclasses
 from collections import defaultdict, deque
 from contextlib import AsyncExitStack, ExitStack
 from inspect import signature
@@ -79,6 +81,7 @@ from langgraph.pregel.algo import (
     GetNextVersion,
     PregelTaskWrites,
     apply_writes,
+    checkpoint_null_version,
     increment,
     prepare_next_tasks,
     prepare_single_task,
@@ -207,6 +210,7 @@ class PregelLoop(LoopProtocol):
         manager: Union[None, AsyncParentRunManager, ParentRunManager] = None,
         input_model: Optional[Type[BaseModel]] = None,
         debug: bool = False,
+        trigger_to_nodes: Optional[Mapping[str, list[str]]] = None,
     ) -> None:
         super().__init__(
             step=0,
@@ -230,6 +234,7 @@ class PregelLoop(LoopProtocol):
             CONFIG_KEY_CHECKPOINT_ID not in config[CONF]
             or CONFIG_KEY_DEDUPE_TASKS in config[CONF]
         )
+        self.trigger_to_nodes = trigger_to_nodes
         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])
@@ -347,12 +352,16 @@ class PregelLoop(LoopProtocol):
         ):
             self.to_interrupt.append(task)
             return
+        checkpoint_id_bytes = binascii.unhexlify(self.checkpoint["id"].replace("-", ""))
+        null_version = checkpoint_null_version(self.checkpoint)
         if pushed := cast(
             Optional[PregelExecutableTask],
             prepare_single_task(
                 (PUSH, task.path, write_idx, task.id, call),
                 None,
                 checkpoint=self.checkpoint,
+                checkpoint_id_bytes=checkpoint_id_bytes,
+                checkpoint_null_version=null_version,
                 pending_writes=self.checkpoint_pending_writes,
                 processes=self.nodes,
                 channels=self.channels,
@@ -400,6 +409,8 @@ class PregelLoop(LoopProtocol):
         if self.status != "pending":
             raise RuntimeError("Cannot tick when status is no longer 'pending'")
 
+        updated_channels: set[str] | None = None
+
         if self.input not in (INPUT_DONE, INPUT_RESUMING, INPUT_SHOULD_VALIDATE):
             self._first(input_keys=input_keys)
         elif self.to_interrupt:
@@ -421,7 +432,7 @@ class PregelLoop(LoopProtocol):
                     ),
                 )
             # all tasks have finished
-            mv_writes = apply_writes(
+            mv_writes, updated_channels = apply_writes(
                 self.checkpoint,
                 self.channels,
                 self.tasks.values(),
@@ -487,6 +498,8 @@ class PregelLoop(LoopProtocol):
             manager=self.manager,
             store=self.store,
             checkpointer=self.checkpointer,
+            trigger_to_nodes=self.trigger_to_nodes,
+            updated_channels=updated_channels,
         )
         self.to_interrupt = []
 
@@ -565,7 +578,7 @@ class PregelLoop(LoopProtocol):
                         self.checkpoint["versions_seen"].get(INTERRUPT, {}).values(),
                         default=None,
                     ):
-                        self.tasks[tid] = task._replace(scheduled=True)
+                        self.tasks[tid] = dataclasses.replace(task, scheduled=True)
                 else:
                     task.writes.append((k, v))
 
@@ -587,15 +600,6 @@ class PregelLoop(LoopProtocol):
             )
         )
 
-        # take resume value from parent
-        if scratchpad := cast(
-            Optional[PregelScratchpad], configurable.get(CONFIG_KEY_SCRATCHPAD)
-        ):
-            if (
-                isinstance(scratchpad, PregelScratchpad)
-                and scratchpad.null_resume is not None
-            ):
-                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:
@@ -615,7 +619,7 @@ class PregelLoop(LoopProtocol):
         if null_writes := [
             w[1:] for w in self.checkpoint_pending_writes if w[0] == NULL_TASK_ID
         ]:
-            mv_writes = apply_writes(
+            mv_writes, _ = apply_writes(
                 self.checkpoint,
                 self.channels,
                 [PregelTaskWrites((), INPUT, null_writes, [])],
@@ -664,7 +668,7 @@ class PregelLoop(LoopProtocol):
                 manager=None,
             )
             # apply input writes
-            mv_writes = apply_writes(
+            mv_writes, _ = apply_writes(
                 self.checkpoint,
                 self.channels,
                 [
@@ -779,7 +783,7 @@ class PregelLoop(LoopProtocol):
                 and self.checkpoint_pending_writes
                 and any(task.writes for task in self.tasks.values())
             ):
-                mv_writes = apply_writes(
+                mv_writes, _ = apply_writes(
                     self.checkpoint,
                     self.channels,
                     self.tasks.values(),
@@ -794,11 +798,14 @@ class PregelLoop(LoopProtocol):
                     [w for t in self.tasks.values() for w in t.writes],
                     self.channels,
                 )
-            # emit INTERRUPT event
-            self._emit(
-                "updates",
-                lambda: iter([{INTERRUPT: cast(GraphInterrupt, exc_value).args[0]}]),
-            )
+            # emit INTERRUPT if exception is empty (otherwise emitted by put_writes)
+            if exc_value is not None and (not exc_value.args or not exc_value.args[0]):
+                self._emit(
+                    "updates",
+                    lambda: iter(
+                        [{INTERRUPT: cast(GraphInterrupt, exc_value).args[0]}]
+                    ),
+                )
             # save final output
             self.output = read_channels(self.channels, self.output_keys)
             # suppress interrupt
@@ -829,7 +836,25 @@ class PregelLoop(LoopProtocol):
                 "tags", EMPTY_SEQ
             ):
                 return
-            if writes[0][0] != ERROR and writes[0][0] != INTERRUPT:
+            if writes[0][0] == INTERRUPT:
+                self._emit(
+                    "updates",
+                    lambda: iter(
+                        [
+                            {
+                                INTERRUPT: tuple(
+                                    v
+                                    for w in writes
+                                    if w[0] == INTERRUPT
+                                    for v in (
+                                        w[1] if isinstance(w[1], Sequence) else (w[1],)
+                                    )
+                                )
+                            }
+                        ]
+                    ),
+                )
+            elif writes[0][0] != ERROR:
                 self._emit(
                     "updates",
                     map_output_updates,
@@ -865,6 +890,7 @@ class SyncPregelLoop(PregelLoop, ContextManager):
         stream_keys: Union[str, Sequence[str]] = EMPTY_SEQ,
         input_model: Optional[Type[BaseModel]] = None,
         debug: bool = False,
+        trigger_to_nodes: Optional[Mapping[str, list[str]]] = None,
     ) -> None:
         super().__init__(
             input,
@@ -881,6 +907,7 @@ class SyncPregelLoop(PregelLoop, ContextManager):
             interrupt_before=interrupt_before,
             manager=manager,
             debug=debug,
+            trigger_to_nodes=trigger_to_nodes,
         )
         self.stack = ExitStack()
         if checkpointer:
@@ -1006,6 +1033,7 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager):
         stream_keys: Union[str, Sequence[str]] = EMPTY_SEQ,
         input_model: Optional[Type[BaseModel]] = None,
         debug: bool = False,
+        trigger_to_nodes: Optional[Mapping[str, list[str]]] = None,
     ) -> None:
         super().__init__(
             input,
@@ -1022,6 +1050,7 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager):
             interrupt_before=interrupt_before,
             manager=manager,
             debug=debug,
+            trigger_to_nodes=trigger_to_nodes,
         )
         self.stack = AsyncExitStack()
         if checkpointer:

libs/langgraph/langgraph/pregel/protocol.py

@@ -12,7 +12,7 @@ from langchain_core.runnables import Runnable, RunnableConfig
 from langchain_core.runnables.graph import Graph as DrawableGraph
 from typing_extensions import Self
 
-from langgraph.pregel.types import All, StateSnapshot, StreamMode
+from langgraph.pregel.types import All, StateSnapshot, StateUpdate, StreamMode
 
 
 class PregelProtocol(
@@ -69,6 +69,20 @@ class PregelProtocol(
         limit: Optional[int] = None,
     ) -> AsyncIterator[StateSnapshot]: ...
 
+    @abstractmethod
+    def bulk_update_state(
+        self,
+        config: RunnableConfig,
+        updates: Sequence[Sequence[StateUpdate]],
+    ) -> RunnableConfig: ...
+
+    @abstractmethod
+    async def abulk_update_state(
+        self,
+        config: RunnableConfig,
+        updates: Sequence[Sequence[StateUpdate]],
+    ) -> RunnableConfig: ...
+
     @abstractmethod
     def update_state(
         self,

libs/langgraph/langgraph/pregel/remote.py

@@ -457,6 +457,20 @@ class RemoteGraph(PregelProtocol):
         for state in states:
             yield self._create_state_snapshot(state)
 
+    def bulk_update_state(
+        self,
+        config: RunnableConfig,
+        updates: list[tuple[Optional[dict[str, Any]], Optional[str]]],
+    ) -> RunnableConfig:
+        raise NotImplementedError
+
+    async def abulk_update_state(
+        self,
+        config: RunnableConfig,
+        updates: list[tuple[Optional[dict[str, Any]], Optional[str]]],
+    ) -> RunnableConfig:
+        raise NotImplementedError
+
     def update_state(
         self,
         config: RunnableConfig,

libs/langgraph/langgraph/pregel/runner.py

@@ -2,6 +2,7 @@ import asyncio
 import concurrent.futures
 import threading
 import time
+import weakref
 from functools import partial
 from typing import (
     Any,
@@ -25,12 +26,10 @@ from langgraph.constants import (
     CONF,
     CONFIG_KEY_CALL,
     CONFIG_KEY_SCRATCHPAD,
-    CONFIG_KEY_SEND,
     ERROR,
     INTERRUPT,
     MISSING,
     NO_WRITES,
-    PUSH,
     RESUME,
     RETURN,
     TAG_HIDDEN,
@@ -48,7 +47,9 @@ E = TypeVar("E", threading.Event, asyncio.Event)
 
 class FuturesDict(Generic[F, E], dict[F, Optional[PregelExecutableTask]]):
     event: E
-    callback: Callable[[PregelExecutableTask, Optional[BaseException]], None]
+    callback: weakref.ref[
+        Callable[[PregelExecutableTask, Optional[BaseException]], None]
+    ]
     counter: int
     done: set[F]
     lock: threading.Lock
@@ -56,7 +57,9 @@ class FuturesDict(Generic[F, E], dict[F, Optional[PregelExecutableTask]]):
     def __init__(
         self,
         event: E,
-        callback: Callable[[PregelExecutableTask, Optional[BaseException]], None],
+        callback: weakref.ref[
+            Callable[[PregelExecutableTask, Optional[BaseException]], None]
+        ],
         future_type: Type[F],
         # used for generic typing, newer py supports FutureDict[...](...)
     ) -> None:
@@ -85,7 +88,7 @@ class FuturesDict(Generic[F, E], dict[F, Optional[PregelExecutableTask]]):
         fut: F,
     ) -> None:
         try:
-            self.callback(task, _exception(fut))
+            self.callback()(task, _exception(fut))  # type: ignore[misc]
         finally:
             with self.lock:
                 self.done.add(fut)
@@ -102,10 +105,13 @@ class PregelRunner:
     def __init__(
         self,
         *,
-        submit: Submit,
-        put_writes: Callable[[str, Sequence[tuple[str, Any]]], None],
-        schedule_task: Callable[
-            [PregelExecutableTask, int, Optional[Call]], Optional[PregelExecutableTask]
+        submit: weakref.ref[Submit],
+        put_writes: weakref.ref[Callable[[str, Sequence[tuple[str, Any]]], None]],
+        schedule_task: weakref.ref[
+            Callable[
+                [PregelExecutableTask, int, Optional[Call]],
+                Optional[PregelExecutableTask],
+            ]
         ],
         use_astream: bool = False,
         node_finished: Optional[Callable[[str], None]] = None,
@@ -125,99 +131,9 @@ class PregelRunner:
         retry_policy: Optional[RetryPolicy] = None,
         get_waiter: Optional[Callable[[], concurrent.futures.Future[None]]] = None,
     ) -> Iterator[None]:
-        def writer(
-            task: PregelExecutableTask,
-            writes: Sequence[tuple[str, Any]],
-            *,
-            calls: Optional[Sequence[Call]] = None,
-        ) -> Sequence[Optional[concurrent.futures.Future]]:
-            if all(w[0] != PUSH for w in writes):
-                return task.config[CONF][CONFIG_KEY_SEND](writes)
-
-            # schedule PUSH tasks, collect futures
-            scratchpad: PregelScratchpad = task.config[CONF][CONFIG_KEY_SCRATCHPAD]
-            rtn: dict[int, Optional[concurrent.futures.Future]] = {}
-            for idx, w in enumerate(writes):
-                # bail if not a PUSH write
-                if w[0] != PUSH:
-                    continue
-                # schedule the next task, if the callback returns one
-                wcall = calls[idx] if calls else None
-                if next_task := self.schedule_task(
-                    task, scratchpad.call_counter(), wcall
-                ):
-                    if fut := next(
-                        (
-                            f
-                            for f, t in futures.items()
-                            if t is not None and t == next_task.id
-                        ),
-                        None,
-                    ):
-                        # if the parent task was retried,
-                        # the next task might already be running
-                        rtn[idx] = fut
-                    elif next_task.writes:
-                        # if it already ran, return the result
-                        fut = concurrent.futures.Future()
-                        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
-                        ):
-                            fut.set_exception(
-                                exc
-                                if isinstance(exc, BaseException)
-                                else Exception(exc)
-                            )
-                        else:
-                            fut.set_result(None)
-                        rtn[idx] = fut
-                    else:
-                        # schedule the next task
-                        fut = self.submit(
-                            run_with_retry,
-                            next_task,
-                            retry_policy,
-                            configurable={
-                                CONFIG_KEY_SEND: partial(writer, next_task),
-                                CONFIG_KEY_CALL: partial(call, next_task),
-                            },
-                            __reraise_on_exit__=reraise,
-                            # starting a new task in the next tick ensures
-                            # updates from this tick are committed/streamed first
-                            __next_tick__=True,
-                        )
-                        futures[fut] = next_task
-                        rtn[idx] = fut
-            return [rtn.get(i) for i in range(len(writes))]
-
-        def call(
-            task: PregelExecutableTask,
-            func: Callable[[Any], Union[Awaitable[Any], Any]],
-            input: Any,
-            *,
-            retry: Optional[RetryPolicy] = None,
-            callbacks: Callbacks = None,
-        ) -> concurrent.futures.Future[Any]:
-            if asyncio.iscoroutinefunction(func):
-                raise RuntimeError("In an sync context async tasks cannot be called")
-            (fut,) = writer(
-                task,
-                [(PUSH, None)],
-                calls=[Call(func, input, retry=retry, callbacks=callbacks)],
-            )
-            assert fut is not None, "writer did not return a future for call"
-            # return a chained future to ensure commit() callback is called
-            # before the returned future is resolved, to ensure stream order etc
-            return chain_future(fut, concurrent.futures.Future())
-
         tasks = tuple(tasks)
         futures = FuturesDict(
-            callback=self.commit,
+            callback=weakref.WeakMethod(self.commit),
             event=threading.Event(),
             future_type=concurrent.futures.Future,
         )
@@ -231,8 +147,15 @@ class PregelRunner:
                     t,
                     retry_policy,
                     configurable={
-                        CONFIG_KEY_SEND: partial(writer, t),
-                        CONFIG_KEY_CALL: partial(call, t),
+                        CONFIG_KEY_CALL: partial(
+                            _call,
+                            weakref.ref(t),
+                            retry=retry_policy,
+                            futures=weakref.ref(futures),
+                            schedule_task=self.schedule_task,
+                            submit=self.submit,
+                            reraise=reraise,
+                        ),
                     },
                 )
                 self.commit(t, None)
@@ -255,13 +178,20 @@ class PregelRunner:
         # schedule tasks
         for t in tasks:
             if not t.writes:
-                fut = self.submit(
+                fut = self.submit()(  # type: ignore[misc]
                     run_with_retry,
                     t,
                     retry_policy,
                     configurable={
-                        CONFIG_KEY_SEND: partial(writer, t),
-                        CONFIG_KEY_CALL: partial(call, t),
+                        CONFIG_KEY_CALL: partial(
+                            _call,
+                            weakref.ref(t),
+                            retry=retry_policy,
+                            futures=weakref.ref(futures),
+                            schedule_task=self.schedule_task,
+                            submit=self.submit,
+                            reraise=reraise,
+                        ),
                     },
                     __reraise_on_exit__=reraise,
                 )
@@ -313,125 +243,10 @@ class PregelRunner:
         retry_policy: Optional[RetryPolicy] = None,
         get_waiter: Optional[Callable[[], asyncio.Future[None]]] = None,
     ) -> AsyncIterator[None]:
-        def writer(
-            task: PregelExecutableTask,
-            writes: Sequence[tuple[str, Any]],
-            *,
-            calls: Optional[Sequence[Call]] = None,
-        ) -> Sequence[Optional[asyncio.Future]]:
-            if all(w[0] != PUSH for w in writes):
-                return task.config[CONF][CONFIG_KEY_SEND](writes)
-
-            # schedule PUSH tasks, collect futures
-            scratchpad: PregelScratchpad = task.config[CONF][CONFIG_KEY_SCRATCHPAD]
-            rtn: dict[int, Optional[asyncio.Future]] = {}
-            for idx, w in enumerate(writes):
-                # bail if not a PUSH write
-                if w[0] != PUSH:
-                    continue
-                # schedule the next task, if the callback returns one
-                wcall = calls[idx] if calls is not None else None
-                if next_task := self.schedule_task(
-                    task, scratchpad.call_counter(), wcall
-                ):
-                    # if the parent task was retried,
-                    # the next task might already be running
-                    if fut := next(
-                        (
-                            f
-                            for f, t in futures.items()
-                            if t is not None and t == next_task.id
-                        ),
-                        None,
-                    ):
-                        # if the parent task was retried,
-                        # the next task might already be running
-                        rtn[idx] = fut
-                    elif next_task.writes:
-                        # if it already ran, return the result
-                        fut = asyncio.Future(loop=loop)
-                        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
-                        ):
-                            fut.set_exception(
-                                exc
-                                if isinstance(exc, BaseException)
-                                else Exception(exc)
-                            )
-                        else:
-                            fut.set_result(None)
-                        rtn[idx] = fut
-                    else:
-                        # schedule the next task
-                        fut = cast(
-                            asyncio.Future,
-                            self.submit(
-                                arun_with_retry,
-                                next_task,
-                                retry_policy,
-                                stream=self.use_astream,
-                                configurable={
-                                    CONFIG_KEY_SEND: partial(writer, next_task),
-                                    CONFIG_KEY_CALL: partial(call, next_task),
-                                },
-                                __name__=t.name,
-                                __cancel_on_exit__=True,
-                                __reraise_on_exit__=reraise,
-                                # starting a new task in the next tick ensures
-                                # updates from this tick are committed/streamed first
-                                __next_tick__=True,
-                            ),
-                        )
-                        futures[fut] = next_task
-                        rtn[idx] = fut
-            return [rtn.get(i) for i in range(len(writes))]
-
-        def call(
-            task: PregelExecutableTask,
-            func: Callable[[Any], Union[Awaitable[Any], Any]],
-            input: Any,
-            *,
-            retry: Optional[RetryPolicy] = None,
-            callbacks: Callbacks = None,
-        ) -> Union[asyncio.Future[Any], concurrent.futures.Future[Any]]:
-            (fut,) = writer(
-                task,
-                [(PUSH, None)],
-                calls=[Call(func, input, retry=retry, callbacks=callbacks)],
-            )
-            assert fut is not None, "writer did not return a future for call"
-            # return a chained future to ensure commit() callback is called
-            # before the returned future is resolved, to ensure stream order etc
-            try:
-                in_async = asyncio.current_task() is not None
-            except RuntimeError:
-                in_async = False
-            # if in async context return an async future
-            # otherwise return a chained sync future
-            if in_async:
-                if isinstance(fut, asyncio.Task):
-                    sfut: Union[asyncio.Future[Any], concurrent.futures.Future[Any]] = (
-                        asyncio.Future(loop=loop)
-                    )
-                    loop.call_soon_threadsafe(chain_future, fut, sfut)
-                    return sfut
-                else:
-                    # already wrapped in a future
-                    return fut
-            else:
-                sfut = concurrent.futures.Future()
-                loop.call_soon_threadsafe(chain_future, fut, sfut)
-                return sfut
-
         loop = asyncio.get_event_loop()
         tasks = tuple(tasks)
         futures = FuturesDict(
-            callback=self.commit,
+            callback=weakref.WeakMethod(self.commit),
             event=asyncio.Event(),
             future_type=asyncio.Future,
         )
@@ -446,8 +261,17 @@ class PregelRunner:
                     retry_policy,
                     stream=self.use_astream,
                     configurable={
-                        CONFIG_KEY_SEND: partial(writer, t),
-                        CONFIG_KEY_CALL: partial(call, t),
+                        CONFIG_KEY_CALL: partial(
+                            _acall,
+                            weakref.ref(t),
+                            stream=self.use_astream,
+                            retry=retry_policy,
+                            futures=weakref.ref(futures),
+                            schedule_task=self.schedule_task,
+                            submit=self.submit,
+                            reraise=reraise,
+                            loop=loop,
+                        ),
                     },
                 )
                 self.commit(t, None)
@@ -472,14 +296,23 @@ class PregelRunner:
             if not t.writes:
                 fut = cast(
                     asyncio.Future,
-                    self.submit(
+                    self.submit()(  # type: ignore[misc]
                         arun_with_retry,
                         t,
                         retry_policy,
                         stream=self.use_astream,
                         configurable={
-                            CONFIG_KEY_SEND: partial(writer, t),
-                            CONFIG_KEY_CALL: partial(call, t),
+                            CONFIG_KEY_CALL: partial(
+                                _acall,
+                                weakref.ref(t),
+                                retry=retry_policy,
+                                stream=self.use_astream,
+                                futures=weakref.ref(futures),
+                                schedule_task=self.schedule_task,
+                                submit=self.submit,
+                                reraise=reraise,
+                                loop=loop,
+                            ),
                         },
                         __name__=t.name,
                         __cancel_on_exit__=True,
@@ -539,19 +372,20 @@ class PregelRunner:
             # for cancelled tasks, also save error in task,
             # so loop can finish super-step
             task.writes.append((ERROR, exception))
-            self.put_writes(task.id, task.writes)
+            self.put_writes()(task.id, task.writes)  # type: ignore[misc]
         elif exception:
             if isinstance(exception, GraphInterrupt):
                 # save interrupt to checkpointer
-                if interrupts := [(INTERRUPT, i) for i in exception.args[0]]:
+                if exception.args[0]:
+                    writes = [(INTERRUPT, exception.args[0])]
                     if resumes := [w for w in task.writes if w[0] == RESUME]:
-                        interrupts.extend(resumes)
-                    self.put_writes(task.id, interrupts)
+                        writes.extend(resumes)
+                    self.put_writes()(task.id, writes)  # type: ignore[misc]
             elif isinstance(exception, GraphBubbleUp):
                 raise exception
             else:
                 # save error to checkpointer
-                self.put_writes(task.id, [(ERROR, exception)])
+                self.put_writes()(task.id, [(ERROR, exception)])  # type: ignore[misc]
         else:
             if self.node_finished and (
                 task.config is None or TAG_HIDDEN not in task.config.get("tags", [])
@@ -561,7 +395,7 @@ class PregelRunner:
                 # add no writes marker
                 task.writes.append((NO_WRITES, None))
             # save task writes to checkpointer
-            self.put_writes(task.id, task.writes)
+            self.put_writes()(task.id, task.writes)  # type: ignore[misc]
 
 
 def _should_stop_others(
@@ -608,6 +442,7 @@ def _panic_or_proceed(
             done.add(fut)
         else:
             inflight.add(fut)
+    interrupts: list[GraphInterrupt] = []
     while done:
         # if any task failed
         if exc := _exception(done.pop()):
@@ -616,7 +451,14 @@ def _panic_or_proceed(
                 inflight.pop().cancel()
             # raise the exception
             if panic:
-                raise exc
+                if isinstance(exc, GraphInterrupt):
+                    # collect interrupts
+                    interrupts.append(exc)
+                else:
+                    raise exc
+    # raise combined interrupts
+    if interrupts:
+        raise GraphInterrupt(tuple(i for exc in interrupts for i in exc.args[0]))
     if inflight:
         # if we got here means we timed out
         while inflight:
@@ -624,3 +466,193 @@ def _panic_or_proceed(
             inflight.pop().cancel()
         # raise timeout error
         raise timeout_exc_cls("Timed out")
+
+
+def _call(
+    task: weakref.ref[PregelExecutableTask],
+    func: Callable[[Any], Union[Awaitable[Any], Any]],
+    input: Any,
+    *,
+    retry: Optional[RetryPolicy] = None,
+    callbacks: Callbacks = None,
+    futures: weakref.ref[FuturesDict],
+    schedule_task: weakref.ref[
+        Callable[
+            [PregelExecutableTask, int, Optional[Call]], Optional[PregelExecutableTask]
+        ]
+    ],
+    submit: weakref.ref[Submit],
+    reraise: bool,
+) -> concurrent.futures.Future[Any]:
+    if asyncio.iscoroutinefunction(func):
+        raise RuntimeError("In an sync context async tasks cannot be called")
+
+    fut: Optional[concurrent.futures.Future] = None
+    # schedule PUSH tasks, collect futures
+    scratchpad: PregelScratchpad = task().config[CONF][CONFIG_KEY_SCRATCHPAD]  # type: ignore[union-attr]
+    # schedule the next task, if the callback returns one
+    if next_task := schedule_task()(  # type: ignore[misc]
+        task(),  # type: ignore[arg-type]
+        scratchpad.call_counter(),
+        Call(func, input, retry=retry, callbacks=callbacks),
+    ):
+        if fut := next(
+            (
+                f
+                for f, t in futures().items()  # type: ignore[union-attr]
+                if t is not None and t == next_task.id
+            ),
+            None,
+        ):
+            # if the parent task was retried,
+            # the next task might already be running
+            pass
+        elif next_task.writes:
+            # if it already ran, return the result
+            fut = concurrent.futures.Future()
+            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):
+                fut.set_exception(
+                    exc if isinstance(exc, BaseException) else Exception(exc)
+                )
+            else:
+                fut.set_result(None)
+        else:
+            # schedule the next task
+            fut = submit()(  # type: ignore[misc]
+                run_with_retry,
+                next_task,
+                retry,
+                configurable={
+                    CONFIG_KEY_CALL: partial(
+                        _call,
+                        weakref.ref(next_task),
+                        futures=futures,
+                        retry=retry,
+                        callbacks=callbacks,
+                        schedule_task=schedule_task,
+                        submit=submit,
+                        reraise=reraise,
+                    ),
+                },
+                __reraise_on_exit__=reraise,
+                # starting a new task in the next tick ensures
+                # updates from this tick are committed/streamed first
+                __next_tick__=True,
+            )
+            futures()[fut] = next_task  # type: ignore[index]
+    fut = cast(Union[asyncio.Future, concurrent.futures.Future], fut)
+    # return a chained future to ensure commit() callback is called
+    # before the returned future is resolved, to ensure stream order etc
+    return chain_future(fut, concurrent.futures.Future())
+
+
+def _acall(
+    task: weakref.ref[PregelExecutableTask],
+    func: Callable[[Any], Union[Awaitable[Any], Any]],
+    input: Any,
+    *,
+    retry: Optional[RetryPolicy] = None,
+    callbacks: Callbacks = None,
+    # injected dependencies
+    futures: weakref.ref[FuturesDict],
+    schedule_task: weakref.ref[
+        Callable[
+            [PregelExecutableTask, int, Optional[Call]], Optional[PregelExecutableTask]
+        ]
+    ],
+    submit: weakref.ref[Submit],
+    loop: asyncio.AbstractEventLoop,
+    reraise: bool = False,
+    stream: bool = False,
+) -> Union[asyncio.Future[Any], concurrent.futures.Future[Any]]:
+    fut: Optional[asyncio.Future] = None
+    # schedule PUSH tasks, collect futures
+    scratchpad: PregelScratchpad = task().config[CONF][CONFIG_KEY_SCRATCHPAD]  # type: ignore[union-attr]
+    # schedule the next task, if the callback returns one
+    if next_task := schedule_task()(  # type: ignore[misc]
+        task(),  # type: ignore[arg-type]
+        scratchpad.call_counter(),
+        Call(func, input, retry=retry, callbacks=callbacks),
+    ):
+        if fut := next(
+            (
+                f
+                for f, t in futures().items()  # type: ignore[union-attr]
+                if t is not None and t == next_task.id
+            ),
+            None,
+        ):
+            # if the parent task was retried,
+            # the next task might already be running
+
+            pass
+        elif next_task.writes:
+            # if it already ran, return the result
+            fut = asyncio.Future(loop=loop)
+            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):
+                fut.set_exception(
+                    exc if isinstance(exc, BaseException) else Exception(exc)
+                )
+            else:
+                fut.set_result(None)
+            futures()[fut] = next_task  # type: ignore[index]
+        else:
+            # schedule the next task
+            fut = cast(
+                asyncio.Future,
+                submit()(  # type: ignore[misc]
+                    arun_with_retry,
+                    next_task,
+                    retry,
+                    stream=stream,
+                    configurable={
+                        CONFIG_KEY_CALL: partial(
+                            _acall,
+                            weakref.ref(next_task),
+                            stream=stream,
+                            futures=futures,
+                            schedule_task=schedule_task,
+                            submit=submit,
+                            loop=loop,
+                            reraise=reraise,
+                        ),
+                    },
+                    __name__=task().name,  # type: ignore[union-attr]
+                    __cancel_on_exit__=True,
+                    __reraise_on_exit__=reraise,
+                    # starting a new task in the next tick ensures
+                    # updates from this tick are committed/streamed first
+                    __next_tick__=True,
+                ),
+            )
+            futures()[fut] = next_task  # type: ignore[index]
+
+    fut = cast(Union[asyncio.Future, concurrent.futures.Future], fut)
+    # return a chained future to ensure commit() callback is called
+    # before the returned future is resolved, to ensure stream order etc
+    try:
+        in_async = asyncio.current_task() is not None
+    except RuntimeError:
+        in_async = False
+    # if in async context return an async future
+    # otherwise return a chained sync future
+    if in_async:
+        if isinstance(fut, asyncio.Task):
+            sfut: Union[asyncio.Future[Any], concurrent.futures.Future[Any]] = (
+                asyncio.Future(loop=loop)
+            )
+            loop.call_soon_threadsafe(chain_future, fut, sfut)
+            return sfut
+        else:
+            # already wrapped in a future
+            return fut
+    else:
+        sfut = concurrent.futures.Future()
+        loop.call_soon_threadsafe(chain_future, fut, sfut)
+        return sfut

libs/langgraph/langgraph/pregel/types.py

@@ -7,6 +7,7 @@ from langgraph.types import (
     PregelTask,
     RetryPolicy,
     StateSnapshot,
+    StateUpdate,
     StreamMode,
     StreamWriter,
     default_retry_on,
@@ -14,6 +15,7 @@ from langgraph.types import (
 
 __all__ = [
     "All",
+    "StateUpdate",
     "CachePolicy",
     "PregelExecutableTask",
     "PregelTask",

libs/langgraph/langgraph/pregel/utils.py

@@ -1,7 +1,10 @@
-from typing import Optional
+import ast
+import inspect
+import textwrap
+from typing import Any, Callable, Optional
 
 from langchain_core.runnables import RunnableLambda, RunnableSequence
-from langchain_core.runnables.utils import get_function_nonlocals
+from typing_extensions import override
 
 from langgraph.checkpoint.base import ChannelVersions
 from langgraph.pregel.protocol import PregelProtocol
@@ -55,3 +58,152 @@ def find_subgraph_pregel(candidate: Runnable) -> Optional[PregelProtocol]:
                 )
 
     return None
+
+
+def get_function_nonlocals(func: Callable) -> list[Any]:
+    """Get the nonlocal variables accessed by a function.
+
+    Args:
+        func: The function to check.
+
+    Returns:
+        List[Any]: The nonlocal variables accessed by the function.
+    """
+    try:
+        code = inspect.getsource(func)
+        tree = ast.parse(textwrap.dedent(code))
+        visitor = FunctionNonLocals()
+        visitor.visit(tree)
+        values: list[Any] = []
+        closure = (
+            inspect.getclosurevars(func.__wrapped__)
+            if hasattr(func, "__wrapped__") and callable(func.__wrapped__)
+            else inspect.getclosurevars(func)
+        )
+        candidates = {**closure.globals, **closure.nonlocals}
+        for k, v in candidates.items():
+            if k in visitor.nonlocals:
+                values.append(v)
+            for kk in visitor.nonlocals:
+                if "." in kk and kk.startswith(k):
+                    vv = v
+                    for part in kk.split(".")[1:]:
+                        if vv is None:
+                            break
+                        else:
+                            try:
+                                vv = getattr(vv, part)
+                            except AttributeError:
+                                break
+                    else:
+                        values.append(vv)
+    except (SyntaxError, TypeError, OSError, SystemError):
+        return []
+
+    return values
+
+
+class FunctionNonLocals(ast.NodeVisitor):
+    """Get the nonlocal variables accessed of a function."""
+
+    def __init__(self) -> None:
+        self.nonlocals: set[str] = set()
+
+    @override
+    def visit_FunctionDef(self, node: ast.FunctionDef) -> Any:
+        """Visit a function definition.
+
+        Args:
+            node: The node to visit.
+
+        Returns:
+            Any: The result of the visit.
+        """
+        visitor = NonLocals()
+        visitor.visit(node)
+        self.nonlocals.update(visitor.loads - visitor.stores)
+
+    @override
+    def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef) -> Any:
+        """Visit an async function definition.
+
+        Args:
+            node: The node to visit.
+
+        Returns:
+            Any: The result of the visit.
+        """
+        visitor = NonLocals()
+        visitor.visit(node)
+        self.nonlocals.update(visitor.loads - visitor.stores)
+
+    @override
+    def visit_Lambda(self, node: ast.Lambda) -> Any:
+        """Visit a lambda function.
+
+        Args:
+            node: The node to visit.
+
+        Returns:
+            Any: The result of the visit.
+        """
+        visitor = NonLocals()
+        visitor.visit(node)
+        self.nonlocals.update(visitor.loads - visitor.stores)
+
+
+class NonLocals(ast.NodeVisitor):
+    """Get nonlocal variables accessed."""
+
+    def __init__(self) -> None:
+        self.loads: set[str] = set()
+        self.stores: set[str] = set()
+
+    @override
+    def visit_Name(self, node: ast.Name) -> Any:
+        """Visit a name node.
+
+        Args:
+            node: The node to visit.
+
+        Returns:
+            Any: The result of the visit.
+        """
+        if isinstance(node.ctx, ast.Load):
+            self.loads.add(node.id)
+        elif isinstance(node.ctx, ast.Store):
+            self.stores.add(node.id)
+
+    @override
+    def visit_Attribute(self, node: ast.Attribute) -> Any:
+        """Visit an attribute node.
+
+        Args:
+            node: The node to visit.
+
+        Returns:
+            Any: The result of the visit.
+        """
+        if isinstance(node.ctx, ast.Load):
+            parent = node.value
+            attr_expr = node.attr
+            while isinstance(parent, ast.Attribute):
+                attr_expr = parent.attr + "." + attr_expr
+                parent = parent.value
+            if isinstance(parent, ast.Name):
+                self.loads.add(parent.id + "." + attr_expr)
+                self.loads.discard(parent.id)
+            elif isinstance(parent, ast.Call):
+                if isinstance(parent.func, ast.Name):
+                    self.loads.add(parent.func.id)
+                else:
+                    parent = parent.func
+                    attr_expr = ""
+                    while isinstance(parent, ast.Attribute):
+                        if attr_expr:
+                            attr_expr = parent.attr + "." + attr_expr
+                        else:
+                            attr_expr = parent.attr
+                        parent = parent.value
+                    if isinstance(parent, ast.Name):
+                        self.loads.add(parent.id + "." + attr_expr)

libs/langgraph/langgraph/types.py

@@ -130,7 +130,12 @@ class Interrupt:
     value: Any
     resumable: bool = False
     ns: Optional[Sequence[str]] = None
-    when: Literal["during"] = "during"
+    when: Literal["during"] = dataclasses.field(default="during", repr=False)
+
+
+class StateUpdate(NamedTuple):
+    values: Optional[dict[str, Any]]
+    as_node: Optional[str] = None
 
 
 class PregelTask(NamedTuple):
@@ -143,13 +148,20 @@ class PregelTask(NamedTuple):
     result: Optional[Any] = None
 
 
-class PregelExecutableTask(NamedTuple):
+if sys.version_info > (3, 11):
+    _T_DC_KWARGS = {"weakref_slot": True, "slots": True, "frozen": True}
+else:
+    _T_DC_KWARGS = {"frozen": True}
+
+
+@dataclasses.dataclass(**_T_DC_KWARGS)
+class PregelExecutableTask:
     name: str
     input: Any
     proc: Runnable
     writes: deque[tuple[str, Any]]
     config: RunnableConfig
-    triggers: list[str]
+    triggers: Sequence[str]
     retry_policy: Optional[RetryPolicy]
     cache_policy: Optional[CachePolicy]
     id: str
@@ -351,20 +363,11 @@ class PregelScratchpad:
     call_counter: Callable[[], int]
     # interrupt
     interrupt_counter: Callable[[], int]
+    get_null_resume: Callable[[bool], Any]
     resume: list[Any]
-    null_resume: Optional[Any]
-    _consume_null_resume: Callable[[], None]
     # subgraph
     subgraph_counter: Callable[[], int]
 
-    def consume_null_resume(self) -> Any:
-        if self.null_resume is not None:
-            value = self.null_resume
-            self._consume_null_resume()
-            self.null_resume = None
-            return value
-        raise ValueError("No null resume to consume")
-
 
 def interrupt(value: Any) -> Any:
     """Interrupt the graph with a resumable exception from within a node.
@@ -480,9 +483,9 @@ def interrupt(value: Any) -> Any:
         if idx < len(scratchpad.resume):
             return scratchpad.resume[idx]
     # find current resume value
-    if scratchpad.null_resume is not None:
+    v = scratchpad.get_null_resume(True)
+    if v is not None:
         assert len(scratchpad.resume) == idx, (scratchpad.resume, idx)
-        v = scratchpad.consume_null_resume()
         scratchpad.resume.append(v)
         conf[CONFIG_KEY_SEND]([(RESUME, scratchpad.resume)])
         return v

libs/langgraph/langgraph/utils/runnable.py

@@ -2,8 +2,8 @@ import asyncio
 import enum
 import inspect
 import sys
-from contextlib import AsyncExitStack
-from contextvars import copy_context
+from contextlib import AsyncExitStack, contextmanager
+from contextvars import Context, Token, copy_context
 from functools import partial, wraps
 from typing import (
     Any,
@@ -11,6 +11,7 @@ from typing import (
     Awaitable,
     Callable,
     Coroutine,
+    Generator,
     Iterator,
     Optional,
     Protocol,
@@ -53,13 +54,69 @@ from langgraph.utils.config import (
     patch_config,
 )
 
-try:
-    from langchain_core.runnables.config import _set_config_context
-except ImportError:
-    # For forwards compatibility
-    def _set_config_context(context: RunnableConfig) -> None:  # type: ignore
-        """Set the context for the current thread."""
-        var_child_runnable_config.set(context)
+
+def _set_config_context(
+    config: RunnableConfig,
+) -> tuple[Token[Optional[RunnableConfig]], Optional[dict[str, Any]]]:
+    """Set the child Runnable config + tracing context.
+
+    Args:
+        config (RunnableConfig): The config to set.
+    """
+    from langchain_core.tracers.langchain import LangChainTracer
+
+    config_token = var_child_runnable_config.set(config)
+    current_context = None
+    if (
+        (callbacks := config.get("callbacks"))
+        and (
+            parent_run_id := getattr(callbacks, "parent_run_id", None)
+        )  # Is callback manager
+        and (
+            tracer := next(
+                (
+                    handler
+                    for handler in getattr(callbacks, "handlers", [])
+                    if isinstance(handler, LangChainTracer)
+                ),
+                None,
+            )
+        )
+        and (run := tracer.run_map.get(str(parent_run_id)))
+    ):
+        from langsmith.run_helpers import _set_tracing_context, get_tracing_context
+
+        current_context = get_tracing_context()
+        _set_tracing_context({"parent": run})
+    return config_token, current_context
+
+
+@contextmanager
+def set_config_context(config: RunnableConfig) -> Generator[Context, None, None]:
+    """Set the child Runnable config + tracing context.
+
+    Args:
+        config (RunnableConfig): The config to set.
+    """
+    from langsmith.run_helpers import _set_tracing_context
+
+    ctx = copy_context()
+    config_token, _ = ctx.run(_set_config_context, config)
+    try:
+        yield ctx
+    finally:
+        ctx.run(var_child_runnable_config.reset, config_token)
+        ctx.run(
+            _set_tracing_context,
+            {
+                "parent": None,
+                "project_name": None,
+                "tags": None,
+                "metadata": None,
+                "enabled": None,
+                "client": None,
+            },
+        )
 
 
 # Before Python 3.11 native StrEnum is not available
@@ -286,7 +343,6 @@ class RunnableCallable(Runnable):
 
             kwargs[kw] = _conf.get(config_key, default_value)
 
-        context = copy_context()
         if self.trace:
             callback_manager = get_callback_manager_for_config(config, self.tags)
             run_manager = callback_manager.on_chain_start(
@@ -297,17 +353,16 @@ class RunnableCallable(Runnable):
             )
             try:
                 child_config = patch_config(config, callbacks=run_manager.get_child())
-                context = copy_context()
-                context.run(_set_config_context, child_config)
-                ret = context.run(self.func, *args, **kwargs)
+                with set_config_context(child_config) as context:
+                    ret = context.run(self.func, *args, **kwargs)
             except BaseException as e:
                 run_manager.on_chain_error(e)
                 raise
             else:
                 run_manager.on_chain_end(ret)
         else:
-            context.run(_set_config_context, config)
-            ret = context.run(self.func, *args, **kwargs)
+            with set_config_context(config) as context:
+                ret = context.run(self.func, *args, **kwargs)
         if isinstance(ret, Runnable) and self.recurse:
             return ret.invoke(input, config)
         return ret
@@ -342,7 +397,6 @@ class RunnableCallable(Runnable):
                     f"Missing required config key '{config_key}' for '{self.name}'."
                 )
             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)
             run_manager = await callback_manager.on_chain_start(
@@ -353,24 +407,24 @@ class RunnableCallable(Runnable):
             )
             try:
                 child_config = patch_config(config, callbacks=run_manager.get_child())
-                context.run(_set_config_context, child_config)
-                coro = cast(Coroutine[None, None, Any], self.afunc(*args, **kwargs))
-                if ASYNCIO_ACCEPTS_CONTEXT:
-                    ret = await asyncio.create_task(coro, context=context)
-                else:
-                    ret = await coro
+                with set_config_context(child_config) as context:
+                    coro = cast(Coroutine[None, None, Any], self.afunc(*args, **kwargs))
+                    if ASYNCIO_ACCEPTS_CONTEXT:
+                        ret = await asyncio.create_task(coro, context=context)
+                    else:
+                        ret = await coro
             except BaseException as e:
                 await run_manager.on_chain_error(e)
                 raise
             else:
                 await run_manager.on_chain_end(ret)
         else:
-            context.run(_set_config_context, config)
-            if ASYNCIO_ACCEPTS_CONTEXT:
-                coro = cast(Coroutine[None, None, Any], self.afunc(*args, **kwargs))
-                ret = await asyncio.create_task(coro, context=context)
-            else:
-                ret = await self.afunc(*args, **kwargs)
+            with set_config_context(config) as context:
+                if ASYNCIO_ACCEPTS_CONTEXT:
+                    coro = cast(Coroutine[None, None, Any], self.afunc(*args, **kwargs))
+                    ret = await asyncio.create_task(coro, context=context)
+                else:
+                    ret = await self.afunc(*args, **kwargs)
         if isinstance(ret, Runnable) and self.recurse:
             return await ret.ainvoke(input, config)
         return ret

libs/langgraph/pyproject.toml

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

libs/langgraph/tests/test_large_cases.py

@@ -2483,7 +2483,7 @@ def test_prebuilt_tool_chat(snapshot: SnapshotAssertion) -> None:
             {
                 "langgraph_step": 1,
                 "langgraph_node": "agent",
-                "langgraph_triggers": ["start:agent"],
+                "langgraph_triggers": ("branch:to:agent", "start:agent", "tools"),
                 "langgraph_path": (PULL, "agent"),
                 "langgraph_checkpoint_ns": AnyStr("agent:"),
                 "checkpoint_ns": AnyStr("agent:"),
@@ -2500,7 +2500,7 @@ def test_prebuilt_tool_chat(snapshot: SnapshotAssertion) -> None:
             {
                 "langgraph_step": 2,
                 "langgraph_node": "tools",
-                "langgraph_triggers": ["branch:agent:should_continue:tools"],
+                "langgraph_triggers": ("branch:to:tools",),
                 "langgraph_path": (PULL, "tools"),
                 "langgraph_checkpoint_ns": AnyStr("tools:"),
             },
@@ -2542,7 +2542,7 @@ def test_prebuilt_tool_chat(snapshot: SnapshotAssertion) -> None:
             {
                 "langgraph_step": 3,
                 "langgraph_node": "agent",
-                "langgraph_triggers": ["tools"],
+                "langgraph_triggers": ("branch:to:agent", "start:agent", "tools"),
                 "langgraph_path": (PULL, "agent"),
                 "langgraph_checkpoint_ns": AnyStr("agent:"),
                 "checkpoint_ns": AnyStr("agent:"),
@@ -2559,7 +2559,7 @@ def test_prebuilt_tool_chat(snapshot: SnapshotAssertion) -> None:
             {
                 "langgraph_step": 4,
                 "langgraph_node": "tools",
-                "langgraph_triggers": ["branch:agent:should_continue:tools"],
+                "langgraph_triggers": ("branch:to:tools",),
                 "langgraph_path": (PULL, "tools"),
                 "langgraph_checkpoint_ns": AnyStr("tools:"),
             },
@@ -2573,7 +2573,7 @@ def test_prebuilt_tool_chat(snapshot: SnapshotAssertion) -> None:
             {
                 "langgraph_step": 4,
                 "langgraph_node": "tools",
-                "langgraph_triggers": ["branch:agent:should_continue:tools"],
+                "langgraph_triggers": ("branch:to:tools",),
                 "langgraph_path": (PULL, "tools"),
                 "langgraph_checkpoint_ns": AnyStr("tools:"),
             },
@@ -2585,7 +2585,7 @@ def test_prebuilt_tool_chat(snapshot: SnapshotAssertion) -> None:
             {
                 "langgraph_step": 5,
                 "langgraph_node": "agent",
-                "langgraph_triggers": ["tools"],
+                "langgraph_triggers": ("branch:to:agent", "start:agent", "tools"),
                 "langgraph_path": (PULL, "agent"),
                 "langgraph_checkpoint_ns": AnyStr("agent:"),
                 "checkpoint_ns": AnyStr("agent:"),
@@ -5501,7 +5501,10 @@ def test_in_one_fan_out_out_one_graph_state() -> None:
                     "id": AnyStr(),
                     "name": "rewrite_query",
                     "input": {"query": "what is weather in sf", "docs": []},
-                    "triggers": ["start:rewrite_query"],
+                    "triggers": (
+                        "branch:to:rewrite_query",
+                        "start:rewrite_query",
+                    ),
                 },
             },
         ),
@@ -5532,7 +5535,10 @@ def test_in_one_fan_out_out_one_graph_state() -> None:
                     "id": AnyStr(),
                     "name": "retriever_one",
                     "input": {"query": "query: what is weather in sf", "docs": []},
-                    "triggers": ["rewrite_query"],
+                    "triggers": (
+                        "branch:to:retriever_one",
+                        "rewrite_query",
+                    ),
                 },
             },
         ),
@@ -5546,7 +5552,10 @@ def test_in_one_fan_out_out_one_graph_state() -> None:
                     "id": AnyStr(),
                     "name": "retriever_two",
                     "input": {"query": "query: what is weather in sf", "docs": []},
-                    "triggers": ["rewrite_query"],
+                    "triggers": (
+                        "branch:to:retriever_two",
+                        "rewrite_query",
+                    ),
                 },
             },
         ),
@@ -5608,7 +5617,7 @@ def test_in_one_fan_out_out_one_graph_state() -> None:
                         "query": "query: what is weather in sf",
                         "docs": ["doc1", "doc2", "doc3", "doc4"],
                     },
-                    "triggers": ["retriever_one", "retriever_two"],
+                    "triggers": ("branch:to:qa", "retriever_one", "retriever_two"),
                 },
             },
         ),
@@ -6634,7 +6643,7 @@ def test_branch_then(
                 "id": AnyStr(),
                 "name": "prepare",
                 "input": {"my_key": "value", "market": "DE"},
-                "triggers": ["start:prepare"],
+                "triggers": ("branch:to:prepare", "start:prepare"),
             },
         },
         {
@@ -6706,7 +6715,7 @@ def test_branch_then(
                 "id": AnyStr(),
                 "name": "tool_two_slow",
                 "input": {"my_key": "value prepared", "market": "DE"},
-                "triggers": ["branch:prepare:condition:tool_two_slow"],
+                "triggers": ("branch:to:tool_two_slow",),
             },
         },
         {
@@ -6773,7 +6782,10 @@ def test_branch_then(
                 "id": AnyStr(),
                 "name": "finish",
                 "input": {"my_key": "value prepared slow", "market": "DE"},
-                "triggers": ["branch:prepare:condition::then"],
+                "triggers": (
+                    "branch:prepare:condition::then",
+                    "branch:to:finish",
+                ),
             },
         },
         {
@@ -7783,7 +7795,7 @@ def test_nested_graph_state(
                         "langgraph_node": "inner",
                         "langgraph_path": [PULL, "inner"],
                         "langgraph_step": 2,
-                        "langgraph_triggers": ["outer_1"],
+                        "langgraph_triggers": ["branch:to:inner", "outer_1"],
                         "langgraph_checkpoint_ns": AnyStr("inner:"),
                     },
                     created_at=AnyStr(),
@@ -7978,7 +7990,7 @@ def test_nested_graph_state(
                 "langgraph_node": "inner",
                 "langgraph_path": [PULL, "inner"],
                 "langgraph_step": 2,
-                "langgraph_triggers": ["outer_1"],
+                "langgraph_triggers": ["branch:to:inner", "outer_1"],
                 "langgraph_checkpoint_ns": AnyStr("inner:"),
             },
             created_at=AnyStr(),
@@ -8021,7 +8033,7 @@ def test_nested_graph_state(
                 "langgraph_node": "inner",
                 "langgraph_path": [PULL, "inner"],
                 "langgraph_step": 2,
-                "langgraph_triggers": ["outer_1"],
+                "langgraph_triggers": ["branch:to:inner", "outer_1"],
                 "langgraph_checkpoint_ns": AnyStr("inner:"),
             },
             created_at=AnyStr(),
@@ -8070,7 +8082,7 @@ def test_nested_graph_state(
                 "langgraph_node": "inner",
                 "langgraph_path": [PULL, "inner"],
                 "langgraph_step": 2,
-                "langgraph_triggers": ["outer_1"],
+                "langgraph_triggers": ["branch:to:inner", "outer_1"],
                 "langgraph_checkpoint_ns": AnyStr("inner:"),
             },
             created_at=AnyStr(),
@@ -8504,7 +8516,7 @@ def test_doubly_nested_graph_state(
             "langgraph_node": "child_1",
             "langgraph_path": [PULL, AnyStr("child_1")],
             "langgraph_step": 1,
-            "langgraph_triggers": [AnyStr("start:child_1")],
+            "langgraph_triggers": ["branch:to:child_1", AnyStr("start:child_1")],
         },
         created_at=AnyStr(),
         parent_config=(
@@ -8588,7 +8600,10 @@ def test_doubly_nested_graph_state(
                                         AnyStr("child_1"),
                                     ],
                                     "langgraph_step": 1,
-                                    "langgraph_triggers": [AnyStr("start:child_1")],
+                                    "langgraph_triggers": [
+                                        "branch:to:child_1",
+                                        AnyStr("start:child_1"),
+                                    ],
                                 },
                                 created_at=AnyStr(),
                                 parent_config=(
@@ -8635,7 +8650,7 @@ def test_doubly_nested_graph_state(
                         "langgraph_node": "child",
                         "langgraph_path": [PULL, AnyStr("child")],
                         "langgraph_step": 2,
-                        "langgraph_triggers": [AnyStr("parent_1")],
+                        "langgraph_triggers": ["branch:to:child", AnyStr("parent_1")],
                         "langgraph_checkpoint_ns": AnyStr("child:"),
                     },
                     created_at=AnyStr(),
@@ -8931,7 +8946,7 @@ def test_doubly_nested_graph_state(
                 "langgraph_node": "child",
                 "langgraph_path": [PULL, AnyStr("child")],
                 "langgraph_step": 2,
-                "langgraph_triggers": [AnyStr("parent_1")],
+                "langgraph_triggers": ["branch:to:child", AnyStr("parent_1")],
                 "langgraph_checkpoint_ns": AnyStr("child:"),
             },
             created_at=AnyStr(),
@@ -8970,7 +8985,7 @@ def test_doubly_nested_graph_state(
                 "langgraph_node": "child",
                 "langgraph_path": [PULL, AnyStr("child")],
                 "langgraph_step": 2,
-                "langgraph_triggers": [AnyStr("parent_1")],
+                "langgraph_triggers": ["branch:to:child", AnyStr("parent_1")],
                 "langgraph_checkpoint_ns": AnyStr("child:"),
             },
             created_at=AnyStr(),
@@ -9022,7 +9037,7 @@ def test_doubly_nested_graph_state(
                 "langgraph_node": "child",
                 "langgraph_path": [PULL, AnyStr("child")],
                 "langgraph_step": 2,
-                "langgraph_triggers": [AnyStr("parent_1")],
+                "langgraph_triggers": ["branch:to:child", AnyStr("parent_1")],
                 "langgraph_checkpoint_ns": AnyStr("child:"),
             },
             created_at=AnyStr(),
@@ -9076,7 +9091,7 @@ def test_doubly_nested_graph_state(
                     AnyStr("child_1"),
                 ],
                 "langgraph_step": 1,
-                "langgraph_triggers": [AnyStr("start:child_1")],
+                "langgraph_triggers": ["branch:to:child_1", AnyStr("start:child_1")],
             },
             created_at=AnyStr(),
             parent_config={
@@ -9131,7 +9146,7 @@ def test_doubly_nested_graph_state(
                     AnyStr("child_1"),
                 ],
                 "langgraph_step": 1,
-                "langgraph_triggers": [AnyStr("start:child_1")],
+                "langgraph_triggers": ["branch:to:child_1", AnyStr("start:child_1")],
             },
             created_at=AnyStr(),
             parent_config={
@@ -9193,7 +9208,7 @@ def test_doubly_nested_graph_state(
                     AnyStr("child_1"),
                 ],
                 "langgraph_step": 1,
-                "langgraph_triggers": [AnyStr("start:child_1")],
+                "langgraph_triggers": ["branch:to:child_1", AnyStr("start:child_1")],
             },
             created_at=AnyStr(),
             parent_config={
@@ -9255,7 +9270,7 @@ def test_doubly_nested_graph_state(
                     AnyStr("child_1"),
                 ],
                 "langgraph_step": 1,
-                "langgraph_triggers": [AnyStr("start:child_1")],
+                "langgraph_triggers": ["branch:to:child_1", AnyStr("start:child_1")],
             },
             created_at=AnyStr(),
             parent_config=None,
@@ -10378,9 +10393,7 @@ def test_weather_subgraph(
                         "langgraph_node": "weather_graph",
                         "langgraph_path": [PULL, "weather_graph"],
                         "langgraph_step": 2,
-                        "langgraph_triggers": [
-                            "branch:router_node:route_after_prediction:weather_graph"
-                        ],
+                        "langgraph_triggers": ["branch:to:weather_graph"],
                         "langgraph_checkpoint_ns": AnyStr("weather_graph:"),
                     },
                     created_at=AnyStr(),
@@ -10492,9 +10505,7 @@ def test_weather_subgraph(
                         "langgraph_node": "weather_graph",
                         "langgraph_path": [PULL, "weather_graph"],
                         "langgraph_step": 2,
-                        "langgraph_triggers": [
-                            "branch:router_node:route_after_prediction:weather_graph"
-                        ],
+                        "langgraph_triggers": ["branch:to:weather_graph"],
                         "langgraph_checkpoint_ns": AnyStr("weather_graph:"),
                     },
                     created_at=AnyStr(),

libs/langgraph/tests/test_large_cases_async.py

@@ -2300,7 +2300,11 @@ async def test_prebuilt_tool_chat() -> None:
             {
                 "langgraph_step": 1,
                 "langgraph_node": "agent",
-                "langgraph_triggers": ["start:agent"],
+                "langgraph_triggers": (
+                    "branch:to:agent",
+                    "start:agent",
+                    "tools",
+                ),
                 "langgraph_path": ("__pregel_pull", "agent"),
                 "langgraph_checkpoint_ns": AnyStr("agent:"),
                 "checkpoint_ns": AnyStr("agent:"),
@@ -2317,7 +2321,7 @@ async def test_prebuilt_tool_chat() -> None:
             {
                 "langgraph_step": 2,
                 "langgraph_node": "tools",
-                "langgraph_triggers": ["branch:agent:should_continue:tools"],
+                "langgraph_triggers": ("branch:to:tools",),
                 "langgraph_path": ("__pregel_pull", "tools"),
                 "langgraph_checkpoint_ns": AnyStr("tools:"),
             },
@@ -2359,7 +2363,11 @@ async def test_prebuilt_tool_chat() -> None:
             {
                 "langgraph_step": 3,
                 "langgraph_node": "agent",
-                "langgraph_triggers": ["tools"],
+                "langgraph_triggers": (
+                    "branch:to:agent",
+                    "start:agent",
+                    "tools",
+                ),
                 "langgraph_path": ("__pregel_pull", "agent"),
                 "langgraph_checkpoint_ns": AnyStr("agent:"),
                 "checkpoint_ns": AnyStr("agent:"),
@@ -2376,7 +2384,7 @@ async def test_prebuilt_tool_chat() -> None:
             {
                 "langgraph_step": 4,
                 "langgraph_node": "tools",
-                "langgraph_triggers": ["branch:agent:should_continue:tools"],
+                "langgraph_triggers": ("branch:to:tools",),
                 "langgraph_path": ("__pregel_pull", "tools"),
                 "langgraph_checkpoint_ns": AnyStr("tools:"),
             },
@@ -2390,7 +2398,7 @@ async def test_prebuilt_tool_chat() -> None:
             {
                 "langgraph_step": 4,
                 "langgraph_node": "tools",
-                "langgraph_triggers": ["branch:agent:should_continue:tools"],
+                "langgraph_triggers": ("branch:to:tools",),
                 "langgraph_path": ("__pregel_pull", "tools"),
                 "langgraph_checkpoint_ns": AnyStr("tools:"),
             },
@@ -2402,7 +2410,11 @@ async def test_prebuilt_tool_chat() -> None:
             {
                 "langgraph_step": 5,
                 "langgraph_node": "agent",
-                "langgraph_triggers": ["tools"],
+                "langgraph_triggers": (
+                    "branch:to:agent",
+                    "start:agent",
+                    "tools",
+                ),
                 "langgraph_path": ("__pregel_pull", "agent"),
                 "langgraph_checkpoint_ns": AnyStr("agent:"),
                 "checkpoint_ns": AnyStr("agent:"),
@@ -3883,7 +3895,10 @@ async def test_in_one_fan_out_out_one_graph_state() -> None:
                     "id": AnyStr(),
                     "name": "rewrite_query",
                     "input": {"query": "what is weather in sf", "docs": []},
-                    "triggers": ["start:rewrite_query"],
+                    "triggers": (
+                        "branch:to:rewrite_query",
+                        "start:rewrite_query",
+                    ),
                 },
             },
         ),
@@ -3914,7 +3929,10 @@ async def test_in_one_fan_out_out_one_graph_state() -> None:
                     "id": AnyStr(),
                     "name": "retriever_one",
                     "input": {"query": "query: what is weather in sf", "docs": []},
-                    "triggers": ["rewrite_query"],
+                    "triggers": (
+                        "branch:to:retriever_one",
+                        "rewrite_query",
+                    ),
                 },
             },
         ),
@@ -3928,7 +3946,10 @@ async def test_in_one_fan_out_out_one_graph_state() -> None:
                     "id": AnyStr(),
                     "name": "retriever_two",
                     "input": {"query": "query: what is weather in sf", "docs": []},
-                    "triggers": ["rewrite_query"],
+                    "triggers": (
+                        "branch:to:retriever_two",
+                        "rewrite_query",
+                    ),
                 },
             },
         ),
@@ -3990,7 +4011,7 @@ async def test_in_one_fan_out_out_one_graph_state() -> None:
                         "query": "query: what is weather in sf",
                         "docs": ["doc1", "doc2", "doc3", "doc4"],
                     },
-                    "triggers": ["retriever_one", "retriever_two"],
+                    "triggers": ("branch:to:qa", "retriever_one", "retriever_two"),
                 },
             },
         ),
@@ -4465,7 +4486,10 @@ async def test_branch_then(checkpointer_name: str) -> None:
                     "id": AnyStr(),
                     "name": "prepare",
                     "input": {"my_key": "value", "market": "DE"},
-                    "triggers": ["start:prepare"],
+                    "triggers": (
+                        "branch:to:prepare",
+                        "start:prepare",
+                    ),
                 },
             },
             {
@@ -4537,7 +4561,7 @@ async def test_branch_then(checkpointer_name: str) -> None:
                     "id": AnyStr(),
                     "name": "tool_two_slow",
                     "input": {"my_key": "value prepared", "market": "DE"},
-                    "triggers": ["branch:prepare:condition:tool_two_slow"],
+                    "triggers": ("branch:to:tool_two_slow",),
                 },
             },
             {
@@ -4609,7 +4633,10 @@ async def test_branch_then(checkpointer_name: str) -> None:
                     "id": AnyStr(),
                     "name": "finish",
                     "input": {"my_key": "value prepared slow", "market": "DE"},
-                    "triggers": ["branch:prepare:condition::then"],
+                    "triggers": (
+                        "branch:prepare:condition::then",
+                        "branch:to:finish",
+                    ),
                 },
             },
             {
@@ -4778,7 +4805,10 @@ async def test_branch_then(checkpointer_name: str) -> None:
                     "id": AnyStr(),
                     "name": "prepare",
                     "input": {"my_key": "value", "market": "DE"},
-                    "triggers": ["start:prepare"],
+                    "triggers": (
+                        "branch:to:prepare",
+                        "start:prepare",
+                    ),
                 },
             },
             {
@@ -5333,7 +5363,7 @@ async def test_nested_graph_state(checkpointer_name: str) -> None:
                             "langgraph_node": "inner",
                             "langgraph_path": [PULL, "inner"],
                             "langgraph_step": 2,
-                            "langgraph_triggers": ["outer_1"],
+                            "langgraph_triggers": ["branch:to:inner", "outer_1"],
                             "langgraph_checkpoint_ns": AnyStr("inner:"),
                         },
                         created_at=AnyStr(),
@@ -5530,7 +5560,7 @@ async def test_nested_graph_state(checkpointer_name: str) -> None:
                     "langgraph_node": "inner",
                     "langgraph_path": [PULL, "inner"],
                     "langgraph_step": 2,
-                    "langgraph_triggers": ["outer_1"],
+                    "langgraph_triggers": ["branch:to:inner", "outer_1"],
                     "langgraph_checkpoint_ns": AnyStr("inner:"),
                 },
                 created_at=AnyStr(),
@@ -5573,7 +5603,7 @@ async def test_nested_graph_state(checkpointer_name: str) -> None:
                     "langgraph_node": "inner",
                     "langgraph_path": [PULL, "inner"],
                     "langgraph_step": 2,
-                    "langgraph_triggers": ["outer_1"],
+                    "langgraph_triggers": ["branch:to:inner", "outer_1"],
                     "langgraph_checkpoint_ns": AnyStr("inner:"),
                 },
                 created_at=AnyStr(),
@@ -5622,7 +5652,7 @@ async def test_nested_graph_state(checkpointer_name: str) -> None:
                     "langgraph_node": "inner",
                     "langgraph_path": [PULL, "inner"],
                     "langgraph_step": 2,
-                    "langgraph_triggers": ["outer_1"],
+                    "langgraph_triggers": ["branch:to:inner", "outer_1"],
                     "langgraph_checkpoint_ns": AnyStr("inner:"),
                 },
                 created_at=AnyStr(),
@@ -6060,7 +6090,7 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None:
                 "langgraph_node": "child_1",
                 "langgraph_path": [PULL, AnyStr("child_1")],
                 "langgraph_step": 1,
-                "langgraph_triggers": [AnyStr("start:child_1")],
+                "langgraph_triggers": ["branch:to:child_1", "start:child_1"],
             },
             created_at=AnyStr(),
             parent_config=(
@@ -6146,7 +6176,10 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None:
                                             AnyStr("child_1"),
                                         ],
                                         "langgraph_step": 1,
-                                        "langgraph_triggers": [AnyStr("start:child_1")],
+                                        "langgraph_triggers": [
+                                            "branch:to:child_1",
+                                            "start:child_1",
+                                        ],
                                     },
                                     created_at=AnyStr(),
                                     parent_config=(
@@ -6195,7 +6228,10 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None:
                             "langgraph_node": "child",
                             "langgraph_path": [PULL, AnyStr("child")],
                             "langgraph_step": 2,
-                            "langgraph_triggers": [AnyStr("parent_1")],
+                            "langgraph_triggers": [
+                                "branch:to:child",
+                                AnyStr("parent_1"),
+                            ],
                             "langgraph_checkpoint_ns": AnyStr("child:"),
                         },
                         created_at=AnyStr(),
@@ -6493,7 +6529,7 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None:
                     "langgraph_node": "child",
                     "langgraph_path": [PULL, AnyStr("child")],
                     "langgraph_step": 2,
-                    "langgraph_triggers": [AnyStr("parent_1")],
+                    "langgraph_triggers": ["branch:to:child", AnyStr("parent_1")],
                     "langgraph_checkpoint_ns": AnyStr("child:"),
                 },
                 created_at=AnyStr(),
@@ -6532,7 +6568,7 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None:
                     "langgraph_node": "child",
                     "langgraph_path": [PULL, AnyStr("child")],
                     "langgraph_step": 2,
-                    "langgraph_triggers": [AnyStr("parent_1")],
+                    "langgraph_triggers": ["branch:to:child", AnyStr("parent_1")],
                     "langgraph_checkpoint_ns": AnyStr("child:"),
                 },
                 created_at=AnyStr(),
@@ -6584,7 +6620,7 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None:
                     "langgraph_node": "child",
                     "langgraph_path": [PULL, AnyStr("child")],
                     "langgraph_step": 2,
-                    "langgraph_triggers": [AnyStr("parent_1")],
+                    "langgraph_triggers": ["branch:to:child", AnyStr("parent_1")],
                     "langgraph_checkpoint_ns": AnyStr("child:"),
                 },
                 created_at=AnyStr(),
@@ -6642,7 +6678,10 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None:
                         AnyStr("child_1"),
                     ],
                     "langgraph_step": 1,
-                    "langgraph_triggers": [AnyStr("start:child_1")],
+                    "langgraph_triggers": [
+                        "branch:to:child_1",
+                        AnyStr("start:child_1"),
+                    ],
                 },
                 created_at=AnyStr(),
                 parent_config={
@@ -6697,7 +6736,10 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None:
                         AnyStr("child_1"),
                     ],
                     "langgraph_step": 1,
-                    "langgraph_triggers": [AnyStr("start:child_1")],
+                    "langgraph_triggers": [
+                        "branch:to:child_1",
+                        AnyStr("start:child_1"),
+                    ],
                 },
                 created_at=AnyStr(),
                 parent_config={
@@ -6759,7 +6801,10 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None:
                         AnyStr("child_1"),
                     ],
                     "langgraph_step": 1,
-                    "langgraph_triggers": [AnyStr("start:child_1")],
+                    "langgraph_triggers": [
+                        "branch:to:child_1",
+                        AnyStr("start:child_1"),
+                    ],
                 },
                 created_at=AnyStr(),
                 parent_config={
@@ -6821,7 +6866,10 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None:
                         AnyStr("child_1"),
                     ],
                     "langgraph_step": 1,
-                    "langgraph_triggers": [AnyStr("start:child_1")],
+                    "langgraph_triggers": [
+                        "branch:to:child_1",
+                        AnyStr("start:child_1"),
+                    ],
                 },
                 created_at=AnyStr(),
                 parent_config=None,
@@ -7231,9 +7279,7 @@ async def test_weather_subgraph(
                             "langgraph_node": "weather_graph",
                             "langgraph_path": [PULL, "weather_graph"],
                             "langgraph_step": 2,
-                            "langgraph_triggers": [
-                                "branch:router_node:route_after_prediction:weather_graph"
-                            ],
+                            "langgraph_triggers": ["branch:to:weather_graph"],
                             "langgraph_checkpoint_ns": AnyStr("weather_graph:"),
                         },
                         created_at=AnyStr(),
@@ -7347,9 +7393,7 @@ async def test_weather_subgraph(
                             "langgraph_node": "weather_graph",
                             "langgraph_path": [PULL, "weather_graph"],
                             "langgraph_step": 2,
-                            "langgraph_triggers": [
-                                "branch:router_node:route_after_prediction:weather_graph"
-                            ],
+                            "langgraph_triggers": ["branch:to:weather_graph"],
                             "langgraph_checkpoint_ns": AnyStr("weather_graph:"),
                         },
                         created_at=AnyStr(),

libs/langgraph/tests/test_pregel.py

@@ -1,5 +1,6 @@
 import enum
 import functools
+import gc
 import json
 import logging
 import operator
@@ -62,13 +63,16 @@ from langgraph.graph import END, Graph, StateGraph
 from langgraph.graph.message import MessageGraph, MessagesState, add_messages
 from langgraph.prebuilt.tool_node import ToolNode
 from langgraph.pregel import Channel, GraphRecursionError, Pregel, StateSnapshot
+from langgraph.pregel.loop import SyncPregelLoop
 from langgraph.pregel.retry import RetryPolicy
+from langgraph.pregel.runner import PregelRunner
 from langgraph.store.base import BaseStore
 from langgraph.types import (
     Command,
     Interrupt,
     PregelTask,
     Send,
+    StateUpdate,
     StreamWriter,
     interrupt,
 )
@@ -817,7 +821,7 @@ def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None:
                 "id": AnyStr(),
                 "name": "one",
                 "input": 2,
-                "triggers": ["input"],
+                "triggers": ("input",),
             },
         },
         {
@@ -828,7 +832,7 @@ def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None:
                 "id": AnyStr(),
                 "name": "two",
                 "input": [12],
-                "triggers": ["inbox"],
+                "triggers": ("inbox",),
             },
         },
         {
@@ -863,7 +867,7 @@ def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None:
                 "id": AnyStr(),
                 "name": "two",
                 "input": [3],
-                "triggers": ["inbox"],
+                "triggers": ("inbox",),
             },
         },
         {
@@ -3247,14 +3251,24 @@ def test_in_one_fan_out_state_graph_waiting_edge_plus_regular(
 
     assert [
         c for c in app_w_interrupt.stream({"query": "what is weather in sf"}, config)
-    ] == [
-        {"rewrite_query": {"query": "query: what is weather in sf"}},
-        {"qa": {"answer": ""}},
-        {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}},
-        {"retriever_two": {"docs": ["doc3", "doc4"]}},
-        {"retriever_one": {"docs": ["doc1", "doc2"]}},
-        {"__interrupt__": ()},
-    ]
+    ] in (
+        [
+            {"rewrite_query": {"query": "query: what is weather in sf"}},
+            {"qa": {"answer": ""}},
+            {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}},
+            {"retriever_two": {"docs": ["doc3", "doc4"]}},
+            {"retriever_one": {"docs": ["doc1", "doc2"]}},
+            {"__interrupt__": ()},
+        ],
+        [
+            {"rewrite_query": {"query": "query: what is weather in sf"}},
+            {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}},
+            {"qa": {"answer": ""}},
+            {"retriever_two": {"docs": ["doc3", "doc4"]}},
+            {"retriever_one": {"docs": ["doc1", "doc2"]}},
+            {"__interrupt__": ()},
+        ],
+    )
 
     assert [c for c in app_w_interrupt.stream(None, config)] == [
         {"qa": {"answer": "doc1,doc2,doc3,doc4"}},
@@ -5969,9 +5983,7 @@ def test_falsy_return_from_task(
                     "a": 5,
                 },
                 "name": "graph",
-                "triggers": [
-                    "__start__",
-                ],
+                "triggers": ("__start__",),
             },
             "step": 0,
             "timestamp": AnyStr(),
@@ -5985,9 +5997,7 @@ def test_falsy_return_from_task(
                     {},
                 ),
                 "name": "falsy_task",
-                "triggers": [
-                    "__pregel_push",
-                ],
+                "triggers": ("__pregel_push",),
             },
             "step": 0,
             "timestamp": AnyStr(),
@@ -6094,9 +6104,7 @@ def test_falsy_return_from_task(
                     "a": 5,
                 },
                 "name": "graph",
-                "triggers": [
-                    "__start__",
-                ],
+                "triggers": ("__start__",),
             },
             "step": 0,
             "timestamp": AnyStr(),
@@ -6110,9 +6118,7 @@ def test_falsy_return_from_task(
                     {},
                 ),
                 "name": "falsy_task",
-                "triggers": [
-                    "__pregel_push",
-                ],
+                "triggers": ("__pregel_push",),
             },
             "step": 0,
             "timestamp": AnyStr(),
@@ -6923,7 +6929,10 @@ def test_tags_stream_mode_messages() -> None:
             {
                 "langgraph_step": 1,
                 "langgraph_node": "call_model",
-                "langgraph_triggers": ["start:call_model"],
+                "langgraph_triggers": (
+                    "branch:to:call_model",
+                    "start:call_model",
+                ),
                 "langgraph_path": ("__pregel_pull", "call_model"),
                 "langgraph_checkpoint_ns": AnyStr("call_model:"),
                 "checkpoint_ns": AnyStr("call_model:"),
@@ -7317,3 +7326,614 @@ def test_empty_invoke() -> None:
         "111": 111,
         "222": 222,
     }
+
+
+@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC)
+def test_parallel_interrupts(
+    request: pytest.FixtureRequest, checkpointer_name: str
+) -> None:
+    from pydantic import BaseModel, Field
+
+    checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}")
+
+    # --- CHILD GRAPH ---
+
+    class ChildState(BaseModel):
+        prompt: str = Field(..., description="What is going to be asked to the user?")
+        human_input: Optional[str] = Field(None, description="What the human said")
+        human_inputs: Annotated[List[str], operator.add] = Field(
+            default_factory=list, description="All of my messages"
+        )
+
+    def get_human_input(state: ChildState):
+        human_input = interrupt(state.prompt)
+
+        return dict(
+            human_input=human_input,  # update child state
+            human_inputs=[human_input],  # update parent state
+        )
+
+    child_graph_builder = StateGraph(ChildState)
+    child_graph_builder.add_node("get_human_input", get_human_input)
+    child_graph_builder.add_edge(START, "get_human_input")
+    child_graph_builder.add_edge("get_human_input", END)
+    child_graph = child_graph_builder.compile()
+
+    # --- PARENT GRAPH ---
+
+    class ParentState(BaseModel):
+        prompts: List[str] = Field(
+            ..., description="What is going to be asked to the user?"
+        )
+        human_inputs: Annotated[List[str], operator.add] = Field(
+            default_factory=list, description="All of my messages"
+        )
+
+    def assign_workers(state: ParentState):
+        return [
+            Send(
+                "child_graph",
+                dict(
+                    prompt=prompt,
+                ),
+            )
+            for prompt in state.prompts
+        ]
+
+    def cleanup(state: ParentState):
+        assert len(state.human_inputs) == len(state.prompts)
+
+    parent_graph_builder = StateGraph(ParentState)
+    parent_graph_builder.add_node("child_graph", child_graph)
+    parent_graph_builder.add_node("cleanup", cleanup)
+
+    parent_graph_builder.add_conditional_edges(START, assign_workers, ["child_graph"])
+    parent_graph_builder.add_edge("child_graph", "cleanup")
+    parent_graph_builder.add_edge("cleanup", END)
+
+    parent_graph = parent_graph_builder.compile(checkpointer=checkpointer)
+
+    # --- CLIENT INVOCATION ---
+
+    thread_config = dict(
+        configurable=dict(
+            thread_id=str(uuid.uuid4()),
+        )
+    )
+    current_input = dict(
+        prompts=["a", "b"],
+    )
+
+    invokes = 0
+    events: dict[int, list[dict]] = {}
+    while invokes < 10:
+        # reset interrupt
+        invokes += 1
+        events[invokes] = []
+        current_interrupts: list[Interrupt] = []
+
+        # start / resume the graph
+        for event in parent_graph.stream(
+            input=current_input,
+            config=thread_config,
+            stream_mode="updates",
+        ):
+            events[invokes].append(event)
+            # handle the interrupt
+            if "__interrupt__" in event:
+                current_interrupts.extend(event["__interrupt__"])
+                # assume that it breaks here, because it is an interrupt
+
+        # get human input and resume
+        if any(i.resumable for i in current_interrupts):
+            current_input = Command(resume=f"Resume #{invokes}")
+
+        # not more human input required, must be completed
+        else:
+            break
+    else:
+        assert False, "Detected infinite loop"
+
+    assert invokes == 3
+    assert len(events) == 3
+
+    assert events[1] == UnsortedSequence(
+        {
+            "__interrupt__": (
+                Interrupt(
+                    value="a",
+                    resumable=True,
+                    ns=[
+                        AnyStr("child_graph:"),
+                        AnyStr("get_human_input:"),
+                    ],
+                ),
+            )
+        },
+        {
+            "__interrupt__": (
+                Interrupt(
+                    value="b",
+                    resumable=True,
+                    ns=[
+                        AnyStr("child_graph:"),
+                        AnyStr("get_human_input:"),
+                    ],
+                ),
+            )
+        },
+    )
+    assert events[2] in (
+        UnsortedSequence(
+            {
+                "__interrupt__": (
+                    Interrupt(
+                        value="a",
+                        resumable=True,
+                        ns=[
+                            AnyStr("child_graph:"),
+                            AnyStr("get_human_input:"),
+                        ],
+                    ),
+                )
+            },
+            {"child_graph": {"human_inputs": ["Resume #1"]}},
+        ),
+        UnsortedSequence(
+            {
+                "__interrupt__": (
+                    Interrupt(
+                        value="b",
+                        resumable=True,
+                        ns=[
+                            AnyStr("child_graph:"),
+                            AnyStr("get_human_input:"),
+                        ],
+                    ),
+                )
+            },
+            {"child_graph": {"human_inputs": ["Resume #1"]}},
+        ),
+    )
+    assert events[3] == UnsortedSequence(
+        {
+            "child_graph": {"human_inputs": ["Resume #1"]},
+            "__metadata__": {"cached": True},
+        },
+        {"child_graph": {"human_inputs": ["Resume #2"]}},
+        {"cleanup": None},
+    )
+
+
+@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC)
+def test_parallel_interrupts_double(
+    request: pytest.FixtureRequest, checkpointer_name: str
+) -> None:
+    from pydantic import BaseModel, Field
+
+    checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}")
+
+    # --- CHILD GRAPH ---
+
+    class ChildState(BaseModel):
+        prompt: str = Field(..., description="What is going to be asked to the user?")
+        human_input: Optional[str] = Field(None, description="What the human said")
+        human_inputs: Annotated[List[str], operator.add] = Field(
+            default_factory=list, description="All of my messages"
+        )
+
+    def get_human_input(state: ChildState):
+        human_input = interrupt(state.prompt)
+
+        return dict(
+            human_inputs=[human_input],  # update parent state
+        )
+
+    def get_dolphin_input(state: ChildState):
+        human_input = interrupt(state.prompt)
+
+        return dict(
+            human_inputs=[human_input],  # update parent state
+        )
+
+    child_graph_builder = StateGraph(ChildState)
+    child_graph_builder.add_node("get_human_input", get_human_input)
+    child_graph_builder.add_node("get_dolphin_input", get_dolphin_input)
+    child_graph_builder.add_edge(START, "get_human_input")
+    child_graph_builder.add_edge(START, "get_dolphin_input")
+    child_graph = child_graph_builder.compile()
+
+    # --- PARENT GRAPH ---
+
+    class ParentState(BaseModel):
+        prompts: List[str] = Field(
+            ..., description="What is going to be asked to the user?"
+        )
+        human_inputs: Annotated[List[str], operator.add] = Field(
+            default_factory=list, description="All of my messages"
+        )
+
+    def assign_workers(state: ParentState):
+        return [
+            Send(
+                "child_graph",
+                dict(
+                    prompt=prompt,
+                ),
+            )
+            for prompt in state.prompts
+        ]
+
+    def cleanup(state: ParentState):
+        assert len(state.human_inputs) == len(state.prompts) * 2
+
+    parent_graph_builder = StateGraph(ParentState)
+    parent_graph_builder.add_node("child_graph", child_graph)
+    parent_graph_builder.add_node("cleanup", cleanup)
+
+    parent_graph_builder.add_conditional_edges(START, assign_workers, ["child_graph"])
+    parent_graph_builder.add_edge("child_graph", "cleanup")
+    parent_graph_builder.add_edge("cleanup", END)
+
+    parent_graph = parent_graph_builder.compile(checkpointer=checkpointer)
+
+    # --- CLIENT INVOCATION ---
+
+    thread_config = dict(
+        configurable=dict(
+            thread_id=str(uuid.uuid4()),
+        )
+    )
+    current_input = dict(
+        prompts=["a", "b"],
+    )
+
+    invokes = 0
+    events: dict[int, list[dict]] = {}
+    while invokes < 10:
+        # reset interrupt
+        invokes += 1
+        events[invokes] = []
+        current_interrupts: list[Interrupt] = []
+
+        # start / resume the graph
+        for event in parent_graph.stream(
+            input=current_input,
+            config=thread_config,
+            stream_mode="updates",
+        ):
+            events[invokes].append(event)
+            # handle the interrupt
+            if "__interrupt__" in event:
+                current_interrupts.extend(event["__interrupt__"])
+                # assume that it breaks here, because it is an interrupt
+
+        # get human input and resume
+        if any(i.resumable for i in current_interrupts):
+            current_input = Command(resume=f"Resume #{invokes}")
+
+        # not more human input required, must be completed
+        else:
+            break
+    else:
+        assert False, "Detected infinite loop"
+
+    assert invokes == 5
+    assert len(events) == 5
+
+
+def test_pregel_loop_refcount():
+    gc.collect()
+    try:
+        gc.disable()
+
+        class State(TypedDict):
+            messages: Annotated[list, add_messages]
+
+        graph_builder = StateGraph(State)
+
+        def chatbot(state: State):
+            return {"messages": [("ai", "HIYA")]}
+
+        graph_builder.add_node("chatbot", chatbot)
+        graph_builder.set_entry_point("chatbot")
+        graph_builder.set_finish_point("chatbot")
+        graph = graph_builder.compile()
+
+        for _ in range(5):
+            graph.invoke({"messages": [{"role": "user", "content": "hi"}]})
+            assert (
+                len(
+                    [obj for obj in gc.get_objects() if isinstance(obj, SyncPregelLoop)]
+                )
+                == 0
+            )
+            assert (
+                len([obj for obj in gc.get_objects() if isinstance(obj, PregelRunner)])
+                == 0
+            )
+    finally:
+        gc.enable()
+
+
+@pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_SYNC)
+def test_bulk_state_updates(
+    request: pytest.FixtureRequest, checkpointer_name: str
+) -> None:
+    checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}")
+
+    class State(TypedDict):
+        foo: str
+        baz: str
+
+    def node_a(state: State) -> State:
+        return {"foo": "bar"}
+
+    def node_b(state: State) -> State:
+        return {"baz": "qux"}
+
+    graph = (
+        StateGraph(State)
+        .add_node("node_a", node_a)
+        .add_node("node_b", node_b)
+        .add_edge(START, "node_a")
+        .add_edge("node_a", "node_b")
+        .compile(checkpointer=checkpointer)
+    )
+
+    config = {"configurable": {"thread_id": "1"}}
+
+    # First update with node_a
+    graph.bulk_update_state(
+        config,
+        [
+            [
+                StateUpdate(values={"foo": "bar"}, as_node="node_a"),
+            ]
+        ],
+    )
+
+    # Then bulk update with both nodes
+    graph.bulk_update_state(
+        config,
+        [
+            [
+                StateUpdate(values={"foo": "updated"}, as_node="node_a"),
+                StateUpdate(values={"baz": "new"}, as_node="node_b"),
+            ]
+        ],
+    )
+
+    state = graph.get_state(config)
+    assert state.values == {"foo": "updated", "baz": "new"}
+
+    # Check if there are only two checkpoints
+    checkpoints = list(checkpointer.list(config))
+    assert len(checkpoints) == 2
+    assert checkpoints[0].metadata["writes"] == {
+        "node_a": {"foo": "updated"},
+        "node_b": {"baz": "new"},
+    }
+    assert checkpoints[1].metadata["writes"] == {"node_a": {"foo": "bar"}}
+
+    # perform multiple steps at the same time
+    config = {"configurable": {"thread_id": "2"}}
+
+    graph.bulk_update_state(
+        config,
+        [
+            [
+                StateUpdate(values={"foo": "bar"}, as_node="node_a"),
+            ],
+            [
+                StateUpdate(values={"foo": "updated"}, as_node="node_a"),
+                StateUpdate(values={"baz": "new"}, as_node="node_b"),
+            ],
+        ],
+    )
+
+    state = graph.get_state(config)
+    assert state.values == {"foo": "updated", "baz": "new"}
+
+    checkpoints = list(checkpointer.list(config))
+    assert len(checkpoints) == 2
+    assert checkpoints[0].metadata["writes"] == {
+        "node_a": {"foo": "updated"},
+        "node_b": {"baz": "new"},
+    }
+    assert checkpoints[1].metadata["writes"] == {"node_a": {"foo": "bar"}}
+
+    # Should raise error if updating without as_node
+    with pytest.raises(InvalidUpdateError):
+        graph.bulk_update_state(
+            config,
+            [
+                [
+                    StateUpdate(values={"foo": "error"}, as_node=None),
+                    StateUpdate(values={"bar": "error"}, as_node=None),
+                ]
+            ],
+        )
+
+    # Should raise if no updates are provided
+    with pytest.raises(ValueError, match="No supersteps provided"):
+        graph.bulk_update_state(config, [])
+
+    # Should raise if no updates are provided
+    with pytest.raises(ValueError, match="No updates provided"):
+        graph.bulk_update_state(config, [[], []])
+
+    # Should raise if __end__ or __copy__ update is applied in bulk
+    with pytest.raises(InvalidUpdateError):
+        graph.bulk_update_state(
+            config,
+            [
+                [
+                    StateUpdate(values=None, as_node="__end__"),
+                    StateUpdate(values=None, as_node="__copy__"),
+                ],
+            ],
+        )
+
+
+@pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_SYNC)
+def test_update_as_input(
+    request: pytest.FixtureRequest, checkpointer_name: str
+) -> None:
+    checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}")
+
+    class State(TypedDict):
+        foo: str
+
+    def agent(state: State) -> State:
+        return {"foo": "agent"}
+
+    def tool(state: State) -> State:
+        return {"foo": "tool"}
+
+    graph = (
+        StateGraph(State)
+        .add_node("agent", agent)
+        .add_node("tool", tool)
+        .add_edge(START, "agent")
+        .add_edge("agent", "tool")
+        .compile(checkpointer=checkpointer)
+    )
+
+    assert graph.invoke({"foo": "input"}, {"configurable": {"thread_id": "1"}}) == {
+        "foo": "tool"
+    }
+
+    assert graph.invoke({"foo": "input"}, {"configurable": {"thread_id": "1"}}) == {
+        "foo": "tool"
+    }
+
+    def map_snapshot(i: StateSnapshot) -> dict:
+        return {
+            "values": i.values,
+            "next": i.next,
+            "step": i.metadata.get("step"),
+        }
+
+    history = [
+        map_snapshot(s)
+        for s in graph.get_state_history({"configurable": {"thread_id": "1"}})
+    ]
+
+    graph.bulk_update_state(
+        {"configurable": {"thread_id": "2"}},
+        [
+            # First turn
+            [StateUpdate({"foo": "input"}, "__input__")],
+            [StateUpdate({"foo": "input"}, "__start__")],
+            [StateUpdate({"foo": "agent"}, "agent")],
+            [StateUpdate({"foo": "tool"}, "tool")],
+            # Second turn
+            [StateUpdate({"foo": "input"}, "__input__")],
+            [StateUpdate({"foo": "input"}, "__start__")],
+            [StateUpdate({"foo": "agent"}, "agent")],
+            [StateUpdate({"foo": "tool"}, "tool")],
+        ],
+    )
+
+    state = graph.get_state({"configurable": {"thread_id": "2"}})
+    assert state.values == {"foo": "tool"}
+
+    new_history = [
+        map_snapshot(s)
+        for s in graph.get_state_history({"configurable": {"thread_id": "2"}})
+    ]
+
+    assert new_history == history
+
+
+@pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_SYNC)
+def test_batch_update_as_input(
+    request: pytest.FixtureRequest, checkpointer_name: str
+) -> None:
+    checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}")
+
+    class State(TypedDict):
+        foo: str
+        tasks: Annotated[list[int], operator.add]
+
+    def agent(state: State) -> State:
+        return {"foo": "agent"}
+
+    def map(state: State) -> Command["task"]:
+        return Command(
+            goto=[
+                Send("task", {"index": 0}),
+                Send("task", {"index": 1}),
+                Send("task", {"index": 2}),
+            ],
+            update={"foo": "map"},
+        )
+
+    def task(state: dict) -> State:
+        return {"tasks": [state["index"]]}
+
+    graph = (
+        StateGraph(State)
+        .add_node("agent", agent)
+        .add_node("map", map)
+        .add_node("task", task)
+        .add_edge(START, "agent")
+        .add_edge("agent", "map")
+        .compile(checkpointer=checkpointer)
+    )
+
+    assert graph.invoke({"foo": "input"}, {"configurable": {"thread_id": "1"}}) == {
+        "foo": "map",
+        "tasks": [0, 1, 2],
+    }
+
+    def map_snapshot(i: StateSnapshot) -> dict:
+        return {
+            "values": i.values,
+            "next": i.next,
+            "step": i.metadata.get("step"),
+            "tasks": [t.name for t in i.tasks],
+        }
+
+    history = [
+        map_snapshot(s)
+        for s in graph.get_state_history({"configurable": {"thread_id": "1"}})
+    ]
+
+    graph.bulk_update_state(
+        {"configurable": {"thread_id": "2"}},
+        [
+            [StateUpdate({"foo": "input"}, "__input__")],
+            [StateUpdate({"foo": "input"}, "__start__")],
+            [StateUpdate({"foo": "agent", "tasks": []}, "agent")],
+            [
+                StateUpdate(
+                    Command(
+                        goto=[
+                            Send("task", {"index": 0}),
+                            Send("task", {"index": 1}),
+                            Send("task", {"index": 2}),
+                        ],
+                        update={"foo": "map"},
+                    ),
+                    "map",
+                )
+            ],
+            [
+                StateUpdate({"tasks": [0]}, "task"),
+                StateUpdate({"tasks": [1]}, "task"),
+                StateUpdate({"tasks": [2]}, "task"),
+            ],
+        ],
+    )
+
+    state = graph.get_state({"configurable": {"thread_id": "2"}})
+    assert state.values == {"foo": "map", "tasks": [0, 1, 2]}
+
+    new_history = [
+        map_snapshot(s)
+        for s in graph.get_state_history({"configurable": {"thread_id": "2"}})
+    ]
+
+    assert new_history == history

libs/langgraph/tests/test_pregel_async.py

@@ -1,5 +1,6 @@
 import asyncio
 import functools
+import gc
 import logging
 import operator
 import random
@@ -27,11 +28,7 @@ from uuid import UUID
 import httpx
 import pytest
 from langchain_core.language_models import GenericFakeChatModel
-from langchain_core.runnables import (
-    RunnableConfig,
-    RunnableLambda,
-    RunnablePassthrough,
-)
+from langchain_core.runnables import RunnableConfig, RunnableLambda, RunnablePassthrough
 from langchain_core.utils.aiter import aclosing
 from pytest_mock import MockerFixture
 from syrupy import SnapshotAssertion
@@ -56,13 +53,16 @@ from langgraph.graph import END, Graph, StateGraph
 from langgraph.graph.message import MessagesState, add_messages
 from langgraph.prebuilt.tool_node import ToolNode
 from langgraph.pregel import Channel, GraphRecursionError, Pregel, StateSnapshot
+from langgraph.pregel.loop import AsyncPregelLoop
 from langgraph.pregel.retry import RetryPolicy
+from langgraph.pregel.runner import PregelRunner
 from langgraph.store.base import BaseStore
 from langgraph.types import (
     Command,
     Interrupt,
     PregelTask,
     Send,
+    StateUpdate,
     StreamWriter,
     interrupt,
 )
@@ -938,10 +938,7 @@ async def test_copy_checkpoint(checkpointer_name: str) -> None:
             async for c in tool_two.astream(
                 {"my_key": "value ⛰️", "market": "DE"}, thread2
             )
-        ] == [
-            {
-                "tool_one": {"my_key": " one"},
-            },
+        ] == UnsortedSequence(
             {
                 "__interrupt__": (
                     Interrupt(
@@ -951,7 +948,10 @@ async def test_copy_checkpoint(checkpointer_name: str) -> None:
                     ),
                 )
             },
-        ]
+            {
+                "tool_one": {"my_key": " one"},
+            },
+        )
         # resume with answer
         assert [
             c async for c in tool_two.astream(Command(resume=" my answer"), thread2)
@@ -1672,7 +1672,7 @@ async def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None:
                 "id": AnyStr(),
                 "name": "one",
                 "input": 2,
-                "triggers": ["input"],
+                "triggers": ("input",),
             },
         },
         {
@@ -1683,7 +1683,7 @@ async def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None:
                 "id": AnyStr(),
                 "name": "two",
                 "input": [12],
-                "triggers": ["inbox"],
+                "triggers": ("inbox",),
             },
         },
         {
@@ -1718,7 +1718,7 @@ async def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None:
                 "id": AnyStr(),
                 "name": "two",
                 "input": [3],
-                "triggers": ["inbox"],
+                "triggers": ("inbox",),
             },
         },
         {
@@ -4524,6 +4524,7 @@ async def test_nested_pydantic_models(version: str) -> None:
     class NestedModel(BaseModel):
         value: int
         name: str
+        something: Optional[str] = None
 
     # Forward reference model
     class RecursiveModel(BaseModel):
@@ -4545,18 +4546,27 @@ async def test_nested_pydantic_models(version: str) -> None:
         name: str
         friends: list[str] = Field(default_factory=list)  # IDs of friends
 
+    class MyTypedDict(TypedDict):
+        x: int
+
     class State(BaseModel):
         # Basic nested model tests
         top_level: str
         nested: NestedModel
         optional_nested: Optional[NestedModel] = None
         dict_nested: dict[str, NestedModel]
+        my_set: set[int]
         list_nested: Annotated[
             Union[dict, list[dict[str, NestedModel]]], lambda x, y: (x or []) + [y]
         ]
+        list_nested_reversed: Annotated[
+            Union[list[dict[str, NestedModel]], NestedModel, dict, list],
+            lambda x, y: (x or []) + [y],
+        ]
         tuple_nested: tuple[str, NestedModel]
         tuple_list_nested: list[tuple[int, NestedModel]]
         complex_tuple: tuple[str, dict[str, tuple[int, NestedModel]]]
+        my_typed_dict: MyTypedDict
 
         # Forward reference test
         recursive: RecursiveModel
@@ -4572,8 +4582,11 @@ async def test_nested_pydantic_models(version: str) -> None:
         "top_level": "initial",
         "nested": {"value": 42, "name": "test"},
         "optional_nested": {"value": 10, "name": "optional"},
+        "my_set": [1, 2, 7],
+        "my_typed_dict": {"x": 1},
         "dict_nested": {"a": {"value": 5, "name": "a"}},
         "list_nested": [{"a": {"value": 6, "name": "b"}}],
+        "list_nested_reversed": ["foo", "bar"],
         "tuple_nested": ["tuple-key", {"value": 7, "name": "tuple-value"}],
         "tuple_list_nested": [[1, {"value": 8, "name": "tuple-in-list"}]],
         "complex_tuple": [
@@ -5882,10 +5895,12 @@ async def test_store_injected_async(checkpointer_name: str, store_name: str) ->
         ):
             assert isinstance(store, BaseStore)
             await store.aput(
-                namespace
-                if self.i is not None
-                and config["configurable"]["thread_id"] in (thread_1, thread_2)
-                else (f"foo_{self.i}", "bar"),
+                (
+                    namespace
+                    if self.i is not None
+                    and config["configurable"]["thread_id"] in (thread_1, thread_2)
+                    else (f"foo_{self.i}", "bar")
+                ),
                 doc_id,
                 {
                     **doc,
@@ -7571,7 +7586,10 @@ async def test_tags_stream_mode_messages() -> None:
             {
                 "langgraph_step": 1,
                 "langgraph_node": "call_model",
-                "langgraph_triggers": ["start:call_model"],
+                "langgraph_triggers": (
+                    "branch:to:call_model",
+                    "start:call_model",
+                ),
                 "langgraph_path": ("__pregel_pull", "call_model"),
                 "langgraph_checkpoint_ns": AnyStr("call_model:"),
                 "checkpoint_ns": AnyStr("call_model:"),
@@ -7827,3 +7845,329 @@ async def test_handles_multiple_interrupts_from_tasks() -> None:
     assert len(result) == 2
     assert result[0] == "Added James!"
     assert result[1] == "Added Will!"
+
+
+async def test_pregel_loop_refcount():
+    gc.collect()
+    try:
+        gc.disable()
+
+        class State(TypedDict):
+            messages: Annotated[list, add_messages]
+
+        graph_builder = StateGraph(State)
+
+        async def chatbot(state: State):
+            return {"messages": [("ai", "HIYA")]}
+
+        graph_builder.add_node("chatbot", chatbot)
+        graph_builder.set_entry_point("chatbot")
+        graph_builder.set_finish_point("chatbot")
+        graph = graph_builder.compile()
+
+        for _ in range(5):
+            await graph.ainvoke({"messages": [{"role": "user", "content": "hi"}]})
+            assert (
+                len(
+                    [
+                        obj
+                        for obj in gc.get_objects()
+                        if isinstance(obj, AsyncPregelLoop)
+                    ]
+                )
+                == 0
+            )
+            assert (
+                len([obj for obj in gc.get_objects() if isinstance(obj, PregelRunner)])
+                == 0
+            )
+    finally:
+        gc.enable()
+
+
+@pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_ASYNC)
+async def test_bulk_state_updates(checkpointer_name: str) -> None:
+    async with awith_checkpointer(checkpointer_name) as checkpointer:
+
+        class State(TypedDict):
+            foo: str
+            baz: str
+
+        def node_a(state: State) -> State:
+            return {"foo": "bar"}
+
+        def node_b(state: State) -> State:
+            return {"baz": "qux"}
+
+        graph = (
+            StateGraph(State)
+            .add_node("node_a", node_a)
+            .add_node("node_b", node_b)
+            .add_edge(START, "node_a")
+            .add_edge("node_a", "node_b")
+            .compile(checkpointer=checkpointer)
+        )
+
+        config = {"configurable": {"thread_id": "1"}}
+
+        # First update with node_a
+        await graph.abulk_update_state(
+            config,
+            [
+                [
+                    StateUpdate({"foo": "bar"}, "node_a"),
+                ]
+            ],
+        )
+
+        # Then bulk update with both nodes
+        await graph.abulk_update_state(
+            config,
+            [
+                [
+                    StateUpdate({"foo": "updated"}, "node_a"),
+                    StateUpdate({"baz": "new"}, "node_b"),
+                ]
+            ],
+        )
+
+        state = await graph.aget_state(config)
+        assert state.values == {"foo": "updated", "baz": "new"}
+
+        # Check if there are only two checkpoints
+        checkpoints = [
+            c async for c in checkpointer.alist({"configurable": {"thread_id": "1"}})
+        ]
+        assert len(checkpoints) == 2
+        assert checkpoints[0].metadata["writes"] == {
+            "node_a": {"foo": "updated"},
+            "node_b": {"baz": "new"},
+        }
+        assert checkpoints[1].metadata["writes"] == {"node_a": {"foo": "bar"}}
+
+        # perform multiple steps at the same time
+        config = {"configurable": {"thread_id": "2"}}
+
+        await graph.abulk_update_state(
+            config,
+            [
+                [
+                    StateUpdate({"foo": "bar"}, "node_a"),
+                ],
+                [
+                    StateUpdate({"foo": "updated"}, "node_a"),
+                    StateUpdate({"baz": "new"}, "node_b"),
+                ],
+            ],
+        )
+
+        state = await graph.aget_state(config)
+        assert state.values == {"foo": "updated", "baz": "new"}
+
+        checkpoints = [
+            c async for c in checkpointer.alist({"configurable": {"thread_id": "1"}})
+        ]
+        assert len(checkpoints) == 2
+        assert checkpoints[0].metadata["writes"] == {
+            "node_a": {"foo": "updated"},
+            "node_b": {"baz": "new"},
+        }
+        assert checkpoints[1].metadata["writes"] == {"node_a": {"foo": "bar"}}
+
+        # Should raise error if updating without as_node
+        with pytest.raises(InvalidUpdateError):
+            await graph.abulk_update_state(
+                config,
+                [
+                    [
+                        StateUpdate(values={"foo": "error"}, as_node=None),
+                        StateUpdate(values={"bar": "error"}, as_node=None),
+                    ]
+                ],
+            )
+
+        # Should raise if no updates are provided
+        with pytest.raises(ValueError, match="No supersteps provided"):
+            await graph.abulk_update_state(config, [])
+
+        # Should raise if no updates are provided
+        with pytest.raises(ValueError, match="No updates provided"):
+            await graph.abulk_update_state(config, [[], []])
+
+        # Should raise if __end__ or __copy__ update is applied in bulk
+        with pytest.raises(InvalidUpdateError):
+            await graph.abulk_update_state(
+                config,
+                [
+                    [
+                        StateUpdate(values=None, as_node="__end__"),
+                        StateUpdate(values=None, as_node="__copy__"),
+                    ],
+                ],
+            )
+
+
+@pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_ASYNC)
+async def test_update_as_input(checkpointer_name: str) -> None:
+    async with awith_checkpointer(checkpointer_name) as checkpointer:
+
+        class State(TypedDict):
+            foo: str
+
+        def agent(state: State) -> State:
+            return {"foo": "agent"}
+
+        def tool(state: State) -> State:
+            return {"foo": "tool"}
+
+        graph = (
+            StateGraph(State)
+            .add_node("agent", agent)
+            .add_node("tool", tool)
+            .add_edge(START, "agent")
+            .add_edge("agent", "tool")
+            .compile(checkpointer=checkpointer)
+        )
+
+        assert await graph.ainvoke(
+            {"foo": "input"}, {"configurable": {"thread_id": "1"}}
+        ) == {"foo": "tool"}
+
+        assert await graph.ainvoke(
+            {"foo": "input"}, {"configurable": {"thread_id": "1"}}
+        ) == {"foo": "tool"}
+
+        def map_snapshot(i: StateSnapshot) -> dict:
+            return {
+                "values": i.values,
+                "next": i.next,
+                "step": i.metadata.get("step"),
+            }
+
+        history = [
+            map_snapshot(s)
+            async for s in graph.aget_state_history(
+                {"configurable": {"thread_id": "1"}}
+            )
+        ]
+
+        await graph.abulk_update_state(
+            {"configurable": {"thread_id": "2"}},
+            [
+                # First turn
+                [StateUpdate({"foo": "input"}, "__input__")],
+                [StateUpdate({"foo": "input"}, "__start__")],
+                [StateUpdate({"foo": "agent"}, "agent")],
+                [StateUpdate({"foo": "tool"}, "tool")],
+                # Second turn
+                [StateUpdate({"foo": "input"}, "__input__")],
+                [StateUpdate({"foo": "input"}, "__start__")],
+                [StateUpdate({"foo": "agent"}, "agent")],
+                [StateUpdate({"foo": "tool"}, "tool")],
+            ],
+        )
+
+        state = await graph.aget_state({"configurable": {"thread_id": "2"}})
+        assert state.values == {"foo": "tool"}
+
+        new_history = [
+            map_snapshot(s)
+            async for s in graph.aget_state_history(
+                {"configurable": {"thread_id": "2"}}
+            )
+        ]
+
+        assert new_history == history
+
+
+@pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_ASYNC)
+async def test_batch_update_as_input(checkpointer_name: str) -> None:
+    async with awith_checkpointer(checkpointer_name) as checkpointer:
+
+        class State(TypedDict):
+            foo: str
+            tasks: Annotated[list[int], operator.add]
+
+        def agent(state: State) -> State:
+            return {"foo": "agent"}
+
+        def map(state: State) -> Command["task"]:
+            return Command(
+                goto=[
+                    Send("task", {"index": 0}),
+                    Send("task", {"index": 1}),
+                    Send("task", {"index": 2}),
+                ],
+                update={"foo": "map"},
+            )
+
+        def task(state: dict) -> State:
+            return {"tasks": [state["index"]]}
+
+        graph = (
+            StateGraph(State)
+            .add_node("agent", agent)
+            .add_node("map", map)
+            .add_node("task", task)
+            .add_edge(START, "agent")
+            .add_edge("agent", "map")
+            .compile(checkpointer=checkpointer)
+        )
+
+        assert await graph.ainvoke(
+            {"foo": "input"}, {"configurable": {"thread_id": "1"}}
+        ) == {"foo": "map", "tasks": [0, 1, 2]}
+
+        def map_snapshot(i: StateSnapshot) -> dict:
+            return {
+                "values": i.values,
+                "next": i.next,
+                "step": i.metadata.get("step"),
+                "tasks": [t.name for t in i.tasks],
+            }
+
+        history = [
+            map_snapshot(s)
+            async for s in graph.aget_state_history(
+                {"configurable": {"thread_id": "1"}}
+            )
+        ]
+
+        await graph.abulk_update_state(
+            {"configurable": {"thread_id": "2"}},
+            [
+                [StateUpdate({"foo": "input"}, "__input__")],
+                [StateUpdate({"foo": "input"}, "__start__")],
+                [StateUpdate({"foo": "agent", "tasks": []}, "agent")],
+                [
+                    StateUpdate(
+                        Command(
+                            goto=[
+                                Send("task", {"index": 0}),
+                                Send("task", {"index": 1}),
+                                Send("task", {"index": 2}),
+                            ],
+                            update={"foo": "map"},
+                        ),
+                        "map",
+                    )
+                ],
+                [
+                    StateUpdate({"tasks": [0]}, "task"),
+                    StateUpdate({"tasks": [1]}, "task"),
+                    StateUpdate({"tasks": [2]}, "task"),
+                ],
+            ],
+        )
+
+        state = await graph.aget_state({"configurable": {"thread_id": "2"}})
+        assert state.values == {"foo": "map", "tasks": [0, 1, 2]}
+
+        new_history = [
+            map_snapshot(s)
+            async for s in graph.aget_state_history(
+                {"configurable": {"thread_id": "2"}}
+            )
+        ]
+
+        assert new_history == history

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

@@ -1,5 +1,7 @@
 import asyncio
+import binascii
 import concurrent.futures
+import weakref
 from collections.abc import Sequence
 from contextlib import (
     AbstractAsyncContextManager,
@@ -19,7 +21,7 @@ import langgraph.scheduler.kafka.serde as serde
 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
+from langgraph.pregel.algo import checkpoint_null_version, prepare_single_task
 from langgraph.pregel.executor import (
     AsyncBackgroundExecutor,
     BackgroundExecutor,
@@ -209,12 +211,17 @@ class AsyncKafkaExecutor(AbstractAsyncContextManager):
                 for_execution=True,
                 checkpointer=self.graph.checkpointer,
                 store=self.graph.store,
+                checkpoint_id_bytes=binascii.unhexlify(
+                    saved.checkpoint["id"].replace("-", "")
+                ),
+                checkpoint_null_version=checkpoint_null_version(saved.checkpoint),
             ):
                 # execute task, saving writes
+                put_writes = partial(self._put_writes, submit, msg["config"])
                 runner = PregelRunner(
-                    submit=submit,
-                    put_writes=partial(self._put_writes, submit, msg["config"]),
-                    schedule_task=self._schedule_task,
+                    submit=weakref.ref(submit),
+                    put_writes=weakref.ref(put_writes),
+                    schedule_task=weakref.WeakMethod(self._schedule_task),
                 )
                 async for _ in runner.atick([task], reraise=False):
                     pass
@@ -421,12 +428,17 @@ class KafkaExecutor(AbstractContextManager):
                 step=saved.metadata["step"] + 1,
                 for_execution=True,
                 checkpointer=self.graph.checkpointer,
+                checkpoint_id_bytes=binascii.unhexlify(
+                    saved.checkpoint["id"].replace("-", "")
+                ),
+                checkpoint_null_version=checkpoint_null_version(saved.checkpoint),
             ):
                 # execute task, saving writes
+                put_writes = partial(self._put_writes, submit, msg["config"])
                 runner = PregelRunner(
-                    submit=submit,
-                    put_writes=partial(self._put_writes, submit, msg["config"]),
-                    schedule_task=self._schedule_task,
+                    submit=weakref.ref(submit),
+                    put_writes=weakref.ref(put_writes),
+                    schedule_task=weakref.WeakMethod(self._schedule_task),
                 )
                 for _ in runner.tick([task], reraise=False):
                     pass

libs/scheduler-kafka/tests/test_subgraph.py

@@ -202,7 +202,7 @@ async def test_subgraph_w_interrupt(
                                 "subgraph_counter": None,
                                 "call_counter": None,
                                 "interrupt_counter": None,
-                                "null_resume": None,
+                                "get_null_resume": None,
                                 "resume": [],
                             },
                             "checkpoint_id": None,
@@ -275,7 +275,7 @@ async def test_subgraph_w_interrupt(
                                 "subgraph_counter": None,
                                 "call_counter": None,
                                 "interrupt_counter": None,
-                                "null_resume": None,
+                                "get_null_resume": None,
                                 "resume": [],
                             },
                             "checkpoint_id": c.config["configurable"]["checkpoint_id"],
@@ -378,7 +378,7 @@ async def test_subgraph_w_interrupt(
                                 "subgraph_counter": None,
                                 "call_counter": None,
                                 "interrupt_counter": None,
-                                "null_resume": None,
+                                "get_null_resume": None,
                                 "resume": [],
                             },
                             "checkpoint_id": c.config["configurable"]["checkpoint_id"],
@@ -491,7 +491,7 @@ async def test_subgraph_w_interrupt(
                                 "subgraph_counter": None,
                                 "call_counter": None,
                                 "interrupt_counter": None,
-                                "null_resume": None,
+                                "get_null_resume": None,
                                 "resume": [],
                             },
                             "checkpoint_id": None,
@@ -559,7 +559,7 @@ async def test_subgraph_w_interrupt(
                                 "subgraph_counter": None,
                                 "call_counter": None,
                                 "interrupt_counter": None,
-                                "null_resume": None,
+                                "get_null_resume": None,
                                 "resume": [],
                             },
                             "checkpoint_id": c.config["configurable"]["checkpoint_id"],
@@ -683,7 +683,7 @@ async def test_subgraph_w_interrupt(
                                 "subgraph_counter": None,
                                 "call_counter": None,
                                 "interrupt_counter": None,
-                                "null_resume": None,
+                                "get_null_resume": None,
                                 "resume": [],
                             },
                             "checkpoint_id": c.config["configurable"]["checkpoint_id"],

libs/scheduler-kafka/tests/test_subgraph_sync.py

@@ -201,7 +201,7 @@ def test_subgraph_w_interrupt(
                                 "subgraph_counter": None,
                                 "call_counter": None,
                                 "interrupt_counter": None,
-                                "null_resume": None,
+                                "get_null_resume": None,
                                 "resume": [],
                             },
                             "checkpoint_id": None,
@@ -274,7 +274,7 @@ def test_subgraph_w_interrupt(
                                 "subgraph_counter": None,
                                 "call_counter": None,
                                 "interrupt_counter": None,
-                                "null_resume": None,
+                                "get_null_resume": None,
                                 "resume": [],
                             },
                             "checkpoint_id": c.config["configurable"]["checkpoint_id"],
@@ -377,7 +377,7 @@ def test_subgraph_w_interrupt(
                                 "subgraph_counter": None,
                                 "call_counter": None,
                                 "interrupt_counter": None,
-                                "null_resume": None,
+                                "get_null_resume": None,
                                 "resume": [],
                             },
                             "checkpoint_id": c.config["configurable"]["checkpoint_id"],
@@ -489,7 +489,7 @@ def test_subgraph_w_interrupt(
                                 "subgraph_counter": None,
                                 "call_counter": None,
                                 "interrupt_counter": None,
-                                "null_resume": None,
+                                "get_null_resume": None,
                                 "resume": [],
                             },
                             "checkpoint_id": None,
@@ -557,7 +557,7 @@ def test_subgraph_w_interrupt(
                                 "subgraph_counter": None,
                                 "call_counter": None,
                                 "interrupt_counter": None,
-                                "null_resume": None,
+                                "get_null_resume": None,
                                 "resume": [],
                             },
                             "checkpoint_id": c.config["configurable"]["checkpoint_id"],
@@ -681,7 +681,7 @@ def test_subgraph_w_interrupt(
                                 "subgraph_counter": None,
                                 "call_counter": None,
                                 "interrupt_counter": None,
-                                "null_resume": None,
+                                "get_null_resume": None,
                                 "resume": [],
                             },
                             "checkpoint_id": c.config["configurable"]["checkpoint_id"],

libs/sdk-js/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@langchain/langgraph-sdk",
-  "version": "0.0.57",
+  "version": "0.0.60",
   "description": "Client library for interacting with the LangGraph API",
   "type": "module",
   "packageManager": "yarn@1.22.19",

libs/sdk-js/src/client.ts

@@ -30,6 +30,7 @@ import type {
   StreamEvent,
   CronsCreatePayload,
   OnConflictBehavior,
+  Command,
 } from "./types.js";
 import { mergeSignals } from "./utils/signals.js";
 import { getEnvironmentVariable } from "./utils/env.js";
@@ -481,15 +482,47 @@ export class ThreadsClient<
      * Metadata for the thread.
      */
     metadata?: Metadata;
+    /**
+     * ID of the thread to create.
+     *
+     * If not provided, a random UUID will be generated.
+     */
     threadId?: string;
+    /**
+     * How to handle duplicate creation.
+     *
+     * @default "raise"
+     */
     ifExists?: OnConflictBehavior;
+    /**
+     * Graph ID to associate with the thread.
+     */
+    graphId?: string;
+    /**
+     * Apply a list of supersteps when creating a thread, each containing a sequence of updates.
+     *
+     * Used for copying a thread between deployments.
+     */
+    supersteps?: Array<{
+      updates: Array<{ values: unknown; command?: Command; asNode: string }>;
+    }>;
   }): Promise<Thread<TStateType>> {
     return this.fetch<Thread<TStateType>>(`/threads`, {
       method: "POST",
       json: {
-        metadata: payload?.metadata,
+        metadata: {
+          ...payload?.metadata,
+          graph_id: payload?.graphId,
+        },
         thread_id: payload?.threadId,
         if_exists: payload?.ifExists,
+        supersteps: payload?.supersteps?.map((s) => ({
+          updates: s.updates.map((u) => ({
+            values: u.values,
+            command: u.command,
+            as_node: u.asNode,
+          })),
+        })),
       },
     });
   }

libs/sdk-py/langgraph_sdk/client.py

@@ -839,6 +839,8 @@ class ThreadsClient:
         metadata: Json = None,
         thread_id: Optional[str] = None,
         if_exists: Optional[OnConflictBehavior] = None,
+        supersteps: Optional[Sequence[dict[str, Sequence[dict[str, Any]]]]] = None,
+        graph_id: Optional[str] = None,
     ) -> Thread:
         """Create a new thread.
 
@@ -848,6 +850,9 @@ class ThreadsClient:
                 If None, ID will be a randomly generated UUID.
             if_exists: How to handle duplicate creation. Defaults to 'raise' under the hood.
                 Must be either 'raise' (raise error if duplicate), or 'do_nothing' (return existing thread).
+            supersteps: Apply a list of supersteps when creating a thread, each containing a sequence of updates.
+                Each update has `values` or `command` and `as_node`. Used for copying a thread between deployments.
+            graph_id: Optional graph ID to associate with the thread.
 
         Returns:
             Thread: The created thread.
@@ -863,10 +868,28 @@ class ThreadsClient:
         payload: Dict[str, Any] = {}
         if thread_id:
             payload["thread_id"] = thread_id
-        if metadata:
-            payload["metadata"] = metadata
+        if metadata or graph_id:
+            payload["metadata"] = {
+                **(metadata or {}),
+                **({"graph_id": graph_id} if graph_id else {}),
+            }
         if if_exists:
             payload["if_exists"] = if_exists
+        if supersteps:
+            payload["supersteps"] = [
+                {
+                    "updates": [
+                        {
+                            "values": u["values"],
+                            "command": u.get("command"),
+                            "as_node": u["as_node"],
+                        }
+                        for u in s["updates"]
+                    ]
+                }
+                for s in supersteps
+            ]
+
         return await self.http.post("/threads", json=payload)
 
     async def update(self, thread_id: str, *, metadata: dict[str, Any]) -> Thread:
@@ -3036,6 +3059,8 @@ class SyncThreadsClient:
         metadata: Json = None,
         thread_id: Optional[str] = None,
         if_exists: Optional[OnConflictBehavior] = None,
+        supersteps: Optional[Sequence[dict[str, Sequence[dict[str, Any]]]]] = None,
+        graph_id: Optional[str] = None,
     ) -> Thread:
         """Create a new thread.
 
@@ -3045,6 +3070,9 @@ class SyncThreadsClient:
                 If None, ID will be a randomly generated UUID.
             if_exists: How to handle duplicate creation. Defaults to 'raise' under the hood.
                 Must be either 'raise' (raise error if duplicate), or 'do_nothing' (return existing thread).
+            supersteps: Apply a list of supersteps when creating a thread, each containing a sequence of updates.
+                Each update has `values` or `command` and `as_node`. Used for copying a thread between deployments.
+            graph_id: Optional graph ID to associate with the thread.
 
         Returns:
             Thread: The created thread.
@@ -3060,10 +3088,28 @@ class SyncThreadsClient:
         payload: Dict[str, Any] = {}
         if thread_id:
             payload["thread_id"] = thread_id
-        if metadata:
-            payload["metadata"] = metadata
+        if metadata or graph_id:
+            payload["metadata"] = {
+                **(metadata or {}),
+                **({"graph_id": graph_id} if graph_id else {}),
+            }
         if if_exists:
             payload["if_exists"] = if_exists
+        if supersteps:
+            payload["supersteps"] = [
+                {
+                    "updates": [
+                        {
+                            "values": u["values"],
+                            "command": u.get("command"),
+                            "as_node": u["as_node"],
+                        }
+                        for u in s["updates"]
+                    ]
+                }
+                for s in supersteps
+            ]
+
         return self.http.post("/threads", json=payload)
 
     def update(self, thread_id: str, *, metadata: dict[str, Any]) -> Thread:
@@ -3307,7 +3353,7 @@ class SyncThreadsClient:
 
         Example Usage:
 
-            response = client.threads.update_state(
+            response = await client.threads.update_state(
                 thread_id="my_thread_id",
                 values={"messages":[{"role": "user", "content": "hello!"}]},
                 as_node="my_node",

libs/sdk-py/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "langgraph-sdk"
-version = "0.1.57"
+version = "0.1.58"
 description = "SDK for interacting with LangGraph API"
 authors = []
 license = "MIT"