Merge branch 'main' into wfh/gc_test_all

Includes:
- Explicit unsetting of runnable context var
- Weakref for PregelExecutableTask

both to reduce the chance of keeping a reference to an internal object
and preventing garbage collection

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/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/how-tos/state-model.ipynb

@@ -463,12 +463,12 @@
    "source": [
     "from langgraph.graph import StateGraph, START, END\n",
     "from pydantic import BaseModel\n",
-    "from langchain_core.messages import HumanMessage, AIMessage, BaseMessage\n",
+    "from langchain_core.messages import HumanMessage, AIMessage, AnyMessage\n",
     "from typing import List\n",
     "\n",
     "\n",
     "class ChatState(BaseModel):\n",
-    "    messages: List[BaseMessage]\n",
+    "    messages: List[AnyMessage]\n",
     "    context: str\n",
     "\n",
     "\n",

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/bench/__main__.py

@@ -6,6 +6,7 @@ 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
@@ -228,6 +229,102 @@ benchmarks = (
         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/langgraph/pregel/algo.py

@@ -517,7 +517,7 @@ def prepare_single_task(
                         CONFIG_KEY_CHECKPOINT_ID: None,
                         CONFIG_KEY_CHECKPOINT_NS: task_checkpoint_ns,
                         CONFIG_KEY_SCRATCHPAD: _scratchpad(
-                            config,
+                            config[CONF].get(CONFIG_KEY_SCRATCHPAD),
                             pending_writes,
                             task_id,
                         ),
@@ -627,7 +627,7 @@ def prepare_single_task(
                             CONFIG_KEY_CHECKPOINT_ID: None,
                             CONFIG_KEY_CHECKPOINT_NS: task_checkpoint_ns,
                             CONFIG_KEY_SCRATCHPAD: _scratchpad(
-                                config,
+                                config[CONF].get(CONFIG_KEY_SCRATCHPAD),
                                 pending_writes,
                                 task_id,
                             ),
@@ -655,13 +655,14 @@ def prepare_single_task(
         if checkpoint_null_version is None:
             return
         # If any of the channels read by this process were updated
-        if triggers := _triggers(
+        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)
@@ -721,7 +722,7 @@ def prepare_single_task(
                                 CONFIG_KEY_SEND: partial(
                                     local_write,
                                     writes.extend,
-                                    processes.keys(),
+                                    tuple(processes.keys()),
                                 ),
                                 CONFIG_KEY_READ: partial(
                                     local_read,
@@ -730,7 +731,10 @@ def prepare_single_task(
                                     channels,
                                     managed,
                                     PregelTaskWrites(
-                                        task_path[:3], name, writes, triggers
+                                        task_path[:3],
+                                        name,
+                                        writes,
+                                        triggers,
                                     ),
                                     config,
                                 ),
@@ -748,7 +752,7 @@ def prepare_single_task(
                                 CONFIG_KEY_CHECKPOINT_ID: None,
                                 CONFIG_KEY_CHECKPOINT_NS: task_checkpoint_ns,
                                 CONFIG_KEY_SCRATCHPAD: _scratchpad(
-                                    config,
+                                    config[CONF].get(CONFIG_KEY_SCRATCHPAD),
                                     pending_writes,
                                     task_id,
                                 ),
@@ -799,7 +803,7 @@ def _triggers(
 
 
 def _scratchpad(
-    config: RunnableConfig,
+    parent_scratchpad: Optional[PregelScratchpad],
     pending_writes: list[PendingWrite],
     task_id: str,
 ) -> PregelScratchpad:
@@ -808,9 +812,6 @@ def _scratchpad(
     null_resume_write = next(
         (w for w in pending_writes if w[0] == NULL_TASK_ID and w[1] == RESUME), None
     )
-    parent_scratchpad: Optional[PregelScratchpad] = config[CONF].get(
-        CONFIG_KEY_SCRATCHPAD
-    )
 
     def get_null_resume(consume: bool = False) -> Any:
         if null_resume_write is None:

libs/langgraph/langgraph/pregel/loop.py

@@ -1,6 +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
@@ -571,7 +572,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))
 

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

@@ -149,7 +149,7 @@ class PregelRunner:
                     configurable={
                         CONFIG_KEY_CALL: partial(
                             _call,
-                            t,
+                            weakref.ref(t),
                             retry=retry_policy,
                             futures=weakref.ref(futures),
                             schedule_task=self.schedule_task,
@@ -185,7 +185,7 @@ class PregelRunner:
                     configurable={
                         CONFIG_KEY_CALL: partial(
                             _call,
-                            t,
+                            weakref.ref(t),
                             retry=retry_policy,
                             futures=weakref.ref(futures),
                             schedule_task=self.schedule_task,
@@ -263,7 +263,7 @@ class PregelRunner:
                     configurable={
                         CONFIG_KEY_CALL: partial(
                             _acall,
-                            t,
+                            weakref.ref(t),
                             stream=self.use_astream,
                             retry=retry_policy,
                             futures=weakref.ref(futures),
@@ -304,7 +304,7 @@ class PregelRunner:
                         configurable={
                             CONFIG_KEY_CALL: partial(
                                 _acall,
-                                t,
+                                weakref.ref(t),
                                 retry=retry_policy,
                                 stream=self.use_astream,
                                 futures=weakref.ref(futures),
@@ -469,7 +469,7 @@ def _panic_or_proceed(
 
 
 def _call(
-    task: PregelExecutableTask,
+    task: weakref.ref[PregelExecutableTask],
     func: Callable[[Any], Union[Awaitable[Any], Any]],
     input: Any,
     *,
@@ -489,10 +489,10 @@ def _call(
 
     fut: Optional[concurrent.futures.Future] = None
     # schedule PUSH tasks, collect futures
-    scratchpad: PregelScratchpad = task.config[CONF][CONFIG_KEY_SCRATCHPAD]
+    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,
+        task(),  # type: ignore[arg-type]
         scratchpad.call_counter(),
         Call(func, input, retry=retry, callbacks=callbacks),
     ):
@@ -528,7 +528,7 @@ def _call(
                 configurable={
                     CONFIG_KEY_CALL: partial(
                         _call,
-                        next_task,
+                        weakref.ref(next_task),
                         futures=futures,
                         retry=retry,
                         callbacks=callbacks,
@@ -550,7 +550,7 @@ def _call(
 
 
 def _acall(
-    task: PregelExecutableTask,
+    task: weakref.ref[PregelExecutableTask],
     func: Callable[[Any], Union[Awaitable[Any], Any]],
     input: Any,
     *,
@@ -570,10 +570,10 @@ def _acall(
 ) -> 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]
+    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,
+        task(),  # type: ignore[arg-type]
         scratchpad.call_counter(),
         Call(func, input, retry=retry, callbacks=callbacks),
     ):
@@ -614,7 +614,7 @@ def _acall(
                     configurable={
                         CONFIG_KEY_CALL: partial(
                             _acall,
-                            next_task,
+                            weakref.ref(next_task),
                             stream=stream,
                             futures=futures,
                             schedule_task=schedule_task,
@@ -623,7 +623,7 @@ def _acall(
                             reraise=reraise,
                         ),
                     },
-                    __name__=task.name,
+                    __name__=task().name,  # type: ignore[union-attr]
                     __cancel_on_exit__=True,
                     __reraise_on_exit__=reraise,
                     # starting a new task in the next tick ensures

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/types.py

@@ -133,6 +133,11 @@ class Interrupt:
     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):
     id: str
     name: str
@@ -143,7 +148,14 @@ 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

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.16"
+version = "0.3.18"
 description = "Building stateful, multi-actor applications with LLMs"
 authors = []
 license = "MIT"

libs/langgraph/tests/conftest.py

@@ -1,3 +1,4 @@
+import gc
 import sys
 from contextlib import asynccontextmanager
 from typing import AsyncIterator, Optional
@@ -19,9 +20,17 @@ from langgraph.checkpoint.postgres.aio import (
 from langgraph.checkpoint.serde.encrypted import EncryptedSerializer
 from langgraph.checkpoint.sqlite import SqliteSaver
 from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver
+from langgraph.pregel.loop import AsyncPregelLoop, PregelTaskWrites, SyncPregelLoop
+from langgraph.pregel.runner import PregelRunner
 from langgraph.store.base import BaseStore
 from langgraph.store.memory import InMemoryStore
 from langgraph.store.postgres import AsyncPostgresStore, PostgresStore
+from langgraph.types import (
+    PregelExecutableTask,
+    PregelScratchpad,
+    PregelTask,
+    StateSnapshot,
+)
 
 pytest.register_assert_rewrite("tests.memory_assert")
 
@@ -440,6 +449,31 @@ async def awith_store(store_name: Optional[str]) -> AsyncIterator[BaseStore]:
         raise NotImplementedError(f"Unknown store {store_name}")
 
 
+@pytest.fixture(autouse=True)
+def check_live_objects() -> None:
+    """Check for live objects after each test."""
+    # TODO: Ideally we should be yielding
+    gc.collect()
+    leaked_objs = [
+        o
+        for o in gc.get_objects()
+        if isinstance(
+            o,
+            (
+                PregelExecutableTask,
+                PregelTask,
+                PregelScratchpad,
+                PregelTaskWrites,
+                StateSnapshot,
+                SyncPregelLoop,
+                AsyncPregelLoop,
+                PregelRunner,
+            ),
+        )
+    ]
+    assert not leaked_objs, f"{len(leaked_objs)} leaked objects at end of test."
+
+
 SHALLOW_CHECKPOINTERS_SYNC = ["postgres_shallow"]
 REGULAR_CHECKPOINTERS_SYNC = [
     "memory",

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:"),
@@ -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:"),
@@ -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": (AnyStr("retriever_"),),
+                    "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"),
             },
         },
         {
@@ -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,

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:"),
@@ -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:"),
@@ -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": (AnyStr("retriever_"),),
+                    "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",
+                    ),
                 },
             },
             {
@@ -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,

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,
 )
@@ -6925,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:"),
@@ -7613,3 +7620,320 @@ def test_parallel_interrupts_double(
 
     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
@@ -52,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,
 )
@@ -7582,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:"),
@@ -7838,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/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"