0.2.58

lib: Performance improvements

docs/docs/concepts/low_level.md

@@ -339,37 +339,6 @@ def my_node(state: State) -> Command[Literal["my_other_node"]]:
     )
 ```
 
-`Command` has the following properties:
-
-| Property | Description |
-| --- | --- |
-| `graph` | Graph to send the command to. Supported values:<br>- `None`: the current graph (default)<br>- `Command.PARENT`: closest parent graph |
-| `update` | Update to apply to the graph's state. |
-| `resume` | Value to resume execution with. To be used together with [`interrupt()`][langgraph.types.interrupt]. |
-| `goto` | Can be one of the following:<br>- name of the node to navigate to next (any node that belongs to the specified `graph`)<br>- sequence of node names to navigate to next<br>- `Send` object (to execute a node with the input provided)<br>- sequence of `Send` objects<br>If `goto` is not specified and there are no other tasks left in the graph, the graph will halt after executing the current superstep. |
-
-```python
-from langgraph.graph import StateGraph, START
-from langgraph.types import Command
-from typing_extensions import Literal, TypedDict
-
-class State(TypedDict):
-    foo: str
-
-def my_node(state: State) -> Command[Literal["my_other_node"]]:
-    return Command(update={"foo": "bar"}, goto="my_other_node")
-
-def my_other_node(state: State):
-    return {"foo": state["foo"] + "baz"}
-
-builder = StateGraph(State)
-builder.add_edge(START, "my_node")
-builder.add_node("my_node", my_node)
-builder.add_node("my_other_node", my_other_node)
-
-graph = builder.compile()
-```
-
 With `Command` you can also achieve dynamic control flow behavior (identical to [conditional edges](#conditional-edges)):
 
 ```python
@@ -380,10 +349,40 @@ def my_node(state: State) -> Command[Literal["my_other_node"]]:
 
 !!! important
 
-    When returning `Command` in your node functions, you must add return type annotations with the list of node names the node is routing to, e.g. `Command[Literal["node_b", "node_c"]]`. This is necessary for the graph compilation and rendering, and tells LangGraph that `node_a` can navigate to `node_b` and `node_c`.
+    When returning `Command` in your node functions, you must add return type annotations with the list of node names the node is routing to, e.g. `Command[Literal["my_other_node"]]`. This is necessary for the graph rendering and tells LangGraph that `my_node` can navigate to `my_other_node`.
 
 Check out this [how-to guide](../how-tos/command.ipynb) for an end-to-end example of how to use `Command`.
 
+### When should I use Command instead of conditional edges?
+
+Use `Command` when you need to **both** update the graph state **and** route to a different node. For example, when implementing [multi-agent handoffs](./multi_agent.md#handoffs) where it's important to route to a different agent and pass some information to that agent.
+
+Use [conditional edges](#conditional-edges) to route between nodes conditionally without updating the state.
+
+### Using inside tools
+
+A common use case is updating graph state from inside a tool. For example, in a customer support application you might want to look up customer information based on their account number or ID in the beginning of the conversation. To update the graph state from the tool, you can return `Command(update={"my_custom_key": "foo", "messages": [...]})` from the tool:
+
+```python
+@tool
+def lookup_user_info(tool_call_id: Annotated[str, InjectedToolCallId], config: RunnableConfig):
+    """Use this to look up user information to better assist them with their questions."""
+    user_info = get_user_info(config.get("configurable", {}).get("user_id"))
+    return Command(
+        update={
+            # update the state keys
+            "user_info": user_info,
+            # update the message history
+            "messages": [ToolMessage("Successfully looked up user information", tool_call_id=tool_call_id)]
+        }
+    )
+```
+
+!!! important
+    You MUST include `messages` (or any state key used for the message history) in `Command.update` when returning `Command` from a tool and the list of messages in `messages` MUST contain a `ToolMessage`. This is necessary for the resulting message history to be valid (LLM providers require AI messages with tool calls to be followed by the tool result messages).
+
+If you are using tools that update state via `Command`, we recommend using prebuilt [`ToolNode`][langgraph.prebuilt.tool_node.ToolNode] which automatically handles tools returning `Command` objects and propagates them to the graph state. If you're writing a custom node that calls tools, you would need to manually propagate `Command` objects returned by the tools as the update from node.
+
 ## Persistence
 
 LangGraph provides built-in persistence for your agent's state using [checkpointers][langgraph.checkpoint.base.BaseCheckpointSaver]. Checkpointers save snapshots of the graph state at every superstep, allowing resumption at any time. This enables features like human-in-the-loop interactions, memory management, and fault-tolerance. You can even directly manipulate a graph's state after its execution using the

docs/docs/concepts/multi_agent.md

@@ -26,13 +26,88 @@ There are several ways to connect agents in a multi-agent system:
 - **Hierarchical**: you can define a multi-agent system with [a supervisor of supervisors](https://langchain-ai.github.io/langgraph/tutorials/multi_agent/hierarchical_agent_teams/). This is a generalization of the supervisor architecture and allows for more complex control flows.
 - **Custom multi-agent workflow**: each agent communicates with only a subset of agents. Parts of the flow are deterministic, and only some agents can decide which other agents to call next.
 
+### Handoffs
+
+In multi-agent architectures, agents can be represented as graph nodes. Each agent node executes its step(s) and decides whether to finish execution or route to another agent, including potentially routing to itself (e.g., running in a loop). A common pattern in multi-agent interactions is handoffs, where one agent hands off control to another. Handoffs allow you to specify:
+
+- __destination__: target agent to navigate to (e.g., name of the node to go to)
+- __payload__: [information to pass to that agent](#communication-between-agents) (e.g., state update)
+
+To implement handoffs in LangGraph, agent nodes can return [`Command`](./low_level.md#command) object that allows you to combine both control flow and state updates:
+
+```python
+def agent(state) -> Command[Literal["agent", "another_agent"]]:
+    # the condition for routing/halting can be anything, e.g. LLM tool call / structured output, etc.
+    goto = get_next_agent(...)  # 'agent' / 'another_agent'
+    return Command(
+        # Specify which agent to call next
+        goto=goto,
+        # Update the graph state
+        update={"my_state_key": "my_state_value"}
+    )
+```
+
+In a more complex scenario where each agent node is itself a graph (i.e., a [subgraph](./low_level.md#subgraphs)), a node in one of the agent subgraphs might want to navigate to a different agent. For example, if you have two agents, `alice` and `bob` (subgraph nodes in a parent graph), and `alice` needs to navigate to `bob`, you can set `graph=Command.PARENT` in the `Command` object:
+
+```python
+def some_node_inside_alice(state)
+    return Command(
+        goto="bob",
+        update={"my_state_key": "my_state_value"},
+        # specify which graph to navigate to (defaults to the current graph)
+        graph=Command.PARENT,
+    )
+```
+
+!!! note
+    If you need to support visualization for subgraphs communicating using `Command(graph=Command.PARENT)` you would need to wrap them in a node function with `Command` annotation, e.g. instead of this:
+
+    ```python
+    builder.add_node(alice)
+    ```
+
+    you would need to do this:
+
+    ```python
+    def call_alice(state) -> Command[Literal["bob"]]:
+        return alice.invoke(state)
+
+    builder.add_node("alice", call_alice)
+    ```
+
+#### Handoffs as tools
+
+One of the most common agent types is a ReAct-style tool-calling agents. For those types of agents, a common pattern is wrapping a handoff in a tool call, e.g.:
+
+```python
+def transfer_to_bob(state):
+    """Transfer to bob."""
+    return Command(
+        goto="bob",
+        update={"my_state_key": "my_state_value"},
+        graph=Command.PARENT,
+    )
+```
+
+This is a special case of updating the graph state from tools where in addition the state update, the control flow is included as well.
+
+!!! important
+
+    If you want to use tools that return `Command`, you can either use prebuilt [`create_react_agent`][langgraph.prebuilt.chat_agent_executor.create_react_agent] / [`ToolNode`][langgraph.prebuilt.tool_node.ToolNode] components, or implement your own tool-executing node that collects `Command` objects returned by the tools and returns a list of them, e.g.:
+
+```python
+def call_tools(state):
+    ...
+    commands = [tools_by_name[call["name"].invoke(call, config={"coerce_tool_content": False}) for tool_call in tool_calls]
+    return commands
+```
+
+Let's now take a closer look at the different multi-agent architectures.
+
 ### Network
 
 In this architecture, agents are defined as graph nodes. Each agent can communicate with every other agent (many-to-many connections) and can decide which agent to call next. This architecture is good for problems that do not have a clear hierarchy of agents or a specific sequence in which agents should be called.
 
-### Supervisor
-
-In this architecture, we define agents as nodes and add a supervisor node (LLM) that decides which agent nodes should be called next. We use [conditional edges](./low_level.md#conditional-edges) to route execution to the appropriate agent node based on supervisor's decision. This architecture also lends itself well to running multiple agents in parallel or using [map-reduce](../how-tos/map-reduce.ipynb) pattern.
 
 ```python
 from typing import Literal
@@ -41,39 +116,83 @@ from langgraph.graph import StateGraph, MessagesState, START
 
 model = ChatOpenAI()
 
-class AgentState(MessagesState):
-    next: Literal["agent_1", "agent_2", "__end__"]
-
-def supervisor(state: AgentState):
+def agent_1(state: MessagesState) -> Command[Literal["agent_2", "agent_3", END]]:
     # you can pass relevant parts of the state to the LLM (e.g., state["messages"])
     # to determine which agent to call next. a common pattern is to call the model
     # with a structured output (e.g. force it to return an output with a "next_agent" field)
     response = model.invoke(...)
-    # the "next" key will be used by the conditional edges to route execution
-    # to the appropriate agent
-    return {"next": response["next_agent"]}
+    # route to one of the agents or exit based on the LLM's decision
+    # if the LLM returns "__end__", the graph will finish execution
+    return Command(
+        goto=response["next_agent"],
+        update={"messages": [response["content"]]},
+    )
 
-def agent_1(state: AgentState):
+def agent_2(state: MessagesState) -> Command[Literal["agent_1", "agent_3", END]]:
+    response = model.invoke(...)
+    return Command(
+        goto=response["next_agent"],
+        update={"messages": [response["content"]]},
+    )
+
+def agent_3(state: MessagesState) -> Command[Literal["agent_1", "agent_2", END]]:
+    ...
+    return Command(
+        goto=response["next_agent"],
+        update={"messages": [response["content"]]},
+    )
+
+builder = StateGraph(MessagesState)
+builder.add_node(agent_1)
+builder.add_node(agent_2)
+builder.add_node(agent_3)
+
+builder.add_edge(START, "agent_1")
+network = builder.compile()
+```
+
+### Supervisor
+
+In this architecture, we define agents as nodes and add a supervisor node (LLM) that decides which agent nodes should be called next. We use [`Command`](./low_level.md#command) to route execution to the appropriate agent node based on supervisor's decision. This architecture also lends itself well to running multiple agents in parallel or using [map-reduce](../how-tos/map-reduce.ipynb) pattern.
+
+```python
+from typing import Literal
+from langchain_openai import ChatOpenAI
+from langgraph.graph import StateGraph, MessagesState, START, END
+
+model = ChatOpenAI()
+
+def supervisor(state: MessagesState) -> Command[Literal["agent_1", "agent_2", END]]:
+    # you can pass relevant parts of the state to the LLM (e.g., state["messages"])
+    # to determine which agent to call next. a common pattern is to call the model
+    # with a structured output (e.g. force it to return an output with a "next_agent" field)
+    response = model.invoke(...)
+    # route to one of the agents or exit based on the supervisor's decision
+    # if the supervisor returns "__end__", the graph will finish execution
+    return Command(goto=response["next_agent"])
+
+def agent_1(state: MessagesState) -> Command[Literal["supervisor"]]:
     # you can pass relevant parts of the state to the LLM (e.g., state["messages"])
     # and add any additional logic (different models, custom prompts, structured output, etc.)
     response = model.invoke(...)
-    return {"messages": [response]}
+    return Command(
+        goto="supervisor",
+        update={"messages": [response]},
+    )
 
-def agent_2(state: AgentState):
+def agent_2(state: MessagesState) -> Command[Literal["supervisor"]]:
     response = model.invoke(...)
-    return {"messages": [response]}
+    return Command(
+        goto="supervisor",
+        update={"messages": [response]},
+    )
 
-builder = StateGraph(AgentState)
+builder = StateGraph(MessagesState)
 builder.add_node(supervisor)
 builder.add_node(agent_1)
 builder.add_node(agent_2)
 
 builder.add_edge(START, "supervisor")
-# route to one of the agents or exit based on the supervisor's decisiion
-# if the supervisor returns "__end__", the graph will finish execution
-builder.add_conditional_edges("supervisor", lambda state: state["next"])
-builder.add_edge("agent_1", "supervisor")
-builder.add_edge("agent_2", "supervisor")
 
 supervisor = builder.compile()
 ```
@@ -121,37 +240,29 @@ To address this, you can design your system _hierarchically_. For example, you c
 ```python
 from typing import Literal
 from langchain_openai import ChatOpenAI
-from langgraph.graph import StateGraph, MessagesState, START
+from langgraph.graph import StateGraph, MessagesState, START, END
 
 model = ChatOpenAI()
 
 # define team 1 (same as the single supervisor example above)
-class Team1State(MessagesState):
-    next: Literal["team_1_agent_1", "team_1_agent_2", "__end__"]
 
-def team_1_supervisor(state: Team1State):
+def team_1_supervisor(state: MessagesState) -> Command[Literal["team_1_agent_1", "team_1_agent_2", END]]:
     response = model.invoke(...)
-    return {"next": response["next_agent"]}
+    return Command(goto=response["next_agent"])
 
-def team_1_agent_1(state: Team1State):
+def team_1_agent_1(state: MessagesState) -> Command[Literal["team_1_supervisor"]]:
     response = model.invoke(...)
-    return {"messages": [response]}
+    return Command(goto="team_1_supervisor", update={"messages": [response]})
 
-def team_1_agent_2(state: Team1State):
+def team_1_agent_2(state: MessagesState) -> Command[Literal["team_1_supervisor"]]:
     response = model.invoke(...)
-    return {"messages": [response]}
+    return Command(goto="team_1_supervisor", update={"messages": [response]})
 
 team_1_builder = StateGraph(Team1State)
 team_1_builder.add_node(team_1_supervisor)
 team_1_builder.add_node(team_1_agent_1)
 team_1_builder.add_node(team_1_agent_2)
 team_1_builder.add_edge(START, "team_1_supervisor")
-# route to one of the agents or exit based on the supervisor's decisiion
-# if the supervisor returns "__end__", the graph will finish execution
-team_1_builder.add_conditional_edges("team_1_supervisor", lambda state: state["next"])
-team_1_builder.add_edge("team_1_agent_1", "team_1_supervisor")
-team_1_builder.add_edge("team_1_agent_2", "team_1_supervisor")
-
 team_1_graph = team_1_builder.compile()
 
 # define team 2 (same as the single supervisor example above)
@@ -174,31 +285,22 @@ team_2_graph = team_2_builder.compile()
 
 # define top-level supervisor
 
-class TopLevelState(MessagesState):
-    next: Literal["team_1", "team_2", "__end__"]
-
-builder = StateGraph(TopLevelState)
-def top_level_supervisor(state: TopLevelState):
+builder = StateGraph(MessagesState)
+def top_level_supervisor(state: MessagesState):
     # you can pass relevant parts of the state to the LLM (e.g., state["messages"])
     # to determine which team to call next. a common pattern is to call the model
     # with a structured output (e.g. force it to return an output with a "next_team" field)
     response = model.invoke(...)
-    # the "next" key will be used by the conditional edges to route execution
-    # to the appropriate team
-    return {"next": response["next_team"]}
+    # route to one of the teams or exit based on the supervisor's decision
+    # if the supervisor returns "__end__", the graph will finish execution
+    return Command(goto=response["next_team"])
 
-builder = StateGraph(TopLevelState)
+builder = StateGraph(MessagesState)
 builder.add_node(top_level_supervisor)
 builder.add_node(team_1_graph)
 builder.add_node(team_2_graph)
 
 builder.add_edge(START, "top_level_supervisor")
-# route to one of the teams or exit based on the supervisor's decision
-# if the top-level supervisor returns "__end__", the graph will finish execution
-builder.add_conditional_edges("top_level_supervisor", lambda state: state["next"])
-builder.add_edge("team_1_graph", "top_level_supervisor")
-builder.add_edge("team_2_graph", "top_level_supervisor")
-
 graph = builder.compile()
 ```
 
@@ -208,7 +310,7 @@ In this architecture we add individual agents as graph nodes and define the orde
 
 - **Explicit control flow (normal edges)**: LangGraph allows you to explicitly define the control flow of your application (i.e. the sequence of how agents communicate) explicitly, via [normal graph edges](./low_level.md#normal-edges). This is the most deterministic variant of this architecture above — we always know which agent will be called next ahead of time.
 
-- **Dynamic control flow (conditional edges)**: in LangGraph you can allow LLMs to decide parts of your application control flow. This can be achieved by using [conditional edges](./low_level.md#conditional-edges). A special case of this is a [supervisor tool-calling](#supervisor-tool-calling) architecture. In that case, the tool-calling LLM powering the supervisor agent will make decisions about the order in which the tools (agents) are being called.
+- **Dynamic control flow (Command)**: in LangGraph you can allow LLMs to decide parts of your application control flow. This can be achieved by using [`Command`](./low_level.md#command). A special case of this is a [supervisor tool-calling](#supervisor-tool-calling) architecture. In that case, the tool-calling LLM powering the supervisor agent will make decisions about the order in which the tools (agents) are being called.
 
 ```python
 from langchain_openai import ChatOpenAI

docs/docs/how-tos/command.ipynb

@@ -25,7 +25,7 @@
     "\n",
     "```python\n",
     "def my_node(state: State) -> Command[Literal[\"my_other_node\"]]:\n",
-    "    return GraphCommand(\n",
+    "    return Command(\n",
     "        # state update\n",
     "        update={\"foo\": \"bar\"},\n",
     "        # control flow\n",
@@ -144,7 +144,7 @@
    "id": "badc25eb-4876-482e-bb10-d763023cdaad",
    "metadata": {},
    "source": [
-    "We can now create the `StateGraph` with the above nodes. Notice that the graph doesn't have [conditional edges](../../concepts/low_level#conditional-edges) for routing! This is because control flow is defined with `GraphCommand` inside `node_a`."
+    "We can now create the `StateGraph` with the above nodes. Notice that the graph doesn't have [conditional edges](../../concepts/low_level#conditional-edges) for routing! This is because control flow is defined with `Command` inside `node_a`."
    ]
   },
   {
@@ -171,7 +171,7 @@
    "source": [
     "!!! important\n",
     "\n",
-    "    You might have noticed that we used `Command` as a return type annotation, e.g. `Command[Literal[\"node_b\", \"node_c\"]]`. This is necessary for the graph compilation and rendering, and tells LangGraph that `node_a` can navigate to `node_b` and `node_c`."
+    "    You might have noticed that we used `Command` as a return type annotation, e.g. `Command[Literal[\"node_b\", \"node_c\"]]`. This is necessary for the graph rendering and tells LangGraph that `node_a` can navigate to `node_b` and `node_c`."
    ]
   },
   {

docs/docs/how-tos/index.md

@@ -81,6 +81,7 @@ These how-to guides show common patterns for tool calling with LangGraph:
 - [How to handle tool calling errors](tool-calling-errors.ipynb)
 - [How to pass runtime values to tools](pass-run-time-values-to-tools.ipynb)
 - [How to pass config to tools](pass-config-to-tools.ipynb)
+- [How to update graph state from tools](update-state-from-tools.ipynb)
 - [How to handle large numbers of tools](many-tools.ipynb)
 
 ### Subgraphs
@@ -91,6 +92,12 @@ These how-to guides show common patterns for tool calling with LangGraph:
 - [How to view and update state in subgraphs](subgraphs-manage-state.ipynb)
 - [How to transform inputs and outputs of a subgraph](subgraph-transform-state.ipynb)
 
+### Multi-agent
+
+- [How to build a multi-agent network](multi-agent-network.ipynb)
+
+See the [multi-agent tutorials](../tutorials/index.md#multi-agent-systems) for implementations of other multi-agent architectures.
+
 ### State Management
 
 - [How to use Pydantic model as state](state-model.ipynb)

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

@@ -0,0 +1,383 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "7c58c957-83d8-44ff-8580-a9b3dd39a0a9",
+   "metadata": {},
+   "source": [
+    "# How to update graph state from tools"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "id": "95f30587-8dd2-40be-920d-59539089c09f",
+   "metadata": {},
+   "source": [
+    "!!! info \"Prerequisites\"\n",
+    "    This guide assumes familiarity with the following:\n",
+    "    \n",
+    "    - [Command](../../concepts/low_level/#command)\n",
+    "\n",
+    "A common use case is updating graph state from inside a tool. For example, in a customer support application you might want to look up customer account number or ID in the beginning of the conversation. To update the graph state from the tool, you can return `Command(update={\"my_custom_key\": \"foo\", \"messages\": [...]})` from the tool:\n",
+    "\n",
+    "```python\n",
+    "@tool\n",
+    "def lookup_user_info(tool_call_id: Annotated[str, InjectedToolCallId], config: RunnableConfig):\n",
+    "    \"\"\"Use this to look up user information to better assist them with their questions.\"\"\"\n",
+    "    user_info = get_user_info(config)\n",
+    "    return Command(\n",
+    "        update={\n",
+    "            # update the state keys\n",
+    "            \"user_info\": user_info,\n",
+    "            # update the message history\n",
+    "            \"messages\": [ToolMessage(\"Successfully looked up user information\", tool_call_id=tool_call_id)]\n",
+    "        }\n",
+    "    )\n",
+    "```\n",
+    "\n",
+    "!!! important\n",
+    "\n",
+    "    If you want to use tools that return `Command` and update graph state, you can either use prebuilt [`create_react_agent`][langgraph.prebuilt.chat_agent_executor.create_react_agent] / [`ToolNode`][langgraph.prebuilt.tool_node.ToolNode] components, or implement your own tool-executing node that collects `Command` objects returned by the tools and returns a list of them, e.g.:\n",
+    "    \n",
+    "    ```python\n",
+    "    def call_tools(state):\n",
+    "        ...\n",
+    "        commands = [tools_by_name[call[\"name\"].invoke(call, config={\"coerce_tool_content\": False}) for tool_call in tool_calls]\n",
+    "        return commands\n",
+    "    ```\n",
+    "\n",
+    "This guide shows how you can do this using LangGraph's prebuilt components ([`create_react_agent`][langgraph.prebuilt.chat_agent_executor.create_react_agent] / [`ToolNode`][langgraph.prebuilt.tool_node.ToolNode]).\n",
+    "\n",
+    "!!! note\n",
+    "\n",
+    "    Support for tools that return [`Command`][langgraph.types.Command] was added in LangGraph `v0.2.57`.\n",
+    "\n",
+    "## Setup\n",
+    "\n",
+    "First, let's install the required packages and set our API keys:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "64500eca-1cdc-43d9-9401-f4cd9999881f",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%%capture --no-stderr\n",
+    "%pip install -U langgraph langchain-openai"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "a3f92fb2-9175-47fa-9c7d-ad5f44bfd20e",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Please provide your OPENAI_API_KEY ········\n"
+     ]
+    }
+   ],
+   "source": [
+    "import os\n",
+    "import getpass\n",
+    "\n",
+    "\n",
+    "def _set_if_undefined(var: str):\n",
+    "    if not os.environ.get(var):\n",
+    "        os.environ[var] = getpass.getpass(f\"Please provide your {var}\")\n",
+    "\n",
+    "\n",
+    "_set_if_undefined(\"OPENAI_API_KEY\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "caf6ff9f-c1e6-499e-a230-9fa231ea7d2f",
+   "metadata": {},
+   "source": [
+    "<div class=\"admonition tip\">\n",
+    "    <p class=\"admonition-title\">Set up <a href=\"https://smith.langchain.com\">LangSmith</a> for LangGraph development</p>\n",
+    "    <p style=\"padding-top: 5px;\">\n",
+    "        Sign up for LangSmith to quickly spot issues and improve the performance of your LangGraph projects. LangSmith lets you use trace data to debug, test, and monitor your LLM apps built with LangGraph — read more about how to get started <a href=\"https://docs.smith.langchain.com\">here</a>. \n",
+    "    </p>\n",
+    "</div>"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "10e9a9c6-fa3f-416c-bac0-3e58d7259908",
+   "metadata": {},
+   "source": [
+    "Let's create a simple ReAct style agent that can look up user information and personalize the response based on the user info."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "4255b9b9-cf67-4cc3-8018-1708f5dfcfd2",
+   "metadata": {},
+   "source": [
+    "## Define tool"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "7de6b010-aab1-4fe8-8251-907fcae78583",
+   "metadata": {},
+   "source": [
+    "First, let's define the tool that we'll be using to look up user information. We'll use a naive implementation that simply looks user information up using a dictionary:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "8d070c9f-6e61-4724-85dc-ac4531b9c79a",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "USER_INFO = [\n",
+    "    {\"user_id\": \"1\", \"name\": \"Bob Dylan\", \"location\": \"New York, NY\"},\n",
+    "    {\"user_id\": \"2\", \"name\": \"Taylor Swift\", \"location\": \"Beverly Hills, CA\"},\n",
+    "]\n",
+    "\n",
+    "USER_ID_TO_USER_INFO = {info[\"user_id\"]: info for info in USER_INFO}"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "08d1ecca-ee57-4e97-b8d0-e09de85337d4",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from langgraph.prebuilt.chat_agent_executor import AgentState\n",
+    "from langgraph.types import Command\n",
+    "from langchain_core.tools import tool\n",
+    "from langchain_core.tools.base import InjectedToolCallId\n",
+    "from langchain_core.messages import ToolMessage\n",
+    "from langchain_core.runnables import RunnableConfig\n",
+    "\n",
+    "from typing_extensions import Any, Annotated\n",
+    "\n",
+    "\n",
+    "class State(AgentState):\n",
+    "    # user provided\n",
+    "    last_name: str\n",
+    "    # updated by the tool\n",
+    "    user_info: dict[str, Any]\n",
+    "\n",
+    "\n",
+    "@tool\n",
+    "def lookup_user_info(\n",
+    "    tool_call_id: Annotated[str, InjectedToolCallId], config: RunnableConfig\n",
+    "):\n",
+    "    \"\"\"Use this to look up user information to better assist them with their questions.\"\"\"\n",
+    "    user_id = config.get(\"configurable\", {}).get(\"user_id\")\n",
+    "    if user_id is None:\n",
+    "        raise ValueError(\"Please provide user ID\")\n",
+    "\n",
+    "    if user_id not in USER_ID_TO_USER_INFO:\n",
+    "        raise ValueError(f\"User '{user_id}' not found\")\n",
+    "\n",
+    "    user_info = USER_ID_TO_USER_INFO[user_id]\n",
+    "    return Command(\n",
+    "        update={\n",
+    "            # update the state keys\n",
+    "            \"user_info\": user_info,\n",
+    "            # update the message history\n",
+    "            \"messages\": [\n",
+    "                ToolMessage(\n",
+    "                    \"Successfully looked up user information\", tool_call_id=tool_call_id\n",
+    "                )\n",
+    "            ],\n",
+    "        }\n",
+    "    )"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "b99e5f24-5e5e-4a34-baae-467182675bb5",
+   "metadata": {},
+   "source": [
+    "## Define prompt"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "cbb06aea-6654-4245-91f8-af6e8f2b5377",
+   "metadata": {},
+   "source": [
+    "Let's now add personalization: we'll respond differently to the user based on the state values AFTER the state has been updated from the tool. To achieve this, let's define a function that will dynamically construct the system prompt based on the graph state. It will be called ever time the LLM is called and the function output will be passed to the LLM:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "c553d062-d145-4145-84bd-9b798f7c95c2",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def state_modifier(state: State):\n",
+    "    user_info = state.get(\"user_info\")\n",
+    "    if user_info is None:\n",
+    "        return state[\"messages\"]\n",
+    "\n",
+    "    system_msg = (\n",
+    "        f\"User name is {user_info['name']}. User lives in {user_info['location']}\"\n",
+    "    )\n",
+    "    return [{\"role\": \"system\", \"content\": system_msg}] + state[\"messages\"]"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "c5acdd5d-68be-466b-9c21-46cbed91d2bc",
+   "metadata": {},
+   "source": [
+    "## Define graph"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "afb65028-0359-46c8-b09c-ffc90180f759",
+   "metadata": {},
+   "source": [
+    "Finally, let's combine this into a single graph using the prebuilt `create_react_agent`:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "2d59db29-fd51-4d29-9854-21763a4855e3",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from langgraph.prebuilt import create_react_agent\n",
+    "from langchain_openai import ChatOpenAI\n",
+    "\n",
+    "model = ChatOpenAI(model=\"gpt-4o\")\n",
+    "\n",
+    "agent = create_react_agent(\n",
+    "    model,\n",
+    "    # pass the tool that can update state\n",
+    "    [lookup_user_info],\n",
+    "    state_schema=State,\n",
+    "    # pass dynamic prompt function\n",
+    "    state_modifier=state_modifier,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "0782b8ab-a603-47b8-9a76-77f593402678",
+   "metadata": {},
+   "source": [
+    "## Use it!"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "6165e153-ab28-4404-adea-796c7bd0701b",
+   "metadata": {},
+   "source": [
+    "Let's now try running our agent. We'll need to provide user ID in the config so that our tool knows what information to look up:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "de34a58b-1765-4b63-a232-d46790aff884",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "{'agent': {'messages': [AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_7LSUh6ZDvGJAUvlWvXiCK4Gf', 'function': {'arguments': '{}', 'name': 'lookup_user_info'}, 'type': 'function'}], 'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 11, 'prompt_tokens': 56, 'total_tokens': 67, 'completion_tokens_details': {'accepted_prediction_tokens': 0, 'audio_tokens': 0, 'reasoning_tokens': 0, 'rejected_prediction_tokens': 0}, 'prompt_tokens_details': {'audio_tokens': 0, 'cached_tokens': 0}}, 'model_name': 'gpt-4o-2024-08-06', 'system_fingerprint': 'fp_9d50cd990b', 'finish_reason': 'tool_calls', 'logprobs': None}, id='run-57eeb216-e35d-4501-aaac-b5c6b26fb17c-0', tool_calls=[{'name': 'lookup_user_info', 'args': {}, 'id': 'call_7LSUh6ZDvGJAUvlWvXiCK4Gf', 'type': 'tool_call'}], usage_metadata={'input_tokens': 56, 'output_tokens': 11, 'total_tokens': 67, 'input_token_details': {'audio': 0, 'cache_read': 0}, 'output_token_details': {'audio': 0, 'reasoning': 0}})]}}\n",
+      "\n",
+      "\n",
+      "{'tools': {'user_info': {'user_id': '1', 'name': 'Bob Dylan', 'location': 'New York, NY'}, 'messages': [ToolMessage(content='Successfully looked up user information', name='lookup_user_info', id='168d8ff8-b021-4c8b-a11a-3b50c30a072c', tool_call_id='call_7LSUh6ZDvGJAUvlWvXiCK4Gf')]}}\n",
+      "\n",
+      "\n",
+      "{'agent': {'messages': [AIMessage(content=\"Hi Bob! Since you're in New York, NY, there are plenty of exciting things to do over the weekend. Here are some suggestions:\\n\\n1. **Explore Central Park**: Take a leisurely walk, rent a bike, or have a picnic in this iconic park.\\n\\n2. **Visit a Museum**: Check out The Metropolitan Museum of Art or the Museum of Modern Art (MoMA) for an enriching cultural experience.\\n\\n3. **Broadway Show**: Catch a Broadway show or an off-Broadway performance for some world-class entertainment.\\n\\n4. **Food Tour**: Explore different neighborhoods like Greenwich Village or Williamsburg for diverse culinary experiences.\\n\\n5. **Brooklyn Bridge Walk**: Take a walk across the Brooklyn Bridge for stunning views of the city skyline.\\n\\n6. **Visit a Rooftop Bar**: Enjoy a drink with a view at one of New York’s many rooftop bars.\\n\\n7. **Explore a New Neighborhood**: Discover the unique charm of areas like SoHo, Chelsea, or Astoria.\\n\\n8. **Live Music**: Check out live music venues for a night of great performances.\\n\\n9. **Art Galleries**: Visit some of the smaller art galleries around Chelsea or the Lower East Side.\\n\\n10. **Attend a Local Event**: Look up any local events or festivals happening this weekend.\\n\\nFeel free to let me know if you want more details on any of these activities!\", additional_kwargs={'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 285, 'prompt_tokens': 95, 'total_tokens': 380, 'completion_tokens_details': {'accepted_prediction_tokens': 0, 'audio_tokens': 0, 'reasoning_tokens': 0, 'rejected_prediction_tokens': 0}, 'prompt_tokens_details': {'audio_tokens': 0, 'cached_tokens': 0}}, 'model_name': 'gpt-4o-2024-08-06', 'system_fingerprint': 'fp_9d50cd990b', 'finish_reason': 'stop', 'logprobs': None}, id='run-f13ce15b-02b6-40e6-8264-c4d9edd0d03a-0', usage_metadata={'input_tokens': 95, 'output_tokens': 285, 'total_tokens': 380, 'input_token_details': {'audio': 0, 'cache_read': 0}, 'output_token_details': {'audio': 0, 'reasoning': 0}})]}}\n",
+      "\n",
+      "\n"
+     ]
+    }
+   ],
+   "source": [
+    "for chunk in agent.stream(\n",
+    "    {\"messages\": [(\"user\", \"hi, what should i do this weekend?\")]},\n",
+    "    # provide user ID in the config\n",
+    "    {\"configurable\": {\"user_id\": \"1\"}},\n",
+    "):\n",
+    "    print(chunk)\n",
+    "    print(\"\\n\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "d9b2281f-269c-41dd-b6b2-4c743f11ffc9",
+   "metadata": {},
+   "source": [
+    "We can see that the model correctly recommended some New York activities for Bob Dylan! Let's try getting recommendations for Taylor Swift:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "id": "9d71af94-572a-4961-88a7-665e792cf96a",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "{'agent': {'messages': [AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_5HLtJtzcgmKbtmK6By21wW5Y', 'function': {'arguments': '{}', 'name': 'lookup_user_info'}, 'type': 'function'}], 'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 11, 'prompt_tokens': 56, 'total_tokens': 67, 'completion_tokens_details': {'accepted_prediction_tokens': 0, 'audio_tokens': 0, 'reasoning_tokens': 0, 'rejected_prediction_tokens': 0}, 'prompt_tokens_details': {'audio_tokens': 0, 'cached_tokens': 0}}, 'model_name': 'gpt-4o-2024-08-06', 'system_fingerprint': 'fp_c7ca0ebaca', 'finish_reason': 'tool_calls', 'logprobs': None}, id='run-bacacd7d-76cc-4f6b-9e9b-d9e6f00b9391-0', tool_calls=[{'name': 'lookup_user_info', 'args': {}, 'id': 'call_5HLtJtzcgmKbtmK6By21wW5Y', 'type': 'tool_call'}], usage_metadata={'input_tokens': 56, 'output_tokens': 11, 'total_tokens': 67, 'input_token_details': {'audio': 0, 'cache_read': 0}, 'output_token_details': {'audio': 0, 'reasoning': 0}})]}}\n",
+      "\n",
+      "\n",
+      "{'tools': {'user_info': {'user_id': '2', 'name': 'Taylor Swift', 'location': 'Beverly Hills, CA'}, 'messages': [ToolMessage(content='Successfully looked up user information', name='lookup_user_info', id='d81ef31e-6d77-4f13-ae86-e2e6ba567e3d', tool_call_id='call_5HLtJtzcgmKbtmK6By21wW5Y')]}}\n",
+      "\n",
+      "\n",
+      "{'agent': {'messages': [AIMessage(content=\"Hi Taylor! Since you're in Beverly Hills, here are a few suggestions for a fun weekend:\\n\\n1. **Hiking at Runyon Canyon**: Enjoy a scenic hike with beautiful views of Los Angeles. It's a great way to get some exercise and enjoy the outdoors.\\n\\n2. **Visit Rodeo Drive**: Spend some time shopping or window shopping at the famous Rodeo Drive. You might even spot some celebrities!\\n\\n3. **Explore the Getty Center**: Check out the art collections and beautiful gardens at the Getty Center. The architecture and views are stunning.\\n\\n4. **Relax at a Spa**: Treat yourself to a relaxing day at one of Beverly Hills' luxurious spas.\\n\\n5. **Dining Out**: Try a new restaurant or visit your favorite spot for a delicious meal. Beverly Hills has a fantastic dining scene.\\n\\n6. **Attend a Local Event**: Check out any local events or concerts happening this weekend. Beverly Hills often hosts exciting events.\\n\\nEnjoy your weekend!\", additional_kwargs={'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 198, 'prompt_tokens': 95, 'total_tokens': 293, 'completion_tokens_details': {'accepted_prediction_tokens': 0, 'audio_tokens': 0, 'reasoning_tokens': 0, 'rejected_prediction_tokens': 0}, 'prompt_tokens_details': {'audio_tokens': 0, 'cached_tokens': 0}}, 'model_name': 'gpt-4o-2024-08-06', 'system_fingerprint': 'fp_c7ca0ebaca', 'finish_reason': 'stop', 'logprobs': None}, id='run-2057df76-f192-4c69-a66a-1f0a86bf5d66-0', usage_metadata={'input_tokens': 95, 'output_tokens': 198, 'total_tokens': 293, 'input_token_details': {'audio': 0, 'cache_read': 0}, 'output_token_details': {'audio': 0, 'reasoning': 0}})]}}\n",
+      "\n",
+      "\n"
+     ]
+    }
+   ],
+   "source": [
+    "for chunk in agent.stream(\n",
+    "    {\"messages\": [(\"user\", \"hi, what should i do this weekend?\")]},\n",
+    "    {\"configurable\": {\"user_id\": \"2\"}},\n",
+    "):\n",
+    "    print(chunk)\n",
+    "    print(\"\\n\")"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.12.3"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

docs/docs/tutorials/multi_agent/hierarchical_agent_teams.ipynb

@@ -289,15 +289,10 @@
     "from langchain_core.language_models.chat_models import BaseChatModel\n",
     "\n",
     "from langgraph.graph import StateGraph, MessagesState, START, END\n",
+    "from langgraph.types import Command\n",
     "from langchain_core.messages import HumanMessage, trim_messages\n",
     "\n",
     "\n",
-    "# The agent state is the input to each node in the graph\n",
-    "class AgentState(MessagesState):\n",
-    "    # The 'next' field indicates where to route to next\n",
-    "    next: str\n",
-    "\n",
-    "\n",
     "def make_supervisor_node(llm: BaseChatModel, members: list[str]) -> str:\n",
     "    options = [\"FINISH\"] + members\n",
     "    system_prompt = (\n",
@@ -313,17 +308,17 @@
     "\n",
     "        next: Literal[*options]\n",
     "\n",
-    "    def supervisor_node(state: MessagesState) -> MessagesState:\n",
+    "    def supervisor_node(state: MessagesState) -> Command[Literal[*members, \"__end__\"]]:\n",
     "        \"\"\"An LLM-based router.\"\"\"\n",
     "        messages = [\n",
     "            {\"role\": \"system\", \"content\": system_prompt},\n",
     "        ] + state[\"messages\"]\n",
     "        response = llm.with_structured_output(Router).invoke(messages)\n",
-    "        next_ = response[\"next\"]\n",
-    "        if next_ == \"FINISH\":\n",
-    "            next_ = END\n",
+    "        goto = response[\"next\"]\n",
+    "        if goto == \"FINISH\":\n",
+    "            goto = END\n",
     "\n",
-    "        return {\"next\": next_}\n",
+    "        return Command(goto=goto)\n",
     "\n",
     "    return supervisor_node"
    ]
@@ -363,25 +358,33 @@
     "search_agent = create_react_agent(llm, tools=[tavily_tool])\n",
     "\n",
     "\n",
-    "def search_node(state: AgentState) -> AgentState:\n",
+    "def search_node(state: MessagesState) -> Command[Literal[\"supervisor\"]]:\n",
     "    result = search_agent.invoke(state)\n",
-    "    return {\n",
-    "        \"messages\": [\n",
-    "            HumanMessage(content=result[\"messages\"][-1].content, name=\"search\")\n",
-    "        ]\n",
-    "    }\n",
+    "    return Command(\n",
+    "        update={\n",
+    "            \"messages\": [\n",
+    "                HumanMessage(content=result[\"messages\"][-1].content, name=\"search\")\n",
+    "            ]\n",
+    "        },\n",
+    "        # We want our workers to ALWAYS \"report back\" to the supervisor when done\n",
+    "        goto=\"supervisor\",\n",
+    "    )\n",
     "\n",
     "\n",
     "web_scraper_agent = create_react_agent(llm, tools=[scrape_webpages])\n",
     "\n",
     "\n",
-    "def web_scraper_node(state: AgentState) -> AgentState:\n",
+    "def web_scraper_node(state: MessagesState) -> Command[Literal[\"supervisor\"]]:\n",
     "    result = web_scraper_agent.invoke(state)\n",
-    "    return {\n",
-    "        \"messages\": [\n",
-    "            HumanMessage(content=result[\"messages\"][-1].content, name=\"web_scraper\")\n",
-    "        ]\n",
-    "    }\n",
+    "    return Command(\n",
+    "        update={\n",
+    "            \"messages\": [\n",
+    "                HumanMessage(content=result[\"messages\"][-1].content, name=\"web_scraper\")\n",
+    "            ]\n",
+    "        },\n",
+    "        # We want our workers to ALWAYS \"report back\" to the supervisor when done\n",
+    "        goto=\"supervisor\",\n",
+    "    )\n",
     "\n",
     "\n",
     "research_supervisor_node = make_supervisor_node(llm, [\"search\", \"web_scraper\"])"
@@ -412,14 +415,7 @@
     "research_builder.add_node(\"search\", search_node)\n",
     "research_builder.add_node(\"web_scraper\", web_scraper_node)\n",
     "\n",
-    "# Define the control flow\n",
     "research_builder.add_edge(START, \"supervisor\")\n",
-    "# We want our workers to ALWAYS \"report back\" to the supervisor when done\n",
-    "research_builder.add_edge(\"search\", \"supervisor\")\n",
-    "research_builder.add_edge(\"web_scraper\", \"supervisor\")\n",
-    "# Add the edges where routing applies\n",
-    "research_builder.add_conditional_edges(\"supervisor\", lambda state: state[\"next\"])\n",
-    "\n",
     "research_graph = research_builder.compile()"
    ]
   },
@@ -532,13 +528,17 @@
     ")\n",
     "\n",
     "\n",
-    "def doc_writing_node(state: AgentState) -> AgentState:\n",
+    "def doc_writing_node(state: MessagesState) -> Command[Literal[\"supervisor\"]]:\n",
     "    result = doc_writer_agent.invoke(state)\n",
-    "    return {\n",
-    "        \"messages\": [\n",
-    "            HumanMessage(content=result[\"messages\"][-1].content, name=\"doc_writer\")\n",
-    "        ]\n",
-    "    }\n",
+    "    return Command(\n",
+    "        update={\n",
+    "            \"messages\": [\n",
+    "                HumanMessage(content=result[\"messages\"][-1].content, name=\"doc_writer\")\n",
+    "            ]\n",
+    "        },\n",
+    "        # We want our workers to ALWAYS \"report back\" to the supervisor when done\n",
+    "        goto=\"supervisor\",\n",
+    "    )\n",
     "\n",
     "\n",
     "note_taking_agent = create_react_agent(\n",
@@ -551,13 +551,17 @@
     ")\n",
     "\n",
     "\n",
-    "def note_taking_node(state: AgentState) -> AgentState:\n",
+    "def note_taking_node(state: MessagesState) -> Command[Literal[\"supervisor\"]]:\n",
     "    result = note_taking_agent.invoke(state)\n",
-    "    return {\n",
-    "        \"messages\": [\n",
-    "            HumanMessage(content=result[\"messages\"][-1].content, name=\"note_taker\")\n",
-    "        ]\n",
-    "    }\n",
+    "    return Command(\n",
+    "        update={\n",
+    "            \"messages\": [\n",
+    "                HumanMessage(content=result[\"messages\"][-1].content, name=\"note_taker\")\n",
+    "            ]\n",
+    "        },\n",
+    "        # We want our workers to ALWAYS \"report back\" to the supervisor when done\n",
+    "        goto=\"supervisor\",\n",
+    "    )\n",
     "\n",
     "\n",
     "chart_generating_agent = create_react_agent(\n",
@@ -565,13 +569,19 @@
     ")\n",
     "\n",
     "\n",
-    "def chart_generating_node(state: AgentState) -> AgentState:\n",
+    "def chart_generating_node(state: MessagesState) -> Command[Literal[\"supervisor\"]]:\n",
     "    result = chart_generating_agent.invoke(state)\n",
-    "    return {\n",
-    "        \"messages\": [\n",
-    "            HumanMessage(content=result[\"messages\"][-1].content, name=\"chart_generator\")\n",
-    "        ]\n",
-    "    }\n",
+    "    return Command(\n",
+    "        update={\n",
+    "            \"messages\": [\n",
+    "                HumanMessage(\n",
+    "                    content=result[\"messages\"][-1].content, name=\"chart_generator\"\n",
+    "                )\n",
+    "            ]\n",
+    "        },\n",
+    "        # We want our workers to ALWAYS \"report back\" to the supervisor when done\n",
+    "        goto=\"supervisor\",\n",
+    "    )\n",
     "\n",
     "\n",
     "doc_writing_supervisor_node = make_supervisor_node(\n",
@@ -600,21 +610,13 @@
    "outputs": [],
    "source": [
     "# Create the graph here\n",
-    "paper_writing_builder = StateGraph(AgentState)\n",
+    "paper_writing_builder = StateGraph(MessagesState)\n",
     "paper_writing_builder.add_node(\"supervisor\", doc_writing_supervisor_node)\n",
     "paper_writing_builder.add_node(\"doc_writer\", doc_writing_node)\n",
     "paper_writing_builder.add_node(\"note_taker\", note_taking_node)\n",
     "paper_writing_builder.add_node(\"chart_generator\", chart_generating_node)\n",
     "\n",
-    "# Define the control flow\n",
     "paper_writing_builder.add_edge(START, \"supervisor\")\n",
-    "# We want our workers to ALWAYS \"report back\" to the supervisor when done\n",
-    "paper_writing_builder.add_edge(\"doc_writer\", \"supervisor\")\n",
-    "paper_writing_builder.add_edge(\"note_taker\", \"supervisor\")\n",
-    "paper_writing_builder.add_edge(\"chart_generator\", \"supervisor\")\n",
-    "# Add the edges where routing applies\n",
-    "paper_writing_builder.add_conditional_edges(\"supervisor\", lambda state: state[\"next\"])\n",
-    "\n",
     "paper_writing_graph = paper_writing_builder.compile()"
    ]
   },
@@ -728,37 +730,41 @@
    },
    "outputs": [],
    "source": [
-    "def call_research_team(state: AgentState) -> AgentState:\n",
+    "def call_research_team(state: MessagesState) -> Command[Literal[\"supervisor\"]]:\n",
     "    response = research_graph.invoke({\"messages\": state[\"messages\"][-1]})\n",
-    "    return {\n",
-    "        \"messages\": [\n",
-    "            HumanMessage(content=response[\"messages\"][-1].content, name=\"research_team\")\n",
-    "        ]\n",
-    "    }\n",
+    "    return Command(\n",
+    "        update={\n",
+    "            \"messages\": [\n",
+    "                HumanMessage(\n",
+    "                    content=response[\"messages\"][-1].content, name=\"research_team\"\n",
+    "                )\n",
+    "            ]\n",
+    "        },\n",
+    "        goto=\"supervisor\",\n",
+    "    )\n",
     "\n",
     "\n",
-    "def call_paper_writing_team(state: AgentState) -> AgentState:\n",
+    "def call_paper_writing_team(state: MessagesState) -> Command[Literal[\"supervisor\"]]:\n",
     "    response = paper_writing_graph.invoke({\"messages\": state[\"messages\"][-1]})\n",
-    "    return {\n",
-    "        \"messages\": [\n",
-    "            HumanMessage(content=response[\"messages\"][-1].content, name=\"writing_team\")\n",
-    "        ]\n",
-    "    }\n",
+    "    return Command(\n",
+    "        update={\n",
+    "            \"messages\": [\n",
+    "                HumanMessage(\n",
+    "                    content=response[\"messages\"][-1].content, name=\"writing_team\"\n",
+    "                )\n",
+    "            ]\n",
+    "        },\n",
+    "        goto=\"supervisor\",\n",
+    "    )\n",
     "\n",
     "\n",
     "# Define the graph.\n",
-    "super_builder = StateGraph(AgentState)\n",
+    "super_builder = StateGraph(MessagesState)\n",
     "super_builder.add_node(\"supervisor\", teams_supervisor_node)\n",
     "super_builder.add_node(\"research_team\", call_research_team)\n",
     "super_builder.add_node(\"writing_team\", call_paper_writing_team)\n",
     "\n",
-    "# Define the control flow\n",
     "super_builder.add_edge(START, \"supervisor\")\n",
-    "# We want our teams to ALWAYS \"report back\" to the top-level supervisor when done\n",
-    "super_builder.add_edge(\"research_team\", \"supervisor\")\n",
-    "super_builder.add_edge(\"writing_team\", \"supervisor\")\n",
-    "# Add the edges where routing applies\n",
-    "super_builder.add_conditional_edges(\"supervisor\", lambda state: state[\"next\"])\n",
     "super_graph = super_builder.compile()"
    ]
   },

docs/mkdocs.yml

@@ -192,6 +192,7 @@ nav:
             - how-tos/tool-calling.ipynb
             - how-tos/tool-calling-errors.ipynb
             - how-tos/pass-run-time-values-to-tools.ipynb
+            - how-tos/update-state-from-tools.ipynb
             - how-tos/pass-config-to-tools.ipynb
             - how-tos/many-tools.ipynb
         - Subgraphs:
@@ -199,6 +200,8 @@ nav:
             - how-tos/subgraph.ipynb
             - how-tos/subgraphs-manage-state.ipynb
             - how-tos/subgraph-transform-state.ipynb
+        - Multi-agent:
+            - how-tos/multi-agent-network.ipynb
         - State Management:
             - State Management: how-tos#state-management
             - how-tos/state-model.ipynb

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

@@ -57,6 +57,17 @@ MIGRATIONS = [
     PRIMARY KEY (thread_id, checkpoint_ns, checkpoint_id, task_id, idx)
 );""",
     "ALTER TABLE checkpoint_blobs ALTER COLUMN blob DROP not null;",
+    """
+    """,
+    """
+    CREATE INDEX CONCURRENTLY IF NOT EXISTS checkpoints_thread_id_idx ON checkpoints(thread_id);
+    """,
+    """
+    CREATE INDEX CONCURRENTLY IF NOT EXISTS checkpoint_blobs_thread_id_idx ON checkpoint_blobs(thread_id);
+    """,
+    """
+    CREATE INDEX CONCURRENTLY IF NOT EXISTS checkpoint_writes_thread_id_idx ON checkpoint_writes(thread_id);
+    """,
 ]
 
 SELECT_SQL = f"""

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

@@ -6,7 +6,6 @@ from typing import Any, Callable, Optional, Union, cast
 
 import orjson
 from psycopg import AsyncConnection, AsyncCursor, AsyncPipeline, Capabilities
-from psycopg.errors import UndefinedTable
 from psycopg.rows import DictRow, dict_row
 from psycopg_pool import AsyncConnectionPool
 
@@ -19,9 +18,7 @@ from langgraph.store.base import (
     Result,
     SearchOp,
 )
-from langgraph.store.base.batch import (
-    BatchedBaseStore,
-)
+from langgraph.store.base.batch import AsyncBatchedBaseStore
 from langgraph.store.postgres.base import (
     _PLACEHOLDER,
     BasePostgresStore,
@@ -38,7 +35,7 @@ from langgraph.store.postgres.base import (
 logger = logging.getLogger(__name__)
 
 
-class AsyncPostgresStore(BatchedBaseStore, BasePostgresStore[_ainternal.Conn]):
+class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[_ainternal.Conn]):
     """Asynchronous Postgres-backed store with optional vector search using pgvector.
 
     !!! example "Examples"
@@ -158,9 +155,6 @@ class AsyncPostgresStore(BatchedBaseStore, BasePostgresStore[_ainternal.Conn]):
 
         return results
 
-    def batch(self, ops: Iterable[Op]) -> list[Result]:
-        return asyncio.run_coroutine_threadsafe(self.abatch(ops), self.loop).result()
-
     @classmethod
     @asynccontextmanager
     async def from_conn_string(
@@ -221,22 +215,19 @@ class AsyncPostgresStore(BatchedBaseStore, BasePostgresStore[_ainternal.Conn]):
         """
 
         async def _get_version(cur: AsyncCursor[DictRow], table: str) -> int:
-            try:
-                await cur.execute(f"SELECT v FROM {table} ORDER BY v DESC LIMIT 1")
-                row = await cur.fetchone()
-                if row is None:
-                    version = -1
-                else:
-                    version = row["v"]
-            except UndefinedTable:
-                version = -1
-                await cur.execute(
-                    f"""
-                    CREATE TABLE IF NOT EXISTS {table} (
-                        v INTEGER PRIMARY KEY
-                    )
-                """
+            await cur.execute(
+                f"""
+                CREATE TABLE IF NOT EXISTS {table} (
+                    v INTEGER PRIMARY KEY
                 )
+            """
+            )
+            await cur.execute(f"SELECT v FROM {table} ORDER BY v DESC LIMIT 1")
+            row = cast(dict, await cur.fetchone())
+            if row is None:
+                version = -1
+            else:
+                version = row["v"]
             return version
 
         async with self._cursor() as cur:

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

@@ -21,7 +21,6 @@ from typing import (
 
 import orjson
 from psycopg import Capabilities, Connection, Cursor, Pipeline
-from psycopg.errors import UndefinedTable
 from psycopg.rows import DictRow, dict_row
 from psycopg.types.json import Jsonb
 from psycopg_pool import ConnectionPool
@@ -30,6 +29,7 @@ from typing_extensions import TypedDict
 from langgraph.checkpoint.postgres import _ainternal as _ainternal
 from langgraph.checkpoint.postgres import _internal as _pg_internal
 from langgraph.store.base import (
+    BaseStore,
     GetOp,
     IndexConfig,
     Item,
@@ -43,7 +43,6 @@ from langgraph.store.base import (
     get_text_at_path,
     tokenize_path,
 )
-from langgraph.store.base.batch import SyncBatchedBaseStore
 
 if TYPE_CHECKING:
     from langchain_core.embeddings import Embeddings
@@ -73,7 +72,7 @@ CREATE TABLE IF NOT EXISTS store (
 """,
     """
 -- For faster lookups by prefix
-CREATE INDEX IF NOT EXISTS store_prefix_idx ON store USING btree (prefix text_pattern_ops);
+CREATE INDEX CONCURRENTLY IF NOT EXISTS store_prefix_idx ON store USING btree (prefix text_pattern_ops);
 """,
 ]
 
@@ -107,7 +106,7 @@ CREATE TABLE IF NOT EXISTS store_vectors (
     ),
     Migration(
         """
-CREATE INDEX IF NOT EXISTS store_vectors_embedding_idx ON store_vectors 
+CREATE INDEX CONCURRENTLY IF NOT EXISTS store_vectors_embedding_idx ON store_vectors 
     USING %(index_type)s (embedding %(ops)s)%(index_params)s;
 """,
         condition=lambda store: bool(
@@ -533,7 +532,7 @@ class BasePostgresStore(Generic[C]):
             raise ValueError(f"Unsupported operator: {op}")
 
 
-class PostgresStore(SyncBatchedBaseStore, BasePostgresStore[_pg_internal.Conn]):
+class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]):
     """Postgres-backed store with optional vector search using pgvector.
 
     !!! example "Examples"
@@ -847,22 +846,19 @@ class PostgresStore(SyncBatchedBaseStore, BasePostgresStore[_pg_internal.Conn]):
         """
 
         def _get_version(cur: Cursor[dict[str, Any]], table: str) -> int:
-            try:
-                cur.execute(f"SELECT v FROM {table} ORDER BY v DESC LIMIT 1")
-                row = cast(dict, cur.fetchone())
-                if row is None:
-                    version = -1
-                else:
-                    version = row["v"]
-            except UndefinedTable:
-                version = -1
-                cur.execute(
-                    f"""
-                    CREATE TABLE IF NOT EXISTS {table} (
-                        v INTEGER PRIMARY KEY
-                    )
-                """
+            cur.execute(
+                f"""
+                CREATE TABLE IF NOT EXISTS {table} (
+                    v INTEGER PRIMARY KEY
                 )
+            """
+            )
+            cur.execute(f"SELECT v FROM {table} ORDER BY v DESC LIMIT 1")
+            row = cast(dict, cur.fetchone())
+            if row is None:
+                version = -1
+            else:
+                version = row["v"]
             return version
 
         with self._cursor() as cur:

libs/checkpoint-postgres/poetry.lock

@@ -13,24 +13,24 @@ files = [
 
 [[package]]
 name = "anyio"
-version = "4.6.2.post1"
+version = "4.7.0"
 description = "High level compatibility layer for multiple asynchronous event loop implementations"
 optional = false
 python-versions = ">=3.9"
 files = [
-    {file = "anyio-4.6.2.post1-py3-none-any.whl", hash = "sha256:6d170c36fba3bdd840c73d3868c1e777e33676a69c3a72cf0a0d5d6d8009b61d"},
-    {file = "anyio-4.6.2.post1.tar.gz", hash = "sha256:4c8bc31ccdb51c7f7bd251f51c609e038d63e34219b44aa86e47576389880b4c"},
+    {file = "anyio-4.7.0-py3-none-any.whl", hash = "sha256:ea60c3723ab42ba6fff7e8ccb0488c898ec538ff4df1f1d5e642c3601d07e352"},
+    {file = "anyio-4.7.0.tar.gz", hash = "sha256:2f834749c602966b7d456a7567cafcb309f96482b5081d14ac93ccd457f9dd48"},
 ]
 
 [package.dependencies]
 exceptiongroup = {version = ">=1.0.2", markers = "python_version < \"3.11\""}
 idna = ">=2.8"
 sniffio = ">=1.1"
-typing-extensions = {version = ">=4.1", markers = "python_version < \"3.11\""}
+typing_extensions = {version = ">=4.5", markers = "python_version < \"3.13\""}
 
 [package.extras]
-doc = ["Sphinx (>=7.4,<8.0)", "packaging", "sphinx-autodoc-typehints (>=1.2.0)", "sphinx-rtd-theme"]
-test = ["anyio[trio]", "coverage[toml] (>=7)", "exceptiongroup (>=1.2.0)", "hypothesis (>=4.0)", "psutil (>=5.9)", "pytest (>=7.0)", "pytest-mock (>=3.6.1)", "trustme", "truststore (>=0.9.1)", "uvloop (>=0.21.0b1)"]
+doc = ["Sphinx (>=7.4,<8.0)", "packaging", "sphinx-autodoc-typehints (>=1.2.0)", "sphinx_rtd_theme"]
+test = ["anyio[trio]", "coverage[toml] (>=7)", "exceptiongroup (>=1.2.0)", "hypothesis (>=4.0)", "psutil (>=5.9)", "pytest (>=7.0)", "pytest-mock (>=3.6.1)", "trustme", "truststore (>=0.9.1)", "uvloop (>=0.21)"]
 trio = ["trio (>=0.26.1)"]
 
 [[package]]
@@ -244,13 +244,13 @@ trio = ["trio (>=0.22.0,<1.0)"]
 
 [[package]]
 name = "httpx"
-version = "0.27.2"
+version = "0.28.0"
 description = "The next generation HTTP client."
 optional = false
 python-versions = ">=3.8"
 files = [
-    {file = "httpx-0.27.2-py3-none-any.whl", hash = "sha256:7bb2708e112d8fdd7829cd4243970f0c223274051cb35ee80c03301ee29a3df0"},
-    {file = "httpx-0.27.2.tar.gz", hash = "sha256:f7c2be1d2f3c3c3160d441802406b206c2b76f5947b11115e6df10c6c65e66c2"},
+    {file = "httpx-0.28.0-py3-none-any.whl", hash = "sha256:dc0b419a0cfeb6e8b34e85167c0da2671206f5095f1baa9663d23bcfd6b535fc"},
+    {file = "httpx-0.28.0.tar.gz", hash = "sha256:0858d3bab51ba7e386637f22a61d8ccddaeec5f3fe4209da3a6168dbb91573e0"},
 ]
 
 [package.dependencies]
@@ -258,7 +258,6 @@ anyio = "*"
 certifi = "*"
 httpcore = "==1.*"
 idna = "*"
-sniffio = "*"
 
 [package.extras]
 brotli = ["brotli", "brotlicffi"]
@@ -342,7 +341,7 @@ typing-extensions = ">=4.7"
 
 [[package]]
 name = "langgraph-checkpoint"
-version = "2.0.7"
+version = "2.0.8"
 description = "Library with base interfaces for LangGraph checkpoint savers."
 optional = false
 python-versions = "^3.9.0,<4.0"
@@ -741,13 +740,13 @@ typing-extensions = ">=4.6"
 
 [[package]]
 name = "pydantic"
-version = "2.10.2"
+version = "2.10.3"
 description = "Data validation using Python type hints"
 optional = false
 python-versions = ">=3.8"
 files = [
-    {file = "pydantic-2.10.2-py3-none-any.whl", hash = "sha256:cfb96e45951117c3024e6b67b25cdc33a3cb7b2fa62e239f7af1378358a1d99e"},
-    {file = "pydantic-2.10.2.tar.gz", hash = "sha256:2bc2d7f17232e0841cbba4641e65ba1eb6fafb3a08de3a091ff3ce14a197c4fa"},
+    {file = "pydantic-2.10.3-py3-none-any.whl", hash = "sha256:be04d85bbc7b65651c5f8e6b9976ed9c6f41782a55524cef079a34a0bb82144d"},
+    {file = "pydantic-2.10.3.tar.gz", hash = "sha256:cb5ac360ce894ceacd69c403187900a02c4b20b693a9dd1d643e1effab9eadf9"},
 ]
 
 [package.dependencies]

libs/checkpoint-postgres/pyproject.toml

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

libs/checkpoint-postgres/tests/test_async.py

@@ -63,9 +63,8 @@ async def _pipe_saver():
             prepare_threshold=0,
             row_factory=dict_row,
         ) as conn:
-            async with conn.pipeline() as pipe:
-                checkpointer = AsyncPostgresSaver(conn, pipe=pipe)
-                await checkpointer.setup()
+            checkpointer = AsyncPostgresSaver(conn)
+            await checkpointer.setup()
             async with conn.pipeline() as pipe:
                 checkpointer = AsyncPostgresSaver(conn, pipe=pipe)
                 yield checkpointer

libs/checkpoint-postgres/tests/test_async_store.py

@@ -12,7 +12,13 @@ import pytest
 from langchain_core.embeddings import Embeddings
 from psycopg import AsyncConnection
 
-from langgraph.store.base import GetOp, Item, ListNamespacesOp, PutOp, SearchOp
+from langgraph.store.base import (
+    GetOp,
+    Item,
+    ListNamespacesOp,
+    PutOp,
+    SearchOp,
+)
 from langgraph.store.postgres import AsyncPostgresStore
 from tests.conftest import (
     DEFAULT_URI,
@@ -65,14 +71,40 @@ async def store(request) -> AsyncIterator[AsyncPostgresStore]:
             await conn.execute(f"DROP DATABASE {database}")
 
 
-def test_large_batches(store: AsyncPostgresStore) -> None:
-    N = 1000
+async def test_no_running_loop(store: AsyncPostgresStore) -> None:
+    with pytest.raises(asyncio.InvalidStateError):
+        store.put(("foo", "bar"), "baz", {"val": "baz"})
+    with pytest.raises(asyncio.InvalidStateError):
+        store.get(("foo", "bar"), "baz")
+    with pytest.raises(asyncio.InvalidStateError):
+        store.delete(("foo", "bar"), "baz")
+    with pytest.raises(asyncio.InvalidStateError):
+        store.search(("foo", "bar"))
+    with pytest.raises(asyncio.InvalidStateError):
+        store.list_namespaces(prefix=("foo",))
+    with pytest.raises(asyncio.InvalidStateError):
+        store.batch([PutOp(namespace=("foo", "bar"), key="baz", value={"val": "baz"})])
+    with ThreadPoolExecutor(max_workers=1) as executor:
+        future = executor.submit(store.put, ("foo", "bar"), "baz", {"val": "baz"})
+        result = await asyncio.wrap_future(future)
+        assert result is None
+        future = executor.submit(store.get, ("foo", "bar"), "baz")
+        result = await asyncio.wrap_future(future)
+        assert result.value == {"val": "baz"}
+        result = await asyncio.wrap_future(
+            executor.submit(store.list_namespaces, prefix=("foo",))
+        )
+
+
+async def test_large_batches(request: Any, store: AsyncPostgresStore) -> None:
+    N = 100  # less important that we are performant here
     M = 10
 
     with ThreadPoolExecutor(max_workers=10) as executor:
+        futures = []
         for m in range(M):
             for i in range(N):
-                _ = [
+                futures += [
                     executor.submit(
                         store.put,
                         ("test", "foo", "bar", "baz", str(m % 2)),
@@ -107,6 +139,11 @@ def test_large_batches(store: AsyncPostgresStore) -> None:
                     ),
                 ]
 
+        results = await asyncio.gather(
+            *(asyncio.wrap_future(future) for future in futures)
+        )
+    assert len(results) == M * N * 6
+
 
 async def test_large_batches_async(store: AsyncPostgresStore) -> None:
     N = 1000
@@ -152,7 +189,8 @@ async def test_large_batches_async(store: AsyncPostgresStore) -> None:
                 )
             )
 
-    await asyncio.gather(*coros)
+    results = await asyncio.gather(*coros)
+    assert len(results) == M * N * 6
 
 
 async def test_abatch_order(store: AsyncPostgresStore) -> None:

libs/checkpoint-postgres/tests/test_store.py

@@ -1,7 +1,5 @@
 # type: ignore
 
-import asyncio
-from concurrent.futures import ThreadPoolExecutor
 from contextlib import contextmanager
 from typing import Any, Optional
 from uuid import uuid4
@@ -19,7 +17,11 @@ from langgraph.store.base import (
     SearchOp,
 )
 from langgraph.store.postgres import PostgresStore
-from tests.conftest import DEFAULT_URI, VECTOR_TYPES, CharacterEmbeddings
+from tests.conftest import (
+    DEFAULT_URI,
+    VECTOR_TYPES,
+    CharacterEmbeddings,
+)
 
 
 @pytest.fixture(scope="function", params=["default", "pipe", "pool"])
@@ -57,96 +59,6 @@ def store(request) -> PostgresStore:
             conn.execute(f"DROP DATABASE {database}")
 
 
-def test_large_batches(store: PostgresStore) -> None:
-    N = 1000
-    M = 10
-
-    with ThreadPoolExecutor(max_workers=10) as executor:
-        for m in range(M):
-            for i in range(N):
-                _ = [
-                    executor.submit(
-                        store.put,
-                        ("test", "foo", "bar", "baz", str(m % 2)),
-                        f"key{i}",
-                        value={"foo": "bar" + str(i)},
-                    ),
-                    executor.submit(
-                        store.get,
-                        ("test", "foo", "bar", "baz", str(m % 2)),
-                        f"key{i}",
-                    ),
-                    executor.submit(
-                        store.list_namespaces,
-                        prefix=None,
-                        max_depth=m + 1,
-                    ),
-                    executor.submit(
-                        store.search,
-                        ("test",),
-                    ),
-                    executor.submit(
-                        store.put,
-                        ("test", "foo", "bar", "baz", str(m % 2)),
-                        f"key{i}",
-                        value={"foo": "bar" + str(i)},
-                    ),
-                    executor.submit(
-                        store.put,
-                        ("test", "foo", "bar", "baz", str(m % 2)),
-                        f"key{i}",
-                        None,
-                    ),
-                ]
-
-
-async def test_large_batches_async(store: PostgresStore) -> None:
-    N = 1000
-    M = 10
-    coros = []
-    for m in range(M):
-        for i in range(N):
-            coros.append(
-                store.aput(
-                    ("test", "foo", "bar", "baz", str(m % 2)),
-                    f"key{i}",
-                    value={"foo": "bar" + str(i)},
-                )
-            )
-            coros.append(
-                store.aget(
-                    ("test", "foo", "bar", "baz", str(m % 2)),
-                    f"key{i}",
-                )
-            )
-            coros.append(
-                store.alist_namespaces(
-                    prefix=None,
-                    max_depth=m + 1,
-                )
-            )
-            coros.append(
-                store.asearch(
-                    ("test",),
-                )
-            )
-            coros.append(
-                store.aput(
-                    ("test", "foo", "bar", "baz", str(m % 2)),
-                    f"key{i}",
-                    value={"foo": "bar" + str(i)},
-                )
-            )
-            coros.append(
-                store.adelete(
-                    ("test", "foo", "bar", "baz", str(m % 2)),
-                    f"key{i}",
-                )
-            )
-
-    await asyncio.gather(*coros)
-
-
 def test_batch_order(store: PostgresStore) -> None:
     # Setup test data
     store.put(("test", "foo"), "key1", {"data": "value1"})

libs/checkpoint-postgres/tests/test_sync.py

@@ -57,9 +57,8 @@ def _pipe_saver():
             prepare_threshold=0,
             row_factory=dict_row,
         ) as conn:
-            with conn.pipeline() as pipe:
-                checkpointer = PostgresSaver(conn, pipe=pipe)
-                checkpointer.setup()
+            checkpointer = PostgresSaver(conn)
+            checkpointer.setup()
             with conn.pipeline() as pipe:
                 checkpointer = PostgresSaver(conn, pipe=pipe)
                 yield checkpointer

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

@@ -1,4 +1,3 @@
-import asyncio
 import logging
 import os
 import pickle
@@ -6,7 +5,6 @@ import random
 import shutil
 from collections import defaultdict
 from contextlib import AbstractAsyncContextManager, AbstractContextManager, ExitStack
-from functools import partial
 from types import TracebackType
 from typing import Any, AsyncIterator, Dict, Iterator, Optional, Sequence, Tuple, Type
 
@@ -395,9 +393,7 @@ class MemorySaver(
         Returns:
             Optional[CheckpointTuple]: The retrieved checkpoint tuple, or None if no matching checkpoint was found.
         """
-        return await asyncio.get_running_loop().run_in_executor(
-            None, self.get_tuple, config
-        )
+        return self.get_tuple(config)
 
     async def alist(
         self,
@@ -418,24 +414,8 @@ class MemorySaver(
         Yields:
             AsyncIterator[CheckpointTuple]: An asynchronous iterator of checkpoint tuples.
         """
-        loop = asyncio.get_running_loop()
-        iter = await loop.run_in_executor(
-            None,
-            partial(
-                self.list,
-                before=before,
-                limit=limit,
-                filter=filter,
-            ),
-            config,
-        )
-        while True:
-            # handling StopIteration exception inside coroutine won't work
-            # as expected, so using next() with default value to break the loop
-            if item := await loop.run_in_executor(None, next, iter, None):
-                yield item
-            else:
-                break
+        for item in self.list(config, filter=filter, before=before, limit=limit):
+            yield item
 
     async def aput(
         self,
@@ -455,9 +435,7 @@ class MemorySaver(
         Returns:
             RunnableConfig: The updated config containing the saved checkpoint's timestamp.
         """
-        return await asyncio.get_running_loop().run_in_executor(
-            None, self.put, config, checkpoint, metadata, new_versions
-        )
+        return self.put(config, checkpoint, metadata, new_versions)
 
     async def aput_writes(
         self,
@@ -474,10 +452,9 @@ class MemorySaver(
             config (RunnableConfig): The config to associate with the writes.
             writes (List[Tuple[str, Any]]): The writes to save, each as a (channel, value) pair.
             task_id (str): Identifier for the task creating the writes.
+        return self.put_writes(config, writes, task_id)
         """
-        return await asyncio.get_running_loop().run_in_executor(
-            None, self.put_writes, config, writes, task_id
-        )
+        return self.put_writes(config, writes, task_id)
 
     def get_next_version(self, current: Optional[str], channel: ChannelProtocol) -> str:
         if current is None:

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

@@ -808,8 +808,6 @@ class BaseStore(ABC):
             # [("a", "b", "c"), ("a", "b", "d"), ("a", "b", "f")]
             ```
         """
-        if max_depth is not None and max_depth <= 0:
-            raise ValueError("If provided, max_depth must be greater than 0")
         match_conditions = []
         if prefix:
             match_conditions.append(MatchCondition(match_type="prefix", path=prefix))
@@ -1006,8 +1004,6 @@ class BaseStore(ABC):
             # Returns: [("a", "b", "c"), ("a", "b", "d"), ("a", "b", "f")]
             ```
         """
-        if max_depth is not None and max_depth <= 0:
-            raise ValueError("If provided, max_depth must be greater than 0")
         match_conditions = []
         if prefix:
             match_conditions.append(MatchCondition(match_type="prefix", path=prefix))

libs/checkpoint/langgraph/store/base/batch.py

@@ -1,9 +1,7 @@
 import asyncio
-import threading
-import time
+import functools
 import weakref
-from concurrent.futures import Future
-from typing import Any, Iterable, Literal, Optional, Union
+from typing import Any, Callable, Iterable, Literal, Optional, TypeVar, Union
 
 from langgraph.store.base import (
     BaseStore,
@@ -20,13 +18,47 @@ from langgraph.store.base import (
     _validate_namespace,
 )
 
+F = TypeVar("F", bound=Callable)
 
-class AsyncBatchedBaseStoreMixin:
+
+def _check_loop(func: F) -> F:
+    @functools.wraps(func)
+    def wrapper(store: "AsyncBatchedBaseStore", *args: Any, **kwargs: Any) -> Any:
+        method_name: str = func.__name__
+        try:
+            current_loop = asyncio.get_running_loop()
+            if current_loop is store._loop:
+                replacement_str = (
+                    f"Specifically, replace `store.{method_name}(...)` with `await store.a{method_name}(...)"
+                    if method_name
+                    else "For example, replace `store.get(...)` with `await store.aget(...)`"
+                )
+                raise asyncio.InvalidStateError(
+                    f"Synchronous calls to {store.__class__.__name__} detected in the main event loop. "
+                    "This can lead to deadlocks or performance issues. "
+                    "Please use the asynchronous interface for main thread operations. "
+                    f"{replacement_str} "
+                )
+        except RuntimeError:
+            pass
+        return func(store, *args, **kwargs)
+
+    return wrapper
+
+
+class AsyncBatchedBaseStore(BaseStore):
     """Efficiently batch operations in a background task."""
 
-    _loop: asyncio.AbstractEventLoop
-    _aqueue: dict[asyncio.Future, Op]
-    _task: asyncio.Task
+    __slots__ = ("_loop", "_aqueue", "_task")
+
+    def __init__(self) -> None:
+        super().__init__()
+        self._loop = asyncio.get_running_loop()
+        self._aqueue: dict[asyncio.Future, Op] = {}
+        self._task = self._loop.create_task(_run(self._aqueue, weakref.ref(self)))
+
+    def __del__(self) -> None:
+        self._task.cancel()
 
     async def aget(
         self,
@@ -97,28 +129,81 @@ class AsyncBatchedBaseStoreMixin:
         self._aqueue[fut] = op
         return await fut
 
-
-class AsyncBatchedBaseStore(AsyncBatchedBaseStoreMixin, BaseStore):
-    """Efficiently batch operations in a background task."""
-
-    __slots__ = ("_loop", "_aqueue", "_task")
-
-    def __init__(self) -> None:
-        super().__init__()
-        self._loop = asyncio.get_running_loop()
-        self._aqueue: dict[asyncio.Future, Op] = {}
-        self._task = self._loop.create_task(_run(self._aqueue, weakref.ref(self)))
-
-    def __del__(self) -> None:
-        self._task.cancel()
-
+    @_check_loop
     def batch(self, ops: Iterable[Op]) -> list[Result]:
-        futures = []
-        for op in ops:
-            fut = self._loop.create_future()
-            self._aqueue[fut] = op
-            futures.append(fut)
-        return [fut.result() for fut in asyncio.as_completed(futures)]
+        return asyncio.run_coroutine_threadsafe(self.abatch(ops), self._loop).result()
+
+    @_check_loop
+    def get(
+        self,
+        namespace: tuple[str, ...],
+        key: str,
+    ) -> Optional[Item]:
+        return asyncio.run_coroutine_threadsafe(
+            self.aget(namespace, key=key), self._loop
+        ).result()
+
+    @_check_loop
+    def search(
+        self,
+        namespace_prefix: tuple[str, ...],
+        /,
+        *,
+        query: Optional[str] = None,
+        filter: Optional[dict[str, Any]] = None,
+        limit: int = 10,
+        offset: int = 0,
+    ) -> list[SearchItem]:
+        return asyncio.run_coroutine_threadsafe(
+            self.asearch(
+                namespace_prefix, query=query, filter=filter, limit=limit, offset=offset
+            ),
+            self._loop,
+        ).result()
+
+    @_check_loop
+    def put(
+        self,
+        namespace: tuple[str, ...],
+        key: str,
+        value: dict[str, Any],
+        index: Optional[Union[Literal[False], list[str]]] = None,
+    ) -> None:
+        _validate_namespace(namespace)
+        asyncio.run_coroutine_threadsafe(
+            self.aput(namespace, key=key, value=value, index=index), self._loop
+        ).result()
+
+    @_check_loop
+    def delete(
+        self,
+        namespace: tuple[str, ...],
+        key: str,
+    ) -> None:
+        asyncio.run_coroutine_threadsafe(
+            self.adelete(namespace, key=key), self._loop
+        ).result()
+
+    @_check_loop
+    def list_namespaces(
+        self,
+        *,
+        prefix: Optional[NamespacePath] = None,
+        suffix: Optional[NamespacePath] = None,
+        max_depth: Optional[int] = None,
+        limit: int = 100,
+        offset: int = 0,
+    ) -> list[tuple[str, ...]]:
+        return asyncio.run_coroutine_threadsafe(
+            self.alist_namespaces(
+                prefix=prefix,
+                suffix=suffix,
+                max_depth=max_depth,
+                limit=limit,
+                offset=offset,
+            ),
+            self._loop,
+        ).result()
 
 
 def _dedupe_ops(values: list[Op]) -> tuple[Optional[list[int]], list[Op]]:
@@ -165,7 +250,8 @@ def _dedupe_ops(values: list[Op]) -> tuple[Optional[list[int]], list[Op]]:
 
 
 async def _run(
-    aqueue: dict[asyncio.Future, Op], store: weakref.ReferenceType[BaseStore]
+    aqueue: dict[asyncio.Future, Op],
+    store: weakref.ReferenceType[BaseStore],
 ) -> None:
     while True:
         await asyncio.sleep(0)
@@ -195,181 +281,3 @@ async def _run(
             break
         # remove strong ref to store
         del s
-
-
-class SyncBatchedBaseStoreMixin(BaseStore):
-    """Efficiently batch operations in a background thread."""
-
-    _sync_queue: dict[Future, Op]
-    _sync_thread: threading.Thread
-
-    def get(
-        self,
-        namespace: tuple[str, ...],
-        key: str,
-    ) -> Optional[Item]:
-        fut: Future[Optional[Item]] = Future()
-        self._sync_queue[fut] = GetOp(namespace, key)
-        return fut.result()
-
-    def search(
-        self,
-        namespace_prefix: tuple[str, ...],
-        /,
-        *,
-        query: Optional[str] = None,
-        filter: Optional[dict[str, Any]] = None,
-        limit: int = 10,
-        offset: int = 0,
-    ) -> list[SearchItem]:
-        fut: Future[list[SearchItem]] = Future()
-        self._sync_queue[fut] = SearchOp(namespace_prefix, filter, limit, offset, query)
-        return fut.result()
-
-    def put(
-        self,
-        namespace: tuple[str, ...],
-        key: str,
-        value: dict[str, Any],
-        index: Optional[Union[Literal[False], list[str]]] = None,
-    ) -> None:
-        _validate_namespace(namespace)
-        fut: Future[None] = Future()
-        self._sync_queue[fut] = PutOp(namespace, key, value, index)
-        return fut.result()
-
-    def delete(
-        self,
-        namespace: tuple[str, ...],
-        key: str,
-    ) -> None:
-        fut: Future[None] = Future()
-        self._sync_queue[fut] = PutOp(namespace, key, None)
-        return fut.result()
-
-    def list_namespaces(
-        self,
-        *,
-        prefix: Optional[NamespacePath] = None,
-        suffix: Optional[NamespacePath] = None,
-        max_depth: Optional[int] = None,
-        limit: int = 100,
-        offset: int = 0,
-    ) -> list[tuple[str, ...]]:
-        fut: Future[list[tuple[str, ...]]] = Future()
-        match_conditions = []
-        if prefix:
-            match_conditions.append(MatchCondition(match_type="prefix", path=prefix))
-        if suffix:
-            match_conditions.append(MatchCondition(match_type="suffix", path=suffix))
-
-        op = ListNamespacesOp(
-            match_conditions=tuple(match_conditions),
-            max_depth=max_depth,
-            limit=limit,
-            offset=offset,
-        )
-        self._sync_queue[fut] = op
-        return fut.result()
-
-
-class SyncBatchedBaseStore(SyncBatchedBaseStoreMixin, BaseStore):
-    """Efficiently batch operations in a background thread."""
-
-    __slots__ = ("_sync_queue", "_sync_thread")
-
-    def __init__(self) -> None:
-        super().__init__()
-        self._sync_queue: dict[Future, Op] = {}
-        self._sync_thread = threading.Thread(
-            target=_sync_run,
-            args=(self._sync_queue, weakref.ref(self)),
-            daemon=True,
-        )
-        self._sync_thread.start()
-
-    def __del__(self) -> None:
-        # Signal the thread to stop
-        if self._sync_thread.is_alive():
-            empty_future: Future = Future()
-            self._sync_queue[empty_future] = None  # type: ignore
-            self._sync_thread.join()
-
-    async def abatch(self, ops: Iterable[Op]) -> list[Result]:
-        futures = []
-        for op in ops:
-            fut: Future[Result] = Future()
-            self._sync_queue[fut] = op
-            futures.append(fut)
-        return [fut.result() for fut in futures]
-
-
-class BatchedBaseStore(
-    AsyncBatchedBaseStoreMixin, SyncBatchedBaseStoreMixin, BaseStore
-):
-    __slots__ = (
-        "_sync_queue",
-        "_sync_thread",
-        "_task",
-        "_loop",
-        "_aqueue",
-    )
-
-    def __init__(self) -> None:
-        super().__init__()
-        # Setup async processing
-        self._loop = asyncio.get_running_loop()
-        self._aqueue: dict[asyncio.Future, Op] = {}
-        self._task = self._loop.create_task(_run(self._aqueue, weakref.ref(self)))
-
-        self._sync_queue: dict[Future, Op] = {}
-        self._sync_thread = threading.Thread(
-            target=_sync_run,
-            args=(self._sync_queue, weakref.ref(self)),
-            daemon=True,
-        )
-        self._sync_thread.start()
-
-    def __del__(self) -> None:
-        # Signal the thread to stop
-        if self._sync_thread.is_alive():
-            empty_future: Future[None] = Future()
-            self._sync_queue[empty_future] = None  # type: ignore
-            self._sync_thread.join()
-
-        # Signal the thread to stop
-        if self._task is not None:
-            self._task.cancel()
-
-
-def _sync_run(queue: dict[Future, Op], store: weakref.ReferenceType[BaseStore]) -> None:
-    while True:
-        time.sleep(0.001)  # Yield to other threads
-        if not queue:
-            continue
-        if s := store():
-            # get the operations to run
-            taken = queue.copy()
-            # action each operation
-            try:
-                values = list(taken.values())
-                if None in values:  # Exit signal
-                    break
-                listen, dedupped = _dedupe_ops(values)
-                results = s.batch(dedupped)  # Note: Using sync batch here
-                if listen is not None:
-                    results = [results[ix] for ix in listen]
-
-                # set the results of each operation
-                for fut, result in zip(taken, results):
-                    fut.set_result(result)
-            except Exception as e:
-                for fut in taken:
-                    fut.set_exception(e)
-            # remove the operations from the queue
-            for fut in taken:
-                del queue[fut]
-        else:
-            break
-        # remove strong ref to store
-        del s

libs/langgraph/langgraph/channels/base.py

@@ -1,5 +1,5 @@
 from abc import ABC, abstractmethod
-from typing import Any, Generic, Optional, Sequence, Type, TypeVar
+from typing import Any, Generic, Optional, Sequence, TypeVar
 
 from typing_extensions import Self
 
@@ -13,7 +13,7 @@ C = TypeVar("C")
 class BaseChannel(Generic[Value, Update, C], ABC):
     __slots__ = ("key", "typ")
 
-    def __init__(self, typ: Type[Any], key: str = "") -> None:
+    def __init__(self, typ: Any, key: str = "") -> None:
         self.typ = typ
         self.key = key
 

libs/langgraph/langgraph/constants.py

@@ -40,12 +40,16 @@ SCHEDULED = sys.intern("__scheduled__")
 # marker to signal node was scheduled (in distributed mode)
 TASKS = sys.intern("__pregel_tasks")
 # for Send objects returned by nodes/edges, corresponds to PUSH below
+RETURN = sys.intern("__return__")
+# for writes of a task where we simply record the return value
 
 # --- Reserved config.configurable keys ---
 CONFIG_KEY_SEND = sys.intern("__pregel_send")
 # holds the `write` function that accepts writes to state/edges/reserved keys
 CONFIG_KEY_READ = sys.intern("__pregel_read")
 # holds the `read` function that returns a copy of the current state
+CONFIG_KEY_CALL = sys.intern("__pregel_call")
+# holds the `call` function that accepts a node/func, args and returns a future
 CONFIG_KEY_CHECKPOINTER = sys.intern("__pregel_checkpointer")
 # holds a `BaseCheckpointSaver` passed from parent graph to child graphs
 CONFIG_KEY_STREAM = sys.intern("__pregel_stream")

libs/langgraph/langgraph/func/__init__.py

@@ -0,0 +1,120 @@
+import asyncio
+import concurrent
+import concurrent.futures
+import inspect
+import types
+from functools import partial, update_wrapper
+from typing import (
+    Any,
+    Awaitable,
+    Callable,
+    Optional,
+    TypeVar,
+    Union,
+    overload,
+)
+
+from typing_extensions import ParamSpec
+
+from langgraph.channels.ephemeral_value import EphemeralValue
+from langgraph.channels.last_value import LastValue
+from langgraph.checkpoint.base import BaseCheckpointSaver
+from langgraph.constants import END, START, TAG_HIDDEN
+from langgraph.pregel import Pregel
+from langgraph.pregel.call import get_runnable_for_func
+from langgraph.pregel.read import PregelNode
+from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry
+from langgraph.store.base import BaseStore
+from langgraph.types import RetryPolicy, StreamMode, StreamWriter
+
+P = ParamSpec("P")
+P1 = TypeVar("P1")
+T = TypeVar("T")
+
+
+def call(
+    func: Callable[[P1], T],
+    input: P1,
+    *,
+    retry: Optional[RetryPolicy] = None,
+) -> concurrent.futures.Future[T]:
+    from langgraph.constants import CONFIG_KEY_CALL
+    from langgraph.utils.config import get_configurable
+
+    conf = get_configurable()
+    impl = conf[CONFIG_KEY_CALL]
+    fut = impl(func, input, retry=retry)
+    return fut
+
+
+@overload
+def task(
+    *, retry: Optional[RetryPolicy] = None
+) -> Callable[[Callable[P, Awaitable[T]]], Callable[P, asyncio.Future[T]]]: ...
+
+
+@overload
+def task(  # type: ignore[overload-cannot-match]
+    *, retry: Optional[RetryPolicy] = None
+) -> Callable[[Callable[P, T]], Callable[P, concurrent.futures.Future[T]]]: ...
+
+
+def task(
+    *, retry: Optional[RetryPolicy] = None
+) -> Union[
+    Callable[[Callable[P, Awaitable[T]]], Callable[P, asyncio.Future[T]]],
+    Callable[[Callable[P, T]], Callable[P, concurrent.futures.Future[T]]],
+]:
+    def _task(func: Callable[P, T]) -> Callable[P, concurrent.futures.Future[T]]:
+        return update_wrapper(partial(call, func, retry=retry), func)
+
+    return _task
+
+
+def entrypoint(
+    *,
+    checkpointer: Optional[BaseCheckpointSaver] = None,
+    store: Optional[BaseStore] = None,
+) -> Callable[[types.FunctionType], Pregel]:
+    def _imp(func: types.FunctionType) -> Pregel:
+        if inspect.isgeneratorfunction(func):
+
+            def gen_wrapper(*args: Any, writer: StreamWriter, **kwargs: Any) -> Any:
+                for chunk in func(*args, **kwargs):
+                    writer(chunk)
+
+            bound = get_runnable_for_func(gen_wrapper)
+            stream_mode: StreamMode = "custom"
+        elif inspect.isasyncgenfunction(func):
+
+            async def agen_wrapper(
+                *args: Any, writer: StreamWriter, **kwargs: Any
+            ) -> Any:
+                async for chunk in func(*args, **kwargs):
+                    writer(chunk)
+
+            bound = get_runnable_for_func(agen_wrapper)
+            stream_mode = "custom"
+        else:
+            bound = get_runnable_for_func(func)
+            stream_mode = "updates"
+
+        return Pregel(
+            nodes={
+                func.__name__: PregelNode(
+                    bound=bound,
+                    triggers=[START],
+                    channels=[START],
+                    writers=[ChannelWrite([ChannelWriteEntry(END)], tags=[TAG_HIDDEN])],
+                )
+            },
+            channels={START: EphemeralValue(Any), END: LastValue(Any, END)},
+            input_channels=START,
+            output_channels=END,
+            stream_channels=END,
+            stream_mode=stream_mode,
+            checkpointer=checkpointer,
+            store=store,
+        )
+
+    return _imp

libs/langgraph/langgraph/graph/state.py

@@ -559,6 +559,7 @@ class StateGraph(Graph):
         for key, node in self.nodes.items():
             compiled.attach_node(key, node)
 
+        compiled.attach_branch(START, SELF, CONTROL_BRANCH, with_reader=False)
         for key, node in self.nodes.items():
             compiled.attach_branch(key, SELF, CONTROL_BRANCH, with_reader=False)
 
@@ -613,21 +614,24 @@ class CompiledStateGraph(CompiledGraph):
             ]
 
         def _get_root(input: Any) -> Optional[Sequence[tuple[str, Any]]]:
-            if (
-                isinstance(input, (list, tuple))
-                and input
-                and all(isinstance(i, Command) for i in input)
-            ):
-                updates: list[tuple[str, Any]] = []
-                for i in input:
-                    if i.graph == Command.PARENT:
-                        continue
-                    updates.extend(i._update_as_tuples())
-                return updates
-            elif isinstance(input, Command):
+            if isinstance(input, Command):
                 if input.graph == Command.PARENT:
                     return ()
                 return input._update_as_tuples()
+            elif (
+                isinstance(input, (list, tuple))
+                and input
+                and any(isinstance(i, Command) for i in input)
+            ):
+                updates: list[tuple[str, Any]] = []
+                for i in input:
+                    if isinstance(i, Command):
+                        if i.graph == Command.PARENT:
+                            continue
+                        updates.extend(i._update_as_tuples())
+                    else:
+                        updates.append(("__root__", i))
+                return updates
             elif input is not None:
                 return [("__root__", input)]
 
@@ -645,13 +649,16 @@ class CompiledStateGraph(CompiledGraph):
             elif (
                 isinstance(input, (list, tuple))
                 and input
-                and all(isinstance(i, Command) for i in input)
+                and any(isinstance(i, Command) for i in input)
             ):
                 updates: list[tuple[str, Any]] = []
                 for i in input:
-                    if i.graph == Command.PARENT:
-                        continue
-                    updates.extend(i._update_as_tuples())
+                    if isinstance(i, Command):
+                        if i.graph == Command.PARENT:
+                            continue
+                        updates.extend(i._update_as_tuples())
+                    else:
+                        updates.extend(_get_updates(i) or ())
                 return updates
             elif get_type_hints(type(input)):
                 return [

libs/langgraph/langgraph/prebuilt/tool_node.py

@@ -1,7 +1,7 @@
 import asyncio
 import inspect
 import json
-from copy import copy
+from copy import copy, deepcopy
 from typing import (
     Any,
     Callable,
@@ -20,6 +20,7 @@ from langchain_core.messages import (
     AnyMessage,
     ToolCall,
     ToolMessage,
+    convert_to_messages,
 )
 from langchain_core.runnables import RunnableConfig
 from langchain_core.runnables.config import (
@@ -35,6 +36,7 @@ from typing_extensions import Annotated, get_args, get_origin
 
 from langgraph.errors import GraphBubbleUp
 from langgraph.store.base import BaseStore
+from langgraph.types import Command
 from langgraph.utils.runnable import RunnableCallable
 
 INVALID_TOOL_NAME_ERROR_TEMPLATE = (
@@ -47,7 +49,7 @@ def msg_content_output(output: Any) -> Union[str, list[dict]]:
     recognized_content_block_types = ("image", "image_url", "text", "json")
     if isinstance(output, str):
         return output
-    elif all(
+    elif isinstance(output, list) and all(
         [
             isinstance(x, dict) and x.get("type") in recognized_content_block_types
             for x in output
@@ -210,12 +212,31 @@ class ToolNode(RunnableCallable):
         *,
         store: BaseStore,
     ) -> Any:
-        tool_calls, output_type = self._parse_input(input, store)
+        tool_calls, input_type = self._parse_input(input, store)
         config_list = get_config_list(config, len(tool_calls))
+        input_types = [input_type] * len(tool_calls)
         with get_executor_for_config(config) as executor:
-            outputs = [*executor.map(self._run_one, tool_calls, config_list)]
-        # TypedDict, pydantic, dataclass, etc. should all be able to load from dict
-        return outputs if output_type == "list" else {self.messages_key: outputs}
+            outputs = [
+                *executor.map(self._run_one, tool_calls, input_types, config_list)
+            ]
+
+        # preserve existing behavior for non-command tool outputs for backwards compatibility
+        if not any(isinstance(output, Command) for output in outputs):
+            # TypedDict, pydantic, dataclass, etc. should all be able to load from dict
+            return outputs if input_type == "list" else {self.messages_key: outputs}
+
+        # LangGraph will automatically handle list of Command and non-command node updates
+        combined_outputs: list[
+            Command | list[ToolMessage] | dict[str, list[ToolMessage]]
+        ] = []
+        for output in outputs:
+            if isinstance(output, Command):
+                combined_outputs.append(output)
+            else:
+                combined_outputs.append(
+                    [output] if input_type == "list" else {self.messages_key: [output]}
+                )
+        return combined_outputs
 
     def invoke(
         self, input: Input, config: Optional[RunnableConfig] = None, **kwargs: Any
@@ -242,26 +263,97 @@ class ToolNode(RunnableCallable):
         *,
         store: BaseStore,
     ) -> Any:
-        tool_calls, output_type = self._parse_input(input, store)
+        tool_calls, input_type = self._parse_input(input, store)
         outputs = await asyncio.gather(
-            *(self._arun_one(call, config) for call in tool_calls)
+            *(self._arun_one(call, input_type, config) for call in tool_calls)
         )
-        # TypedDict, pydantic, dataclass, etc. should all be able to load from dict
-        return outputs if output_type == "list" else {self.messages_key: outputs}
 
-    def _run_one(self, call: ToolCall, config: RunnableConfig) -> ToolMessage:
+        # preserve existing behavior for non-command tool outputs for backwards compatibility
+        if not any(isinstance(output, Command) for output in outputs):
+            # TypedDict, pydantic, dataclass, etc. should all be able to load from dict
+            return outputs if input_type == "list" else {self.messages_key: outputs}
+
+        # LangGraph will automatically handle list of Command and non-command node updates
+        combined_outputs: list[
+            Command | list[ToolMessage] | dict[str, list[ToolMessage]]
+        ] = []
+        for output in outputs:
+            if isinstance(output, Command):
+                combined_outputs.append(output)
+            else:
+                combined_outputs.append(
+                    [output] if input_type == "list" else {self.messages_key: [output]}
+                )
+        return combined_outputs
+
+    def _run_one(
+        self,
+        call: ToolCall,
+        input_type: Literal["list", "dict"],
+        config: RunnableConfig,
+    ) -> ToolMessage:
         if invalid_tool_message := self._validate_tool_call(call):
             return invalid_tool_message
 
         try:
             input = {**call, **{"type": "tool_call"}}
-            tool_message: ToolMessage = self.tools_by_name[call["name"]].invoke(
-                input, config
+            response = self.tools_by_name[call["name"]].invoke(input)
+
+        # GraphInterrupt is a special exception that will always be raised.
+        # It can be triggered in the following scenarios:
+        # (1) a NodeInterrupt is raised inside a tool
+        # (2) a NodeInterrupt is raised inside a graph node for a graph called as a tool
+        # (3) a GraphInterrupt is raised when a subgraph is interrupted inside a graph called as a tool
+        # (2 and 3 can happen in a "supervisor w/ tools" multi-agent architecture)
+        except GraphBubbleUp as e:
+            raise e
+        except Exception as e:
+            if isinstance(self.handle_tool_errors, tuple):
+                handled_types: tuple = self.handle_tool_errors
+            elif callable(self.handle_tool_errors):
+                handled_types = _infer_handled_types(self.handle_tool_errors)
+            else:
+                # default behavior is catching all exceptions
+                handled_types = (Exception,)
+
+            # Unhandled
+            if not self.handle_tool_errors or not isinstance(e, handled_types):
+                raise e
+            # Handled
+            else:
+                content = _handle_tool_error(e, flag=self.handle_tool_errors)
+            return ToolMessage(
+                content=content,
+                name=call["name"],
+                tool_call_id=call["id"],
+                status="error",
             )
-            tool_message.content = cast(
-                Union[str, list], msg_content_output(tool_message.content)
+
+        if isinstance(response, Command):
+            return self._validate_tool_command(response, call, input_type)
+        elif isinstance(response, ToolMessage):
+            response.content = cast(
+                Union[str, list], msg_content_output(response.content)
             )
-            return tool_message
+            return response
+        else:
+            raise TypeError(
+                f"Tool {call['name']} returned unexpected type: {type(response)}"
+            )
+
+    async def _arun_one(
+        self,
+        call: ToolCall,
+        input_type: Literal["list", "dict"],
+        config: RunnableConfig,
+    ) -> ToolMessage:
+        if invalid_tool_message := self._validate_tool_call(call):
+            return invalid_tool_message
+
+        try:
+            input = {**call, **{"type": "tool_call"}}
+            response = await self.tools_by_name[call["name"]].ainvoke(input)
+
         # GraphInterrupt is a special exception that will always be raised.
         # It can be triggered in the following scenarios:
         # (1) a NodeInterrupt is raised inside a tool
@@ -286,50 +378,24 @@ class ToolNode(RunnableCallable):
             else:
                 content = _handle_tool_error(e, flag=self.handle_tool_errors)
 
-        return ToolMessage(
-            content=content, name=call["name"], tool_call_id=call["id"], status="error"
-        )
-
-    async def _arun_one(self, call: ToolCall, config: RunnableConfig) -> ToolMessage:
-        if invalid_tool_message := self._validate_tool_call(call):
-            return invalid_tool_message
-
-        try:
-            input = {**call, **{"type": "tool_call"}}
-            tool_message: ToolMessage = await self.tools_by_name[call["name"]].ainvoke(
-                input, config
+            return ToolMessage(
+                content=content,
+                name=call["name"],
+                tool_call_id=call["id"],
+                status="error",
             )
-            tool_message.content = cast(
-                Union[str, list], msg_content_output(tool_message.content)
+
+        if isinstance(response, Command):
+            return self._validate_tool_command(response, call, input_type)
+        elif isinstance(response, ToolMessage):
+            response.content = cast(
+                Union[str, list], msg_content_output(response.content)
+            )
+            return response
+        else:
+            raise TypeError(
+                f"Tool {call['name']} returned unexpected type: {type(response)}"
             )
-            return tool_message
-        # GraphInterrupt is a special exception that will always be raised.
-        # It can be triggered in the following scenarios:
-        # (1) a NodeInterrupt is raised inside a tool
-        # (2) a NodeInterrupt is raised inside a graph node for a graph called as a tool
-        # (3) a GraphInterrupt is raised when a subgraph is interrupted inside a graph called as a tool
-        # (2 and 3 can happen in a "supervisor w/ tools" multi-agent architecture)
-        except GraphBubbleUp as e:
-            raise e
-        except Exception as e:
-            if isinstance(self.handle_tool_errors, tuple):
-                handled_types: tuple = self.handle_tool_errors
-            elif callable(self.handle_tool_errors):
-                handled_types = _infer_handled_types(self.handle_tool_errors)
-            else:
-                # default behavior is catching all exceptions
-                handled_types = (Exception,)
-
-            # Unhandled
-            if not self.handle_tool_errors or not isinstance(e, handled_types):
-                raise e
-            # Handled
-            else:
-                content = _handle_tool_error(e, flag=self.handle_tool_errors)
-
-        return ToolMessage(
-            content=content, name=call["name"], tool_call_id=call["id"], status="error"
-        )
 
     def _parse_input(
         self,
@@ -341,14 +407,14 @@ class ToolNode(RunnableCallable):
         store: BaseStore,
     ) -> Tuple[list[ToolCall], Literal["list", "dict"]]:
         if isinstance(input, list):
-            output_type = "list"
+            input_type = "list"
             message: AnyMessage = input[-1]
         elif isinstance(input, dict) and (messages := input.get(self.messages_key, [])):
-            output_type = "dict"
+            input_type = "dict"
             message = messages[-1]
         elif messages := getattr(input, self.messages_key, None):
             # Assume dataclass-like state that can coerce from dict
-            output_type = "dict"
+            input_type = "dict"
             message = messages[-1]
         else:
             raise ValueError("No message found in input")
@@ -359,7 +425,7 @@ class ToolNode(RunnableCallable):
         tool_calls = [
             self._inject_tool_args(call, input, store) for call in message.tool_calls
         ]
-        return tool_calls, output_type
+        return tool_calls, input_type
 
     def _validate_tool_call(self, call: ToolCall) -> Optional[ToolMessage]:
         if (requested_tool := call["name"]) not in self.tools_by_name:
@@ -453,6 +519,67 @@ class ToolNode(RunnableCallable):
         tool_call_with_store = self._inject_store(tool_call_with_state, store)
         return tool_call_with_store
 
+    def _validate_tool_command(
+        self, command: Command, call: ToolCall, input_type: Literal["list", "dict"]
+    ) -> Command:
+        if isinstance(command.update, dict):
+            # input type is dict when ToolNode is invoked with a dict input (e.g. {"messages": [AIMessage(..., tool_calls=[...])]})
+            if input_type != "dict":
+                raise ValueError(
+                    f"Tools can provide a dict in Command.update only when using dict with '{self.messages_key}' key as ToolNode input, "
+                    f"got: {command.update} for tool '{call['name']}'"
+                )
+
+            updated_command = deepcopy(command)
+            state_update = cast(dict[str, Any], updated_command.update) or {}
+            messages_update = state_update.get(self.messages_key, [])
+        elif isinstance(command.update, list):
+            # input type is list when ToolNode is invoked with a list input (e.g. [AIMessage(..., tool_calls=[...])])
+            if input_type != "list":
+                raise ValueError(
+                    f"Tools can provide a list of messages in Command.update only when using list of messages as ToolNode input, "
+                    f"got: {command.update} for tool '{call['name']}'"
+                )
+
+            updated_command = deepcopy(command)
+            messages_update = updated_command.update
+        else:
+            return command
+
+        # convert to message objects if updates are in a dict format
+        messages_update = convert_to_messages(messages_update)
+        have_seen_tool_messages = False
+        for message in messages_update:
+            if not isinstance(message, ToolMessage):
+                continue
+
+            if have_seen_tool_messages:
+                raise ValueError(
+                    f"Expected at most one ToolMessage in Command.update for tool '{call['name']}', got multiple: {messages_update}."
+                )
+
+            if message.tool_call_id != call["id"]:
+                raise ValueError(
+                    f"ToolMessage.tool_call_id must match the tool call id. Expected: {call['id']}, got: {message.tool_call_id} for tool '{call['name']}'."
+                )
+
+            message.name = call["name"]
+            have_seen_tool_messages = True
+
+        # validate that we always have exactly one ToolMessage in Command.update if command is sent to the CURRENT graph
+        if updated_command.graph is None and not have_seen_tool_messages:
+            example_update = (
+                '`Command(update={"messages": [ToolMessage("Success", tool_call_id=tool_call_id), ...]}, ...)`'
+                if input_type == "dict"
+                else '`Command(update=[ToolMessage("Success", tool_call_id=tool_call_id), ...], ...)`'
+            )
+            raise ValueError(
+                f"Expected exactly one message (ToolMessage) in Command.update for tool '{call['name']}', got: {messages_update}. "
+                "Every tool call (LLM requesting to call a tool) in the message history MUST have a corresponding ToolMessage. "
+                f"You can fix it by modifying the tool to return {example_update}."
+            )
+        return updated_command
+
 
 def tools_condition(
     state: Union[list[AnyMessage], dict[str, Any], BaseModel],

libs/langgraph/langgraph/pregel/__init__.py

@@ -18,7 +18,6 @@ from typing import (
     Type,
     Union,
     cast,
-    get_type_hints,
     overload,
 )
 from uuid import UUID, uuid5
@@ -117,6 +116,7 @@ from langgraph.utils.config import (
     patch_config,
     patch_configurable,
 )
+from langgraph.utils.fields import get_enhanced_type_hints
 from langgraph.utils.pydantic import create_model
 from langgraph.utils.queue import AsyncQueue, SyncQueue  # type: ignore[attr-defined]
 
@@ -319,8 +319,15 @@ class Pregel(PregelProtocol):
                 )
                 + (
                     [
-                        ConfigurableFieldSpec(id=name, annotation=typ)
-                        for name, typ in get_type_hints(self.config_type).items()
+                        ConfigurableFieldSpec(
+                            id=name,
+                            annotation=typ,
+                            default=default,
+                            description=description,
+                        )
+                        for name, typ, default, description in get_enhanced_type_hints(
+                            self.config_type
+                        )
                     ]
                     if self.config_type is not None
                     else []

libs/langgraph/langgraph/pregel/algo.py

@@ -43,6 +43,7 @@ from langgraph.constants import (
     CONFIG_KEY_TASK_ID,
     CONFIG_KEY_WRITES,
     EMPTY_SEQ,
+    ERROR,
     INTERRUPT,
     NO_WRITES,
     NS_END,
@@ -52,18 +53,26 @@ from langgraph.constants import (
     PUSH,
     RESERVED,
     RESUME,
+    RETURN,
     TAG_HIDDEN,
     TASKS,
     Send,
 )
 from langgraph.errors import EmptyChannelError, InvalidUpdateError
 from langgraph.managed.base import ManagedValueMapping
+from langgraph.pregel.call import get_runnable_for_func
 from langgraph.pregel.io import read_channel, read_channels
 from langgraph.pregel.log import logger
 from langgraph.pregel.manager import ChannelsManager
 from langgraph.pregel.read import PregelNode
 from langgraph.store.base import BaseStore
-from langgraph.types import All, LoopProtocol, PregelExecutableTask, PregelTask
+from langgraph.types import (
+    All,
+    LoopProtocol,
+    PregelExecutableTask,
+    PregelTask,
+    RetryPolicy,
+)
 from langgraph.utils.config import merge_configs, patch_config
 
 GetNextVersion = Callable[[Optional[V], BaseChannel], V]
@@ -97,6 +106,21 @@ class PregelTaskWrites(NamedTuple):
     triggers: Sequence[str]
 
 
+class Call:
+    __slots__ = ("func", "input", "retry")
+
+    func: Callable
+    input: Any
+    retry: Optional[RetryPolicy]
+
+    def __init__(
+        self, func: Callable, input: Any, *, retry: Optional[RetryPolicy]
+    ) -> None:
+        self.func = func
+        self.input = input
+        self.retry = retry
+
+
 def should_interrupt(
     checkpoint: Checkpoint,
     interrupt_nodes: Union[All, Sequence[str]],
@@ -179,7 +203,7 @@ def local_write(
     """Function injected under CONFIG_KEY_SEND in task config, to write to channels.
     Validates writes and forwards them to `commit` function."""
     for chan, value in writes:
-        if chan in (PUSH, TASKS):
+        if chan in (PUSH, TASKS) and value is not None:
             if not isinstance(value, Send):
                 raise InvalidUpdateError(f"Expected Send, got {value}")
             if value.node not in process_keys:
@@ -247,7 +271,7 @@ def apply_writes(
     pending_writes_by_managed: dict[str, list[Any]] = defaultdict(list)
     for task in tasks:
         for chan, val in task.writes:
-            if chan in (NO_WRITES, PUSH, RESUME, INTERRUPT):
+            if chan in (NO_WRITES, PUSH, RESUME, INTERRUPT, RETURN, ERROR):
                 pass
             elif chan == TASKS:  # TODO: remove branch in 1.0
                 checkpoint["pending_sends"].append(val)
@@ -438,7 +462,7 @@ def prepare_next_tasks(
 
 
 def prepare_single_task(
-    task_path: tuple[Union[str, int, tuple], ...],
+    task_path: tuple[Any, ...],
     task_id_checksum: Optional[str],
     *,
     checkpoint: Checkpoint,
@@ -459,7 +483,94 @@ def prepare_single_task(
     configurable = config.get(CONF, {})
     parent_ns = configurable.get(CONFIG_KEY_CHECKPOINT_NS, "")
 
-    if task_path[0] == PUSH:
+    if task_path[0] == PUSH and isinstance(task_path[-1], Call):
+        # (PUSH, parent task path, idx of PUSH write, id of parent task, Call)
+        task_path_t = cast(tuple[str, tuple, int, str, Call], task_path)
+        call = task_path_t[-1]
+        proc_ = get_runnable_for_func(call.func)
+        name = proc_.name
+        if name is None:
+            raise ValueError("`call` functions must have a `__name__` attribute")
+        # create task id
+        triggers = [PUSH]
+        checkpoint_ns = f"{parent_ns}{NS_SEP}{name}" if parent_ns else name
+        task_id = _uuid5_str(
+            checkpoint_id,
+            checkpoint_ns,
+            str(step),
+            name,
+            PUSH,
+            _tuple_str(task_path[1]),
+            str(task_path[2]),
+        )
+        task_checkpoint_ns = f"{checkpoint_ns}:{task_id}"
+        metadata = {
+            "langgraph_step": step,
+            "langgraph_node": name,
+            "langgraph_triggers": triggers,
+            "langgraph_path": task_path[:3],
+            "langgraph_checkpoint_ns": task_checkpoint_ns,
+        }
+        if task_id_checksum is not None:
+            assert task_id == task_id_checksum, f"{task_id} != {task_id_checksum}"
+        if for_execution:
+            writes: deque[tuple[str, Any]] = deque()
+            return PregelExecutableTask(
+                name,
+                call.input,
+                proc_,
+                writes,
+                patch_config(
+                    merge_configs(config, {"metadata": metadata}),
+                    run_name=name,
+                    callbacks=(
+                        manager.get_child(f"graph:step:{step}") if manager else None
+                    ),
+                    configurable={
+                        CONFIG_KEY_TASK_ID: task_id,
+                        # deque.extend is thread-safe
+                        CONFIG_KEY_SEND: partial(
+                            local_write,
+                            writes.extend,
+                            processes.keys(),
+                        ),
+                        CONFIG_KEY_READ: partial(
+                            local_read,
+                            step,
+                            checkpoint,
+                            channels,
+                            managed,
+                            PregelTaskWrites(task_path[:3], name, writes, triggers),
+                            config,
+                        ),
+                        CONFIG_KEY_STORE: (store or configurable.get(CONFIG_KEY_STORE)),
+                        CONFIG_KEY_CHECKPOINTER: (
+                            checkpointer or configurable.get(CONFIG_KEY_CHECKPOINTER)
+                        ),
+                        CONFIG_KEY_CHECKPOINT_MAP: {
+                            **configurable.get(CONFIG_KEY_CHECKPOINT_MAP, {}),
+                            parent_ns: checkpoint["id"],
+                        },
+                        CONFIG_KEY_CHECKPOINT_ID: None,
+                        CONFIG_KEY_CHECKPOINT_NS: task_checkpoint_ns,
+                        CONFIG_KEY_WRITES: [
+                            w
+                            for w in pending_writes
+                            + configurable.get(CONFIG_KEY_WRITES, [])
+                            if w[0] in (NULL_TASK_ID, task_id)
+                        ],
+                        CONFIG_KEY_SCRATCHPAD: {},
+                    },
+                ),
+                triggers,
+                call.retry,
+                None,
+                task_id,
+                task_path[:3],
+            )
+        else:
+            return PregelTask(task_id, name, task_path[:3])
+    elif task_path[0] == PUSH:
         if len(task_path) == 2:  # TODO: remove branch in 1.0
             # legacy SEND tasks, executed in superstep n+1
             # (PUSH, idx of pending send)
@@ -490,17 +601,19 @@ def prepare_single_task(
                 PUSH,
                 str(idx),
             )
-        elif len(task_path) == 4:
+        elif len(task_path) >= 4:
             # new PUSH tasks, executed in superstep n
             # (PUSH, parent task path, idx of PUSH write, id of parent task)
-            task_path_t = cast(tuple[str, tuple, int, str], task_path)
-            writes_for_path = [w for w in pending_writes if w[0] == task_path_t[3]]
-            if task_path_t[2] >= len(writes_for_path):
+            task_path_tt = cast(tuple[str, tuple, int, str], task_path)
+            writes_for_path = [w for w in pending_writes if w[0] == task_path_tt[3]]
+            if task_path_tt[2] >= len(writes_for_path):
                 logger.warning(
                     f"Ignoring invalid write index {task_path[2]} in pending writes"
                 )
                 return
-            packet = writes_for_path[task_path_t[2]][2]
+            packet = writes_for_path[task_path_tt[2]][2]
+            if packet is None:
+                return
             if not isinstance(packet, Send):
                 logger.warning(
                     f"Ignoring invalid packet type {type(packet)} in pending writes"
@@ -533,7 +646,7 @@ def prepare_single_task(
             "langgraph_step": step,
             "langgraph_node": packet.node,
             "langgraph_triggers": triggers,
-            "langgraph_path": task_path,
+            "langgraph_path": task_path[:3],
             "langgraph_checkpoint_ns": task_checkpoint_ns,
         }
         if task_id_checksum is not None:
@@ -543,7 +656,7 @@ def prepare_single_task(
             if node := proc.node:
                 if proc.metadata:
                     metadata.update(proc.metadata)
-                writes: deque[tuple[str, Any]] = deque()
+                writes = deque()
                 return PregelExecutableTask(
                     packet.node,
                     packet.arg,
@@ -572,7 +685,7 @@ def prepare_single_task(
                                 channels,
                                 managed,
                                 PregelTaskWrites(
-                                    task_path, packet.node, writes, triggers
+                                    task_path[:3], packet.node, writes, triggers
                                 ),
                                 config,
                             ),
@@ -602,12 +715,11 @@ def prepare_single_task(
                     proc.retry_policy,
                     None,
                     task_id,
-                    task_path,
+                    task_path[:3],
                     writers=proc.flat_writers,
                 )
-
         else:
-            return PregelTask(task_id, packet.node, task_path)
+            return PregelTask(task_id, packet.node, task_path[:3])
     elif task_path[0] == PULL:
         # (PULL, node name)
         name = cast(str, task_path[1])
@@ -657,7 +769,7 @@ def prepare_single_task(
                 "langgraph_step": step,
                 "langgraph_node": name,
                 "langgraph_triggers": triggers,
-                "langgraph_path": task_path,
+                "langgraph_path": task_path[:3],
                 "langgraph_checkpoint_ns": task_checkpoint_ns,
             }
             if task_id_checksum is not None:
@@ -696,7 +808,9 @@ def prepare_single_task(
                                     checkpoint,
                                     channels,
                                     managed,
-                                    PregelTaskWrites(task_path, name, writes, triggers),
+                                    PregelTaskWrites(
+                                        task_path[:3], name, writes, triggers
+                                    ),
                                     config,
                                 ),
                                 CONFIG_KEY_STORE: (
@@ -725,11 +839,11 @@ def prepare_single_task(
                         proc.retry_policy,
                         None,
                         task_id,
-                        task_path,
+                        task_path[:3],
                         writers=proc.flat_writers,
                     )
             else:
-                return PregelTask(task_id, name, task_path)
+                return PregelTask(task_id, name, task_path[:3])
 
 
 def _proc_input(

libs/langgraph/langgraph/pregel/call.py

@@ -0,0 +1,123 @@
+import sys
+import types
+from typing import Any, Callable, Optional
+
+from langgraph.constants import RETURN
+from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry
+from langgraph.utils.runnable import RunnableSeq, coerce_to_runnable
+
+"""
+Utilities borrowed from cloudpickle.
+https://github.com/cloudpipe/cloudpickle/blob/6220b0ce83ffee5e47e06770a1ee38ca9e47c850/cloudpickle/cloudpickle.py#L265
+"""
+
+
+def _getattribute(obj: Any, name: str) -> Any:
+    for subpath in name.split("."):
+        if subpath == "<locals>":
+            raise AttributeError(
+                "Can't get local attribute {!r} on {!r}".format(name, obj)
+            )
+        try:
+            parent = obj
+            obj = getattr(obj, subpath)
+        except AttributeError:
+            raise AttributeError(
+                "Can't get attribute {!r} on {!r}".format(name, obj)
+            ) from None
+    return obj, parent
+
+
+def _whichmodule(obj: Any, name: str) -> Optional[str]:
+    """Find the module an object belongs to.
+
+    This function differs from ``pickle.whichmodule`` in two ways:
+    - it does not mangle the cases where obj's module is __main__ and obj was
+      not found in any module.
+    - Errors arising during module introspection are ignored, as those errors
+      are considered unwanted side effects.
+    """
+    module_name = getattr(obj, "__module__", None)
+
+    if module_name is not None:
+        return module_name
+    # Protect the iteration by using a copy of sys.modules against dynamic
+    # modules that trigger imports of other modules upon calls to getattr or
+    # other threads importing at the same time.
+    for module_name, module in sys.modules.copy().items():
+        # Some modules such as coverage can inject non-module objects inside
+        # sys.modules
+        if (
+            module_name == "__main__"
+            or module_name == "__mp_main__"
+            or module is None
+            or not isinstance(module, types.ModuleType)
+        ):
+            continue
+        try:
+            if _getattribute(module, name)[0] is obj:
+                return module_name
+        except Exception:
+            pass
+    return None
+
+
+def _lookup_module_and_qualname(
+    obj: Any, name: Optional[str] = None
+) -> Optional[tuple[types.ModuleType, str]]:
+    if name is None:
+        name = getattr(obj, "__qualname__", None)
+    if name is None:  # pragma: no cover
+        # This used to be needed for Python 2.7 support but is probably not
+        # needed anymore. However we keep the __name__ introspection in case
+        # users of cloudpickle rely on this old behavior for unknown reasons.
+        name = getattr(obj, "__name__", None)
+    if name is None:
+        return None
+
+    module_name = _whichmodule(obj, name)
+
+    if module_name is None:
+        # In this case, obj.__module__ is None AND obj was not found in any
+        # imported module. obj is thus treated as dynamic.
+        return None
+
+    if module_name == "__main__":
+        return None
+
+    # Note: if module_name is in sys.modules, the corresponding module is
+    # assumed importable at unpickling time. See #357
+    module = sys.modules.get(module_name, None)
+    if module is None:
+        # The main reason why obj's module would not be imported is that this
+        # module has been dynamically created, using for example
+        # types.ModuleType. The other possibility is that module was removed
+        # from sys.modules after obj was created/imported. But this case is not
+        # supported, as the standard pickle does not support it either.
+        return None
+
+    try:
+        obj2, parent = _getattribute(module, name)
+    except AttributeError:
+        # obj was not found inside the module it points to
+        return None
+    if obj2 is not obj:
+        return None
+    return module, name
+
+
+def get_runnable_for_func(func: Callable[..., Any]) -> RunnableSeq:
+    if func in CACHE:
+        return CACHE[func]
+    else:
+        seq = RunnableSeq(
+            coerce_to_runnable(func, name=None, trace=False),
+            ChannelWrite([ChannelWriteEntry(RETURN)]),
+            name=func.__name__,
+        )
+        if not _lookup_module_and_qualname(func):
+            return seq
+        return CACHE.setdefault(func, seq)
+
+
+CACHE: dict[Callable[..., Any], RunnableSeq] = {}

libs/langgraph/langgraph/pregel/executor.py

@@ -1,6 +1,7 @@
 import asyncio
 import concurrent.futures
 import sys
+import time
 from contextlib import ExitStack
 from contextvars import copy_context
 from types import TracebackType
@@ -34,6 +35,7 @@ class Submit(Protocol[P, T]):
         __name__: Optional[str] = None,
         __cancel_on_exit__: bool = False,
         __reraise_on_exit__: bool = True,
+        __next_tick__: bool = False,
         **kwargs: P.kwargs,
     ) -> concurrent.futures.Future[T]: ...
 
@@ -58,9 +60,13 @@ class BackgroundExecutor(ContextManager):
         __name__: Optional[str] = None,  # currently not used in sync version
         __cancel_on_exit__: bool = False,  # for sync, can cancel only if not started
         __reraise_on_exit__: bool = True,
+        __next_tick__: bool = False,
         **kwargs: P.kwargs,
     ) -> concurrent.futures.Future[T]:
-        task = self.executor.submit(fn, *args, **kwargs)
+        if __next_tick__:
+            task = self.executor.submit(next_tick, fn, *args, **kwargs)
+        else:
+            task = self.executor.submit(fn, *args, **kwargs)
         self.tasks[task] = (__cancel_on_exit__, __reraise_on_exit__)
         task.add_done_callback(self.done)
         return task
@@ -137,11 +143,14 @@ class AsyncBackgroundExecutor(AsyncContextManager):
         __name__: Optional[str] = None,
         __cancel_on_exit__: bool = False,
         __reraise_on_exit__: bool = True,
+        __next_tick__: bool = False,
         **kwargs: P.kwargs,
     ) -> asyncio.Task[T]:
         coro = cast(Coroutine[None, None, T], fn(*args, **kwargs))
         if self.semaphore:
             coro = gated(self.semaphore, coro)
+        if __next_tick__:
+            coro = anext_tick(coro)
         if self.context_not_supported:
             task = self.loop.create_task(coro, name=__name__)
         else:
@@ -197,3 +206,15 @@ async def gated(semaphore: asyncio.Semaphore, coro: Coroutine[None, None, T]) ->
     """A coroutine that waits for a semaphore before running another coroutine."""
     async with semaphore:
         return await coro
+
+
+def next_tick(fn: Callable[P, T], *args: P.args, **kwargs: P.kwargs) -> T:
+    """A function that yields control to other threads before running another function."""
+    time.sleep(0)
+    return fn(*args, **kwargs)
+
+
+async def anext_tick(coro: Coroutine[None, None, T]) -> T:
+    """A coroutine that yields control to event loop before running another coroutine."""
+    await asyncio.sleep(0)
+    return await coro

libs/langgraph/langgraph/pregel/io.py

@@ -13,6 +13,9 @@ from langgraph.constants import (
     NULL_TASK_ID,
     PUSH,
     RESUME,
+    RETURN,
+    SELF,
+    START,
     TAG_HIDDEN,
     TASKS,
 )
@@ -78,12 +81,14 @@ def map_command(
         else:
             sends = [cmd.goto]
         for send in sends:
-            if not isinstance(send, Send):
+            if isinstance(send, Send):
+                yield (NULL_TASK_ID, PUSH if FF_SEND_V2 else TASKS, send)
+            elif isinstance(send, str):
+                yield (NULL_TASK_ID, f"branch:{START}:{SELF}:{send}", START)
+            else:
                 raise TypeError(
-                    f"In Command.goto, expected Send, got {type(send).__name__}"
+                    f"In Command.goto, expected Send/str, got {type(send).__name__}"
                 )
-            yield (NULL_TASK_ID, PUSH if FF_SEND_V2 else TASKS, send)
-        # TODO handle goto str for state graph
     if cmd.resume:
         if isinstance(cmd.resume, dict) and all(is_task_id(k) for k in cmd.resume):
             for tid, resume in cmd.resume.items():
@@ -167,22 +172,21 @@ def map_output_updates(
     ]
     if not output_tasks:
         return
-    if isinstance(output_channels, str):
-        updated = (
-            (task.name, value)
-            for task, writes in output_tasks
-            for chan, value in writes
-            if chan == output_channels
-        )
-    else:
-        updated = (
-            (
-                task.name,
-                {chan: value for chan, value in writes if chan in output_channels},
+    updated: list[tuple[str, Any]] = []
+    for task, writes in output_tasks:
+        if rtn := next((value for chan, value in writes if chan == RETURN), None):
+            updated.append((task.name, rtn))
+        elif isinstance(output_channels, str):
+            updated.extend(
+                (task.name, value) for chan, value in writes if chan == output_channels
+            )
+        elif any(chan in output_channels for chan, _ in writes):
+            updated.append(
+                (
+                    task.name,
+                    {chan: value for chan, value in writes if chan in output_channels},
+                )
             )
-            for task, writes in output_tasks
-            if any(chan in output_channels for chan, _ in writes)
-        )
     grouped: dict[str, list[Any]] = {t.name: [] for t, _ in output_tasks}
     for node, value in updated:
         grouped[node].append(value)

libs/langgraph/langgraph/pregel/loop.py

@@ -73,6 +73,7 @@ from langgraph.managed.base import (
     WritableManagedValue,
 )
 from langgraph.pregel.algo import (
+    Call,
     GetNextVersion,
     PregelTaskWrites,
     apply_writes,
@@ -289,16 +290,15 @@ class PregelLoop(LoopProtocol):
         if self.checkpointer_put_writes is not None:
             self.submit(
                 self.checkpointer_put_writes,
-                {
-                    **self.checkpoint_config,
-                    CONF: {
-                        **self.checkpoint_config[CONF],
+                patch_configurable(
+                    self.checkpoint_config,
+                    {
                         CONFIG_KEY_CHECKPOINT_NS: self.config[CONF].get(
                             CONFIG_KEY_CHECKPOINT_NS, ""
                         ),
                         CONFIG_KEY_CHECKPOINT_ID: self.checkpoint["id"],
                     },
-                },
+                ),
                 writes,
                 task_id,
             )
@@ -307,20 +307,19 @@ class PregelLoop(LoopProtocol):
             self._output_writes(task_id, writes)
 
     def accept_push(
-        self, task: PregelExecutableTask, write_idx: int
+        self, task: PregelExecutableTask, write_idx: int, call: Optional[Call] = None
     ) -> Optional[PregelExecutableTask]:
         """Accept a PUSH from a task, potentially returning a new task to start."""
-        # don't start if an earlier PUSH has already triggered an interrupt
-        if self.to_interrupt:
-            return
         # don't start if we should interrupt *after* the original task
-        if should_interrupt(self.checkpoint, self.interrupt_after, [task]):
+        if self.interrupt_after and should_interrupt(
+            self.checkpoint, self.interrupt_after, [task]
+        ):
             self.to_interrupt.append(task)
             return
         if pushed := cast(
             Optional[PregelExecutableTask],
             prepare_single_task(
-                (PUSH, task.path, write_idx, task.id),
+                (PUSH, task.path, write_idx, task.id, call),
                 None,
                 checkpoint=self.checkpoint,
                 pending_writes=[(task.id, *w) for w in task.writes],
@@ -336,7 +335,9 @@ class PregelLoop(LoopProtocol):
             ),
         ):
             # don't start if we should interrupt *before* the new task
-            if should_interrupt(self.checkpoint, self.interrupt_before, [pushed]):
+            if self.interrupt_before and should_interrupt(
+                self.checkpoint, self.interrupt_before, [pushed]
+            ):
                 self.to_interrupt.append(pushed)
                 return
             # produce debug output
@@ -349,9 +350,8 @@ class PregelLoop(LoopProtocol):
             # match any pending writes to the new task
             if self.skip_done_tasks:
                 self._match_writes({pushed.id: pushed})
-            # return the new task, to be started, if not run before
-            if not pushed.writes:
-                return pushed
+            # return the new task, to be started if not run before
+            return pushed
 
     def tick(
         self,
@@ -413,7 +413,7 @@ class PregelLoop(LoopProtocol):
                 }
             )
             # after execution, check if we should interrupt
-            if should_interrupt(
+            if self.interrupt_after and should_interrupt(
                 self.checkpoint, self.interrupt_after, self.tasks.values()
             ):
                 self.status = "interrupt_after"
@@ -426,18 +426,6 @@ class PregelLoop(LoopProtocol):
             self.status = "out_of_steps"
             return False
 
-        # apply NULL writes
-        if null_writes := [
-            w[1:] for w in self.checkpoint_pending_writes if w[0] == NULL_TASK_ID
-        ]:
-            mv_writes = apply_writes(
-                self.checkpoint,
-                self.channels,
-                [PregelTaskWrites((), INPUT, null_writes, [])],
-                self.checkpointer_get_next_version,
-            )
-            for key, values in mv_writes.items():
-                self._update_mv(key, values)
         # prepare next tasks
         self.tasks = prepare_next_tasks(
             self.checkpoint,
@@ -497,7 +485,7 @@ class PregelLoop(LoopProtocol):
             return self.tick(input_keys=input_keys)
 
         # before execution, check if we should interrupt
-        if should_interrupt(
+        if self.interrupt_before and should_interrupt(
             self.checkpoint, self.interrupt_before, self.tasks.values()
         ):
             self.status = "interrupt_before"
@@ -539,9 +527,35 @@ class PregelLoop(LoopProtocol):
         # - receiving None input (outer graph) or RESUMING flag (subgraph)
         configurable = self.config.get(CONF, {})
         is_resuming = bool(self.checkpoint["channel_versions"]) and bool(
-            configurable.get(CONFIG_KEY_RESUMING, self.input is None)
+            configurable.get(
+                CONFIG_KEY_RESUMING,
+                self.input is None or isinstance(self.input, Command),
+            )
         )
 
+        # map command to writes
+        if isinstance(self.input, Command):
+            writes: defaultdict[str, list[tuple[str, Any]]] = defaultdict(list)
+            # group writes by task ID
+            for tid, c, v in map_command(self.input, self.checkpoint_pending_writes):
+                writes[tid].append((c, v))
+            if not writes:
+                raise EmptyInputError("Received empty Command input")
+            # save writes
+            for tid, ws in writes.items():
+                self.put_writes(tid, ws)
+        # apply NULL writes
+        if null_writes := [
+            w[1:] for w in self.checkpoint_pending_writes if w[0] == NULL_TASK_ID
+        ]:
+            mv_writes = apply_writes(
+                self.checkpoint,
+                self.channels,
+                [PregelTaskWrites((), INPUT, null_writes, [])],
+                self.checkpointer_get_next_version,
+            )
+            for key, values in mv_writes.items():
+                self._update_mv(key, values)
         # proceed past previous checkpoint
         if is_resuming:
             self.checkpoint["versions_seen"].setdefault(INTERRUPT, {})
@@ -553,17 +567,6 @@ class PregelLoop(LoopProtocol):
             self._emit(
                 "values", map_output_values, self.output_keys, True, self.channels
             )
-        # map command to writes
-        elif isinstance(self.input, Command):
-            writes: defaultdict[str, list[tuple[str, Any]]] = defaultdict(list)
-            # group writes by task ID
-            for tid, c, v in map_command(self.input, self.checkpoint_pending_writes):
-                writes[tid].append((c, v))
-            if not writes:
-                raise EmptyInputError("Received empty Command input")
-            # save writes
-            for tid, ws in writes.items():
-                self.put_writes(tid, ws)
         # map inputs to channel updates
         elif input_writes := deque(map_input(input_keys, self.input)):
             # TODO shouldn't these writes be passed to put_writes too?

libs/langgraph/langgraph/pregel/retry.py

@@ -4,14 +4,12 @@ import random
 import sys
 import time
 from dataclasses import replace
-from functools import partial
-from typing import Any, Callable, Optional, Sequence
+from typing import Any, Optional, Sequence
 
 from langgraph.constants import (
     CONF,
     CONFIG_KEY_CHECKPOINT_NS,
     CONFIG_KEY_RESUMING,
-    CONFIG_KEY_SEND,
     NS_SEP,
 )
 from langgraph.errors import _SEEN_CHECKPOINT_NS, GraphBubbleUp, ParentCommand
@@ -25,25 +23,21 @@ SUPPORTS_EXC_NOTES = sys.version_info >= (3, 11)
 def run_with_retry(
     task: PregelExecutableTask,
     retry_policy: Optional[RetryPolicy],
-    writer: Optional[
-        Callable[[PregelExecutableTask, Sequence[tuple[str, Any]]], None]
-    ] = None,
+    configurable: Optional[dict[str, Any]] = None,
 ) -> None:
     """Run a task with retries."""
     retry_policy = task.retry_policy or retry_policy
     interval = retry_policy.initial_interval if retry_policy else 0
     attempts = 0
     config = task.config
-    if writer is not None:
-        config = patch_configurable(config, {CONFIG_KEY_SEND: partial(writer, task)})
+    if configurable is not None:
+        config = patch_configurable(config, configurable)
     while True:
         try:
             # clear any writes from previous attempts
             task.writes.clear()
             # run the task
-            task.proc.invoke(task.input, config)
-            # if successful, end
-            break
+            return task.proc.invoke(task.input, config)
         except ParentCommand as exc:
             ns: str = config[CONF][CONFIG_KEY_CHECKPOINT_NS]
             cmd = exc.args[0]
@@ -115,17 +109,15 @@ async def arun_with_retry(
     task: PregelExecutableTask,
     retry_policy: Optional[RetryPolicy],
     stream: bool = False,
-    writer: Optional[
-        Callable[[PregelExecutableTask, Sequence[tuple[str, Any]]], None]
-    ] = None,
+    configurable: Optional[dict[str, Any]] = None,
 ) -> None:
     """Run a task asynchronously with retries."""
     retry_policy = task.retry_policy or retry_policy
     interval = retry_policy.initial_interval if retry_policy else 0
     attempts = 0
     config = task.config
-    if writer is not None:
-        config = patch_configurable(config, {CONFIG_KEY_SEND: partial(writer, task)})
+    if configurable is not None:
+        config = patch_configurable(config, configurable)
     while True:
         try:
             # clear any writes from previous attempts
@@ -134,10 +126,10 @@ async def arun_with_retry(
             if stream:
                 async for _ in task.proc.astream(task.input, config):
                     pass
+                # if successful, end
+                break
             else:
-                await task.proc.ainvoke(task.input, config)
-            # if successful, end
-            break
+                return await task.proc.ainvoke(task.input, config)
         except ParentCommand as exc:
             ns: str = config[CONF][CONFIG_KEY_CHECKPOINT_NS]
             cmd = exc.args[0]

libs/langgraph/langgraph/pregel/runner.py

@@ -1,9 +1,12 @@
 import asyncio
 import concurrent.futures
+import threading
 import time
+from functools import partial
 from typing import (
     Any,
     AsyncIterator,
+    Awaitable,
     Callable,
     Iterable,
     Iterator,
@@ -16,18 +19,22 @@ from typing import (
 
 from langgraph.constants import (
     CONF,
+    CONFIG_KEY_CALL,
     CONFIG_KEY_SEND,
     ERROR,
     INTERRUPT,
     NO_WRITES,
     PUSH,
     RESUME,
+    RETURN,
     TAG_HIDDEN,
 )
 from langgraph.errors import GraphBubbleUp, GraphInterrupt
+from langgraph.pregel.algo import Call
 from langgraph.pregel.executor import Submit
 from langgraph.pregel.retry import arun_with_retry, run_with_retry
 from langgraph.types import PregelExecutableTask, RetryPolicy
+from langgraph.utils.future import chain_future
 
 
 class PregelRunner:
@@ -41,7 +48,7 @@ class PregelRunner:
         submit: Submit,
         put_writes: Callable[[str, Sequence[tuple[str, Any]]], None],
         schedule_task: Callable[
-            [PregelExecutableTask, int], Optional[PregelExecutableTask]
+            [PregelExecutableTask, int, Optional[Call]], Optional[PregelExecutableTask]
         ],
         use_astream: bool = False,
         node_finished: Optional[Callable[[str], None]] = None,
@@ -61,73 +68,143 @@ class PregelRunner:
         retry_policy: Optional[RetryPolicy] = None,
         get_waiter: Optional[Callable[[], concurrent.futures.Future[None]]] = None,
     ) -> Iterator[None]:
+        locks: dict[str, threading.Lock] = {}
+
         def writer(
-            task: PregelExecutableTask, writes: Sequence[tuple[str, Any]]
-        ) -> None:
-            prev_length = len(task.writes)
-            # delegate to the underlying writer
-            task.config[CONF][CONFIG_KEY_SEND](writes)
-            for idx, w in enumerate(task.writes):
-                # find the index for the newly inserted writes
-                if idx < prev_length:
-                    continue
-                assert writes[idx - prev_length] is w
+            task: PregelExecutableTask,
+            writes: Sequence[tuple[str, Any]],
+            *,
+            calls: Optional[Sequence[Call]] = None,
+        ) -> Sequence[Optional[concurrent.futures.Future]]:
+            if all(w[0] != PUSH for w in writes):
+                return task.config[CONF][CONFIG_KEY_SEND](writes)
+
+            if task.id not in locks:
+                locks[task.id] = threading.Lock()
+            with locks[task.id]:
+                prev_length = len(task.writes)
+                # delegate to the underlying writer
+                task.config[CONF][CONFIG_KEY_SEND](writes)
+                # confirm no other concurrent writes were added
+                assert len(task.writes) == prev_length + len(writes)
+            # schedule PUSH tasks, collect futures
+            rtn: dict[int, Optional[concurrent.futures.Future]] = {}
+            for idx, w in enumerate(writes, start=prev_length):
                 # bail if not a PUSH write
                 if w[0] != PUSH:
                     continue
                 # schedule the next task, if the callback returns one
-                if next_task := self.schedule_task(task, idx):
-                    # if the parent task was retried,
-                    # the next task might already be running
-                    if any(
-                        t == next_task.id for t in futures.values() if t is not None
+                if next_task := self.schedule_task(
+                    task, idx, calls[idx - prev_length] if calls else None
+                ):
+                    if fut := next(
+                        (
+                            f
+                            for f, t in futures.items()
+                            if t is not None and t == next_task.id
+                        ),
+                        None,
                     ):
-                        continue
-                    # schedule the next task
-                    futures[
-                        self.submit(
+                        # if the parent task was retried,
+                        # the next task might already be running
+                        rtn[idx - prev_length] = fut
+                    elif next_task.writes:
+                        # if it already ran, return the result
+                        fut = concurrent.futures.Future()
+                        if val := next(v for c, v in next_task.writes if c == RETURN):
+                            fut.set_result(val)
+                        elif exc := next(v for c, v in next_task.writes if c == ERROR):
+                            fut.set_exception(
+                                exc
+                                if isinstance(exc, BaseException)
+                                else Exception(exc)
+                            )
+                        else:
+                            fut.set_result(None)
+                        rtn[idx - prev_length] = fut
+                    else:
+                        # schedule the next task
+                        fut = self.submit(
                             run_with_retry,
                             next_task,
                             retry_policy,
-                            writer=writer,
+                            configurable={
+                                CONFIG_KEY_SEND: partial(writer, next_task),
+                                CONFIG_KEY_CALL: partial(call, next_task),
+                            },
                             __reraise_on_exit__=reraise,
+                            # starting a new task in the next tick ensures
+                            # updates from this tick are committed/streamed first
+                            __next_tick__=True,
                         )
-                    ] = next_task
+                        fut.add_done_callback(partial(self.commit, next_task))
+                        futures[fut] = next_task
+                        rtn[idx - prev_length] = fut
+            return [rtn.get(i) for i in range(len(writes))]
+
+        def call(
+            task: PregelExecutableTask,
+            func: Callable[[Any], Union[Awaitable[Any], Any]],
+            input: Any,
+            *,
+            retry: Optional[RetryPolicy] = None,
+        ) -> concurrent.futures.Future[Any]:
+            (fut,) = writer(
+                task, [(PUSH, None)], calls=[Call(func, input, retry=retry)]
+            )
+            assert fut is not None, "writer did not return a future for call"
+            return fut
 
         tasks = tuple(tasks)
         futures: dict[concurrent.futures.Future, Optional[PregelExecutableTask]] = {}
+        done_futures: set[concurrent.futures.Future] = set()
         # give control back to the caller
         yield
         # fast path if single task with no timeout and no waiter
         if len(tasks) == 1 and timeout is None and get_waiter is None:
             t = tasks[0]
             try:
-                run_with_retry(t, retry_policy, writer=writer)
+                run_with_retry(
+                    t,
+                    retry_policy,
+                    configurable={
+                        CONFIG_KEY_SEND: partial(writer, t),
+                        CONFIG_KEY_CALL: partial(call, t),
+                    },
+                )
                 self.commit(t, None)
             except Exception as exc:
-                self.commit(t, exc)
-                if reraise:
+                self.commit(t, None, exc)
+                if reraise and futures:
+                    # will be re-raised after futures are done
+                    fut: concurrent.futures.Future = concurrent.futures.Future()
+                    fut.set_exception(exc)
+                    done_futures.add(fut)
+                elif reraise:
                     raise
             if not futures:  # maybe `t` schuduled another task
                 return
         # add waiter task if requested
         if get_waiter is not None:
             futures[get_waiter()] = None
+        # schedule tasks
+        for t in tasks:
+            if not t.writes:
+                fut = self.submit(
+                    run_with_retry,
+                    t,
+                    retry_policy,
+                    configurable={
+                        CONFIG_KEY_SEND: partial(writer, t),
+                        CONFIG_KEY_CALL: partial(call, t),
+                    },
+                    __reraise_on_exit__=reraise,
+                )
+                fut.add_done_callback(partial(self.commit, t))
+                futures[fut] = t
         # execute tasks, and wait for one to fail or all to finish.
         # each task is independent from all other concurrent tasks
         # yield updates/debug output as each task finishes
-        for t in tasks:
-            if not t.writes:
-                futures[
-                    self.submit(
-                        run_with_retry,
-                        t,
-                        retry_policy,
-                        writer=writer,
-                        __reraise_on_exit__=reraise,
-                    )
-                ] = t
-        done_futures: set[concurrent.futures.Future] = set()
         end_time = timeout + time.monotonic() if timeout else None
         while len(futures) > (1 if get_waiter is not None else 0):
             done, inflight = concurrent.futures.wait(
@@ -146,8 +223,6 @@ class PregelRunner:
                 else:
                     # store for panic check
                     done_futures.add(fut)
-                    # task finished, commit writes
-                    self.commit(task, _exception(fut))
             else:
                 # remove references to loop vars
                 del fut, task
@@ -156,6 +231,10 @@ class PregelRunner:
                 break
             # give control back to the caller
             yield
+        # wait for pending done callbacks
+        # if a 2nd future finishes while `wait` is returning, it's possible
+        # that done callbacks for the 2nd future aren't called until next tick
+        time.sleep(0)
         # panic on failure or timeout
         _panic_or_proceed(
             done_futures.union(f for f, t in futures.items() if t is not None),
@@ -171,48 +250,109 @@ class PregelRunner:
         retry_policy: Optional[RetryPolicy] = None,
         get_waiter: Optional[Callable[[], asyncio.Future[None]]] = None,
     ) -> AsyncIterator[None]:
+        locks: dict[str, threading.Lock] = {}
+
         def writer(
-            task: PregelExecutableTask, writes: Sequence[tuple[str, Any]]
-        ) -> None:
-            prev_length = len(task.writes)
-            # delegate to the underlying writer
-            task.config[CONF][CONFIG_KEY_SEND](writes)
-            for idx, w in enumerate(task.writes):
-                # find the index for the newly inserted writes
-                if idx < prev_length:
-                    continue
-                assert writes[idx - prev_length] is w
+            task: PregelExecutableTask,
+            writes: Sequence[tuple[str, Any]],
+            *,
+            calls: Optional[Sequence[Call]] = None,
+        ) -> Sequence[Optional[asyncio.Future]]:
+            if all(w[0] != PUSH for w in writes):
+                return task.config[CONF][CONFIG_KEY_SEND](writes)
+
+            if task.id not in locks:
+                locks[task.id] = threading.Lock()
+            with locks[task.id]:
+                prev_length = len(task.writes)
+                # delegate to the underlying writer
+                task.config[CONF][CONFIG_KEY_SEND](writes)
+                # confirm no other concurrent writes were added
+                assert len(task.writes) == prev_length + len(writes)
+            # schedule PUSH tasks, collect futures
+            rtn: dict[int, Optional[asyncio.Future]] = {}
+            for idx, w in enumerate(writes, start=prev_length):
                 # bail if not a PUSH write
                 if w[0] != PUSH:
                     continue
                 # schedule the next task, if the callback returns one
-                if next_task := self.schedule_task(task, idx):
+                wcall = calls[idx - prev_length] if calls is not None else None
+                if next_task := self.schedule_task(task, idx, wcall):
                     # if the parent task was retried,
                     # the next task might already be running
-                    if any(
-                        t == next_task.id for t in futures.values() if t is not None
+                    if fut := next(
+                        (
+                            f
+                            for f, t in futures.items()
+                            if t is not None and t == next_task.id
+                        ),
+                        None,
                     ):
-                        continue
-                    # schedule the next task
-                    futures[
-                        cast(
+                        # if the parent task was retried,
+                        # the next task might already be running
+                        rtn[idx - prev_length] = fut
+                    elif next_task.writes:
+                        # if it already ran, return the result
+                        fut = asyncio.Future()
+                        if val := next(v for c, v in next_task.writes if c == RETURN):
+                            fut.set_result(val)
+                        elif exc := next(v for c, v in next_task.writes if c == ERROR):
+                            fut.set_exception(
+                                exc
+                                if isinstance(exc, BaseException)
+                                else Exception(exc)
+                            )
+                        else:
+                            fut.set_result(None)
+                        rtn[idx - prev_length] = fut
+                    else:
+                        # schedule the next task
+                        fut = cast(
                             asyncio.Future,
                             self.submit(
                                 arun_with_retry,
                                 next_task,
                                 retry_policy,
                                 stream=self.use_astream,
-                                writer=writer,
+                                configurable={
+                                    CONFIG_KEY_SEND: partial(writer, next_task),
+                                    CONFIG_KEY_CALL: partial(call, next_task),
+                                },
                                 __name__=t.name,
                                 __cancel_on_exit__=True,
                                 __reraise_on_exit__=reraise,
+                                # starting a new task in the next tick ensures
+                                # updates from this tick are committed/streamed first
+                                __next_tick__=True,
                             ),
                         )
-                    ] = next_task
+                        fut.add_done_callback(partial(self.commit, next_task))
+                        futures[fut] = next_task
+                        rtn[idx - prev_length] = fut
+            return [rtn.get(i) for i in range(len(writes))]
+
+        def call(
+            task: PregelExecutableTask,
+            func: Callable[[Any], Union[Awaitable[Any], Any]],
+            input: Any,
+            *,
+            retry: Optional[RetryPolicy] = None,
+        ) -> Union[asyncio.Future[Any], concurrent.futures.Future[Any]]:
+            (fut,) = writer(
+                task, [(PUSH, None)], calls=[Call(func, input, retry=retry)]
+            )
+            assert fut is not None, "writer did not return a future for call"
+            if asyncio.iscoroutinefunction(func):
+                return fut
+            # adapted from asyncio.run_coroutine_threadsafe
+            sfut: concurrent.futures.Future = concurrent.futures.Future()
+            loop.call_soon_threadsafe(chain_future, fut, sfut)
+            return sfut
 
         loop = asyncio.get_event_loop()
         tasks = tuple(tasks)
         futures: dict[asyncio.Future, Optional[PregelExecutableTask]] = {}
+        done_futures: set[asyncio.Future] = set()
         # give control back to the caller
         yield
         # fast path if single task with no waiter and no timeout
@@ -220,39 +360,53 @@ class PregelRunner:
             t = tasks[0]
             try:
                 await arun_with_retry(
-                    t, retry_policy, stream=self.use_astream, writer=writer
+                    t,
+                    retry_policy,
+                    stream=self.use_astream,
+                    configurable={
+                        CONFIG_KEY_SEND: partial(writer, t),
+                        CONFIG_KEY_CALL: partial(call, t),
+                    },
                 )
                 self.commit(t, None)
             except Exception as exc:
-                self.commit(t, exc)
-                if reraise:
+                self.commit(t, None, exc)
+                if reraise and futures:
+                    # will be re-raised after futures are done
+                    fut: asyncio.Future = loop.create_future()
+                    fut.set_exception(exc)
+                    done_futures.add(fut)
+                elif reraise:
                     raise
             if not futures:  # maybe `t` schuduled another task
                 return
         # add waiter task if requested
         if get_waiter is not None:
             futures[get_waiter()] = None
+        # schedule tasks
+        for t in tasks:
+            if not t.writes:
+                fut = cast(
+                    asyncio.Future,
+                    self.submit(
+                        arun_with_retry,
+                        t,
+                        retry_policy,
+                        stream=self.use_astream,
+                        configurable={
+                            CONFIG_KEY_SEND: partial(writer, t),
+                            CONFIG_KEY_CALL: partial(call, t),
+                        },
+                        __name__=t.name,
+                        __cancel_on_exit__=True,
+                        __reraise_on_exit__=reraise,
+                    ),
+                )
+                fut.add_done_callback(partial(self.commit, t))
+                futures[fut] = t
         # execute tasks, and wait for one to fail or all to finish.
         # each task is independent from all other concurrent tasks
         # yield updates/debug output as each task finishes
-        for t in tasks:
-            if not t.writes:
-                futures[
-                    cast(
-                        asyncio.Future,
-                        self.submit(
-                            arun_with_retry,
-                            t,
-                            retry_policy,
-                            stream=self.use_astream,
-                            writer=writer,
-                            __name__=t.name,
-                            __cancel_on_exit__=True,
-                            __reraise_on_exit__=reraise,
-                        ),
-                    )
-                ] = t
-        done_futures: set[asyncio.Future] = set()
         end_time = timeout + loop.time() if timeout else None
         while len(futures) > (1 if get_waiter is not None else 0):
             done, inflight = await asyncio.wait(
@@ -271,8 +425,6 @@ class PregelRunner:
                 else:
                     # store for panic check
                     done_futures.add(fut)
-                    # task finished, commit writes
-                    self.commit(task, _exception(fut))
             else:
                 # remove references to loop vars
                 del fut, task
@@ -281,6 +433,10 @@ class PregelRunner:
                 break
             # give control back to the caller
             yield
+        # wait for pending done callbacks
+        # if a 2nd future finishes while `wait` is returning, it's possible
+        # that done callbacks for the 2nd future aren't called until next tick
+        await asyncio.sleep(0)
         # cancel waiter task
         for fut in futures:
             fut.cancel()
@@ -292,9 +448,19 @@ class PregelRunner:
         )
 
     def commit(
-        self, task: PregelExecutableTask, exception: Optional[BaseException]
+        self,
+        task: PregelExecutableTask,
+        fut: Union[None, concurrent.futures.Future[Any], asyncio.Future[Any]],
+        exception: Optional[BaseException] = None,
     ) -> None:
-        if exception:
+        if fut is not None:
+            exception = _exception(fut)
+        if isinstance(exception, asyncio.CancelledError):
+            # for cancelled tasks, also save error in task,
+            # so loop can finish super-step
+            task.writes.append((ERROR, exception))
+            self.put_writes(task.id, task.writes)
+        elif exception:
             if isinstance(exception, GraphInterrupt):
                 # save interrupt to checkpointer
                 if interrupts := [(INTERRUPT, i) for i in exception.args[0]]:
@@ -325,11 +491,12 @@ def _should_stop_others(
     GraphInterrupts are not considered failures."""
     for fut in done:
         if fut.cancelled():
-            return True
-        if exc := fut.exception():
-            return not isinstance(exc, GraphBubbleUp)
-    else:
-        return False
+            continue
+        elif exc := fut.exception():
+            if not isinstance(exc, GraphBubbleUp):
+                return True
+
+    return False
 
 
 def _exception(
@@ -355,7 +522,9 @@ def _panic_or_proceed(
     done: set[Union[concurrent.futures.Future[Any], asyncio.Future[Any]]] = set()
     inflight: set[Union[concurrent.futures.Future[Any], asyncio.Future[Any]]] = set()
     for fut in futs:
-        if fut.done():
+        if fut.cancelled():
+            continue
+        elif fut.done():
             done.add(fut)
         else:
             inflight.add(fut)
@@ -368,8 +537,6 @@ def _panic_or_proceed(
             # raise the exception
             if panic:
                 raise exc
-            else:
-                return
     if inflight:
         # if we got here means we timed out
         while inflight:

libs/langgraph/langgraph/types.py

@@ -32,6 +32,14 @@ if TYPE_CHECKING:
     from langgraph.store.base import BaseStore
 
 
+try:
+    from langchain_core.messages.tool import ToolOutputMixin
+except ImportError:
+
+    class ToolOutputMixin:  # type: ignore[no-redef]
+        pass
+
+
 All = Literal["*"]
 """Special value to indicate that graph should interrupt on all nodes."""
 
@@ -244,7 +252,7 @@ N = TypeVar("N", bound=Hashable)
 
 
 @dataclasses.dataclass(**_DC_KWARGS)
-class Command(Generic[N]):
+class Command(Generic[N], ToolOutputMixin):
     """One or more commands to update the graph's state and send messages to nodes.
 
     Args:

libs/langgraph/langgraph/utils/fields.py

@@ -1,5 +1,5 @@
 import dataclasses
-from typing import Any, Optional, Type, Union
+from typing import Any, Generator, Optional, Type, Union, get_type_hints
 
 from typing_extensions import Annotated, NotRequired, ReadOnly, Required, get_origin
 
@@ -106,3 +106,44 @@ def get_field_default(name: str, type_: Any, schema: Type[Any]) -> Any:
     if _is_optional_type(type_):
         return None
     return ...
+
+
+def get_enhanced_type_hints(
+    type: Type[Any],
+) -> Generator[tuple[str, Any, Any, Optional[str]], None, None]:
+    """Attempt to extract default values and descriptions from provided type, used for config schema."""
+    for name, typ in get_type_hints(type).items():
+        default = None
+        description = None
+
+        # Pydantic models
+        try:
+            if hasattr(type, "__fields__") and name in type.__fields__:
+                field = type.__fields__[name]
+
+                if hasattr(field, "description") and field.description is not None:
+                    description = field.description
+
+                if hasattr(field, "default") and field.default is not None:
+                    default = field.default
+                    if (
+                        hasattr(default, "__class__")
+                        and getattr(default.__class__, "__name__", "")
+                        == "PydanticUndefinedType"
+                    ):
+                        default = None
+
+        except (AttributeError, KeyError, TypeError):
+            pass
+
+        # TypedDict, dataclass
+        try:
+            if hasattr(type, "__dict__"):
+                type_dict = getattr(type, "__dict__")
+
+                if name in type_dict:
+                    default = type_dict[name]
+        except (AttributeError, KeyError, TypeError):
+            pass
+
+        yield name, typ, default, description

libs/langgraph/langgraph/utils/future.py

@@ -0,0 +1,124 @@
+import asyncio
+import concurrent.futures
+from typing import Union
+
+AnyFuture = Union[asyncio.Future, concurrent.futures.Future]
+
+
+def _get_loop(fut: asyncio.Future) -> asyncio.AbstractEventLoop:
+    # Tries to call Future.get_loop() if it's available.
+    # Otherwise fallbacks to using the old '_loop' property.
+    try:
+        get_loop = fut.get_loop
+    except AttributeError:
+        pass
+    else:
+        return get_loop()
+    return fut._loop
+
+
+def _convert_future_exc(exc: BaseException) -> BaseException:
+    exc_class = type(exc)
+    if exc_class is concurrent.futures.CancelledError:
+        return asyncio.CancelledError(*exc.args)
+    elif exc_class is concurrent.futures.TimeoutError:
+        return asyncio.TimeoutError(*exc.args)
+    elif exc_class is concurrent.futures.InvalidStateError:
+        return asyncio.InvalidStateError(*exc.args)
+    else:
+        return exc
+
+
+def _set_concurrent_future_state(
+    concurrent: concurrent.futures.Future,
+    source: AnyFuture,
+) -> None:
+    """Copy state from a future to a concurrent.futures.Future."""
+    assert source.done()
+    if source.cancelled():
+        concurrent.cancel()
+    if not concurrent.set_running_or_notify_cancel():
+        return
+    exception = source.exception()
+    if exception is not None:
+        concurrent.set_exception(_convert_future_exc(exception))
+    else:
+        result = source.result()
+        concurrent.set_result(result)
+
+
+def _copy_future_state(source: AnyFuture, dest: asyncio.Future) -> None:
+    """Internal helper to copy state from another Future.
+
+    The other Future may be a concurrent.futures.Future.
+    """
+    assert source.done()
+    if dest.cancelled():
+        return
+    assert not dest.done()
+    if source.cancelled():
+        dest.cancel()
+    else:
+        exception = source.exception()
+        if exception is not None:
+            dest.set_exception(_convert_future_exc(exception))
+        else:
+            result = source.result()
+            dest.set_result(result)
+
+
+def _chain_future(source: AnyFuture, destination: AnyFuture) -> None:
+    """Chain two futures so that when one completes, so does the other.
+
+    The result (or exception) of source will be copied to destination.
+    If destination is cancelled, source gets cancelled too.
+    Compatible with both asyncio.Future and concurrent.futures.Future.
+    """
+    if not asyncio.isfuture(source) and not isinstance(
+        source, concurrent.futures.Future
+    ):
+        raise TypeError("A future is required for source argument")
+    if not asyncio.isfuture(destination) and not isinstance(
+        destination, concurrent.futures.Future
+    ):
+        raise TypeError("A future is required for destination argument")
+    source_loop = _get_loop(source) if asyncio.isfuture(source) else None
+    dest_loop = _get_loop(destination) if asyncio.isfuture(destination) else None
+
+    def _set_state(future: AnyFuture, other: AnyFuture) -> None:
+        if asyncio.isfuture(future):
+            _copy_future_state(other, future)
+        else:
+            _set_concurrent_future_state(future, other)
+
+    def _call_check_cancel(destination: AnyFuture) -> None:
+        if destination.cancelled():
+            if source_loop is None or source_loop is dest_loop:
+                source.cancel()
+            else:
+                source_loop.call_soon_threadsafe(source.cancel)
+
+    def _call_set_state(source: AnyFuture) -> None:
+        if destination.cancelled() and dest_loop is not None and dest_loop.is_closed():
+            return
+        if dest_loop is None or dest_loop is source_loop:
+            _set_state(destination, source)
+        else:
+            if dest_loop.is_closed():
+                return
+            dest_loop.call_soon_threadsafe(_set_state, destination, source)
+
+    destination.add_done_callback(_call_check_cancel)
+    source.add_done_callback(_call_set_state)
+
+
+def chain_future(source: AnyFuture, destination: concurrent.futures.Future) -> None:
+    # adapted from asyncio.run_coroutine_threadsafe
+    try:
+        _chain_future(source, destination)
+    except (SystemExit, KeyboardInterrupt):
+        raise
+    except BaseException as exc:
+        if destination.set_running_or_notify_cancel():
+            destination.set_exception(exc)
+        raise

libs/langgraph/langgraph/utils/runnable.py

@@ -404,12 +404,10 @@ class RunnableSeq(Runnable):
                 config = patch_config(
                     config, callbacks=run_manager.get_child(f"seq:step:{i+1}")
                 )
-                context = copy_context()
-                context.run(_set_config_context, config)
                 if i == 0:
-                    input = context.run(step.invoke, input, config, **kwargs)
+                    input = step.invoke(input, config, **kwargs)
                 else:
-                    input = context.run(step.invoke, input, config)
+                    input = step.invoke(input, config)
         # finish the root run
         except BaseException as e:
             run_manager.on_chain_error(e)
@@ -443,16 +441,10 @@ class RunnableSeq(Runnable):
                 config = patch_config(
                     config, callbacks=run_manager.get_child(f"seq:step:{i+1}")
                 )
-                context = copy_context()
-                context.run(_set_config_context, config)
                 if i == 0:
-                    coro = step.ainvoke(input, config, **kwargs)
+                    input = await step.ainvoke(input, config, **kwargs)
                 else:
-                    coro = step.ainvoke(input, config)
-                if ASYNCIO_ACCEPTS_CONTEXT:
-                    input = await asyncio.create_task(coro, context=context)
-                else:
-                    input = await asyncio.create_task(coro)
+                    input = await step.ainvoke(input, config)
         # finish the root run
         except BaseException as e:
             await run_manager.on_chain_error(e)

libs/langgraph/poetry.lock

@@ -1325,13 +1325,13 @@ files = [
 
 [[package]]
 name = "langchain-core"
-version = "0.3.15"
+version = "0.3.23"
 description = "Building applications with LLMs through composability"
 optional = false
 python-versions = "<4.0,>=3.9"
 files = [
-    {file = "langchain_core-0.3.15-py3-none-any.whl", hash = "sha256:3d4ca6dbb8ed396a6ee061063832a2451b0ce8c345570f7b086ffa7288e4fa29"},
-    {file = "langchain_core-0.3.15.tar.gz", hash = "sha256:b1a29787a4ffb7ec2103b4e97d435287201da7809b369740dd1e32f176325aba"},
+    {file = "langchain_core-0.3.23-py3-none-any.whl", hash = "sha256:550c0b996990830fa6515a71a1192a8a0343367999afc36d4ede14222941e420"},
+    {file = "langchain_core-0.3.23.tar.gz", hash = "sha256:f9e175e3b82063cc3b160c2ca2b155832e1c6f915312e1204828f97d4aabf6e1"},
 ]
 
 [package.dependencies]
@@ -1382,7 +1382,7 @@ url = "../checkpoint-duckdb"
 
 [[package]]
 name = "langgraph-checkpoint-postgres"
-version = "2.0.7"
+version = "2.0.8"
 description = "Library with a Postgres implementation of LangGraph checkpoint saver."
 optional = false
 python-versions = "^3.9.0,<4.0"
@@ -1418,7 +1418,7 @@ url = "../checkpoint-sqlite"
 
 [[package]]
 name = "langgraph-sdk"
-version = "0.1.42"
+version = "0.1.43"
 description = "SDK for interacting with LangGraph API"
 optional = false
 python-versions = "^3.9.0,<4.0"
@@ -3413,4 +3413,4 @@ type = ["pytest-mypy"]
 [metadata]
 lock-version = "2.0"
 python-versions = ">=3.9.0,<4.0"
-content-hash = "2df4d5d5e61917bdfff0ba430067a17662666eedee2858d841fa02e594cf69d0"
+content-hash = "936530a5f00f329aeff2e6e921fe64480be317fad2c0a59cd54ea9018d089304"

libs/langgraph/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "langgraph"
-version = "0.2.56"
+version = "0.2.58"
 description = "Building stateful, multi-actor applications with LLMs"
 authors = []
 license = "MIT"
@@ -9,7 +9,7 @@ repository = "https://www.github.com/langchain-ai/langgraph"
 
 [tool.poetry.dependencies]
 python = ">=3.9.0,<4.0"
-langchain-core = ">=0.2.43,<0.4.0,!=0.3.0,!=0.3.1,!=0.3.2,!=0.3.3,!=0.3.4,!=0.3.5,!=0.3.6,!=0.3.7,!=0.3.8,!=0.3.9,!=0.3.10,!=0.3.11,!=0.3.12,!=0.3.13,!=0.3.14"
+langchain-core = ">=0.2.43,<0.4.0,!=0.3.0,!=0.3.1,!=0.3.2,!=0.3.3,!=0.3.4,!=0.3.5,!=0.3.6,!=0.3.7,!=0.3.8,!=0.3.9,!=0.3.10,!=0.3.11,!=0.3.12,!=0.3.13,!=0.3.14,!=0.3.15,!=0.3.16,!=0.3.17,!=0.3.18,!=0.3.19,!=0.3.20,!=0.3.21,!=0.3.22"
 langgraph-checkpoint = "^2.0.4"
 langgraph-sdk = "^0.1.42"
 

libs/langgraph/tests/test_prebuilt.py

@@ -46,7 +46,7 @@ from langgraph.prebuilt import (
     create_react_agent,
     tools_condition,
 )
-from langgraph.prebuilt.chat_agent_executor import _validate_chat_history
+from langgraph.prebuilt.chat_agent_executor import AgentState, _validate_chat_history
 from langgraph.prebuilt.tool_node import (
     TOOL_CALL_ERROR_TEMPLATE,
     InjectedState,
@@ -56,7 +56,7 @@ from langgraph.prebuilt.tool_node import (
 )
 from langgraph.store.base import BaseStore
 from langgraph.store.memory import InMemoryStore
-from langgraph.types import Interrupt
+from langgraph.types import Command, Interrupt, interrupt
 from tests.conftest import (
     ALL_CHECKPOINTERS_ASYNC,
     ALL_CHECKPOINTERS_SYNC,
@@ -988,6 +988,645 @@ def test_tool_node_node_interrupt():
     assert task.interrupts == (Interrupt(value="foo", when="during"),)
 
 
+@pytest.mark.skipif(
+    not IS_LANGCHAIN_CORE_030_OR_GREATER,
+    reason="Langchain core 0.3.0 or greater is required",
+)
+async def test_tool_node_command():
+    from langchain_core.tools.base import InjectedToolCallId
+
+    @dec_tool
+    def transfer_to_bob(tool_call_id: Annotated[str, InjectedToolCallId]):
+        """Transfer to Bob"""
+        return Command(
+            update={
+                "messages": [
+                    ToolMessage(content="Transferred to Bob", tool_call_id=tool_call_id)
+                ]
+            },
+            goto="bob",
+            graph=Command.PARENT,
+        )
+
+    @dec_tool
+    async def async_transfer_to_bob(tool_call_id: Annotated[str, InjectedToolCallId]):
+        """Transfer to Bob"""
+        return Command(
+            update={
+                "messages": [
+                    ToolMessage(content="Transferred to Bob", tool_call_id=tool_call_id)
+                ]
+            },
+            goto="bob",
+            graph=Command.PARENT,
+        )
+
+    class CustomToolSchema(BaseModel):
+        tool_call_id: Annotated[str, InjectedToolCallId]
+
+    class MyCustomTool(BaseTool):
+        def _run(*args: Any, **kwargs: Any):
+            return Command(
+                update={
+                    "messages": [
+                        ToolMessage(
+                            content="Transferred to Bob",
+                            tool_call_id=kwargs["tool_call_id"],
+                        )
+                    ]
+                },
+                goto="bob",
+                graph=Command.PARENT,
+            )
+
+        async def _arun(*args: Any, **kwargs: Any):
+            return Command(
+                update={
+                    "messages": [
+                        ToolMessage(
+                            content="Transferred to Bob",
+                            tool_call_id=kwargs["tool_call_id"],
+                        )
+                    ]
+                },
+                goto="bob",
+                graph=Command.PARENT,
+            )
+
+    custom_tool = MyCustomTool(
+        name="custom_transfer_to_bob",
+        description="Transfer to bob",
+        args_schema=CustomToolSchema,
+    )
+    async_custom_tool = MyCustomTool(
+        name="async_custom_transfer_to_bob",
+        description="Transfer to bob",
+        args_schema=CustomToolSchema,
+    )
+
+    # test mixing regular tools and tools returning commands
+    def add(a: int, b: int) -> int:
+        """Add two numbers"""
+        return a + b
+
+    result = ToolNode([add, transfer_to_bob]).invoke(
+        {
+            "messages": [
+                AIMessage(
+                    "",
+                    tool_calls=[
+                        {"args": {"a": 1, "b": 2}, "id": "1", "name": "add"},
+                        {"args": {}, "id": "2", "name": "transfer_to_bob"},
+                    ],
+                )
+            ]
+        }
+    )
+
+    assert result == [
+        {
+            "messages": [
+                ToolMessage(
+                    content="3",
+                    tool_call_id="1",
+                    name="add",
+                )
+            ]
+        },
+        Command(
+            update={
+                "messages": [
+                    ToolMessage(
+                        content="Transferred to Bob",
+                        tool_call_id="2",
+                        name="transfer_to_bob",
+                    )
+                ]
+            },
+            goto="bob",
+            graph=Command.PARENT,
+        ),
+    ]
+
+    # test tools returning commands
+
+    # test sync tools
+    for tool in [transfer_to_bob, custom_tool]:
+        result = ToolNode([tool]).invoke(
+            {
+                "messages": [
+                    AIMessage(
+                        "", tool_calls=[{"args": {}, "id": "1", "name": tool.name}]
+                    )
+                ]
+            }
+        )
+        assert result == [
+            Command(
+                update={
+                    "messages": [
+                        ToolMessage(
+                            content="Transferred to Bob",
+                            tool_call_id="1",
+                            name=tool.name,
+                        )
+                    ]
+                },
+                goto="bob",
+                graph=Command.PARENT,
+            )
+        ]
+
+    # test async tools
+    for tool in [async_transfer_to_bob, async_custom_tool]:
+        result = await ToolNode([tool]).ainvoke(
+            {
+                "messages": [
+                    AIMessage(
+                        "", tool_calls=[{"args": {}, "id": "1", "name": tool.name}]
+                    )
+                ]
+            }
+        )
+        assert result == [
+            Command(
+                update={
+                    "messages": [
+                        ToolMessage(
+                            content="Transferred to Bob",
+                            tool_call_id="1",
+                            name=tool.name,
+                        )
+                    ]
+                },
+                goto="bob",
+                graph=Command.PARENT,
+            )
+        ]
+
+    # test multiple commands
+    result = ToolNode([transfer_to_bob, custom_tool]).invoke(
+        {
+            "messages": [
+                AIMessage(
+                    "",
+                    tool_calls=[
+                        {"args": {}, "id": "1", "name": "transfer_to_bob"},
+                        {"args": {}, "id": "2", "name": "custom_transfer_to_bob"},
+                    ],
+                )
+            ]
+        }
+    )
+    assert result == [
+        Command(
+            update={
+                "messages": [
+                    ToolMessage(
+                        content="Transferred to Bob",
+                        tool_call_id="1",
+                        name="transfer_to_bob",
+                    )
+                ]
+            },
+            goto="bob",
+            graph=Command.PARENT,
+        ),
+        Command(
+            update={
+                "messages": [
+                    ToolMessage(
+                        content="Transferred to Bob",
+                        tool_call_id="2",
+                        name="custom_transfer_to_bob",
+                    )
+                ]
+            },
+            goto="bob",
+            graph=Command.PARENT,
+        ),
+    ]
+
+    # test validation (mismatch between input type and command.update type)
+    with pytest.raises(ValueError):
+
+        @dec_tool
+        def list_update_tool(tool_call_id: Annotated[str, InjectedToolCallId]):
+            """My tool"""
+            return Command(
+                update=[ToolMessage(content="foo", tool_call_id=tool_call_id)]
+            )
+
+        ToolNode([list_update_tool]).invoke(
+            {
+                "messages": [
+                    AIMessage(
+                        "",
+                        tool_calls=[
+                            {"args": {}, "id": "1", "name": "list_update_tool"}
+                        ],
+                    )
+                ]
+            }
+        )
+
+    # test validation (missing tool message in the update for current graph)
+    with pytest.raises(ValueError):
+
+        @dec_tool
+        def no_update_tool():
+            """My tool"""
+            return Command(update={"messages": []})
+
+        ToolNode([no_update_tool]).invoke(
+            {
+                "messages": [
+                    AIMessage(
+                        "",
+                        tool_calls=[{"args": {}, "id": "1", "name": "no_update_tool"}],
+                    )
+                ]
+            }
+        )
+
+    # test validation (missing tool message in the update for parent graph is OK)
+    @dec_tool
+    def node_update_parent_tool():
+        """No update"""
+        return Command(update={"messages": []}, graph=Command.PARENT)
+
+    assert ToolNode([node_update_parent_tool]).invoke(
+        {
+            "messages": [
+                AIMessage(
+                    "",
+                    tool_calls=[
+                        {"args": {}, "id": "1", "name": "node_update_parent_tool"}
+                    ],
+                )
+            ]
+        }
+    ) == [Command(update={"messages": []}, graph=Command.PARENT)]
+
+    # test validation (multiple tool messages)
+    with pytest.raises(ValueError):
+        for graph in (None, Command.PARENT):
+
+            @dec_tool
+            def multiple_tool_messages_tool():
+                """My tool"""
+                return Command(
+                    update={
+                        "messages": [
+                            ToolMessage(content="foo", tool_call_id=""),
+                            ToolMessage(content="bar", tool_call_id=""),
+                        ]
+                    },
+                    graph=graph,
+                )
+
+            ToolNode([multiple_tool_messages_tool]).invoke(
+                {
+                    "messages": [
+                        AIMessage(
+                            "",
+                            tool_calls=[
+                                {
+                                    "args": {},
+                                    "id": "1",
+                                    "name": "multiple_tool_messages_tool",
+                                }
+                            ],
+                        )
+                    ]
+                }
+            )
+
+
+@pytest.mark.skipif(
+    not IS_LANGCHAIN_CORE_030_OR_GREATER,
+    reason="Langchain core 0.3.0 or greater is required",
+)
+async def test_tool_node_command_list_input():
+    from langchain_core.tools.base import InjectedToolCallId
+
+    @dec_tool
+    def transfer_to_bob(tool_call_id: Annotated[str, InjectedToolCallId]):
+        """Transfer to Bob"""
+        return Command(
+            update=[
+                ToolMessage(content="Transferred to Bob", tool_call_id=tool_call_id)
+            ],
+            goto="bob",
+            graph=Command.PARENT,
+        )
+
+    @dec_tool
+    async def async_transfer_to_bob(tool_call_id: Annotated[str, InjectedToolCallId]):
+        """Transfer to Bob"""
+        return Command(
+            update=[
+                ToolMessage(content="Transferred to Bob", tool_call_id=tool_call_id)
+            ],
+            goto="bob",
+            graph=Command.PARENT,
+        )
+
+    class CustomToolSchema(BaseModel):
+        tool_call_id: Annotated[str, InjectedToolCallId]
+
+    class MyCustomTool(BaseTool):
+        def _run(*args: Any, **kwargs: Any):
+            return Command(
+                update=[
+                    ToolMessage(
+                        content="Transferred to Bob",
+                        tool_call_id=kwargs["tool_call_id"],
+                    )
+                ],
+                goto="bob",
+                graph=Command.PARENT,
+            )
+
+        async def _arun(*args: Any, **kwargs: Any):
+            return Command(
+                update=[
+                    ToolMessage(
+                        content="Transferred to Bob",
+                        tool_call_id=kwargs["tool_call_id"],
+                    )
+                ],
+                goto="bob",
+                graph=Command.PARENT,
+            )
+
+    custom_tool = MyCustomTool(
+        name="custom_transfer_to_bob",
+        description="Transfer to bob",
+        args_schema=CustomToolSchema,
+    )
+    async_custom_tool = MyCustomTool(
+        name="async_custom_transfer_to_bob",
+        description="Transfer to bob",
+        args_schema=CustomToolSchema,
+    )
+
+    # test mixing regular tools and tools returning commands
+    def add(a: int, b: int) -> int:
+        """Add two numbers"""
+        return a + b
+
+    result = ToolNode([add, transfer_to_bob]).invoke(
+        [
+            AIMessage(
+                "",
+                tool_calls=[
+                    {"args": {"a": 1, "b": 2}, "id": "1", "name": "add"},
+                    {"args": {}, "id": "2", "name": "transfer_to_bob"},
+                ],
+            )
+        ]
+    )
+
+    assert result == [
+        [
+            ToolMessage(
+                content="3",
+                tool_call_id="1",
+                name="add",
+            )
+        ],
+        Command(
+            update=[
+                ToolMessage(
+                    content="Transferred to Bob",
+                    tool_call_id="2",
+                    name="transfer_to_bob",
+                )
+            ],
+            goto="bob",
+            graph=Command.PARENT,
+        ),
+    ]
+
+    # test tools returning commands
+
+    # test sync tools
+    for tool in [transfer_to_bob, custom_tool]:
+        result = ToolNode([tool]).invoke(
+            [AIMessage("", tool_calls=[{"args": {}, "id": "1", "name": tool.name}])]
+        )
+        assert result == [
+            Command(
+                update=[
+                    ToolMessage(
+                        content="Transferred to Bob",
+                        tool_call_id="1",
+                        name=tool.name,
+                    )
+                ],
+                goto="bob",
+                graph=Command.PARENT,
+            )
+        ]
+
+    # test async tools
+    for tool in [async_transfer_to_bob, async_custom_tool]:
+        result = await ToolNode([tool]).ainvoke(
+            [AIMessage("", tool_calls=[{"args": {}, "id": "1", "name": tool.name}])]
+        )
+        assert result == [
+            Command(
+                update=[
+                    ToolMessage(
+                        content="Transferred to Bob",
+                        tool_call_id="1",
+                        name=tool.name,
+                    )
+                ],
+                goto="bob",
+                graph=Command.PARENT,
+            )
+        ]
+
+    # test multiple commands
+    result = ToolNode([transfer_to_bob, custom_tool]).invoke(
+        [
+            AIMessage(
+                "",
+                tool_calls=[
+                    {"args": {}, "id": "1", "name": "transfer_to_bob"},
+                    {"args": {}, "id": "2", "name": "custom_transfer_to_bob"},
+                ],
+            )
+        ]
+    )
+    assert result == [
+        Command(
+            update=[
+                ToolMessage(
+                    content="Transferred to Bob",
+                    tool_call_id="1",
+                    name="transfer_to_bob",
+                )
+            ],
+            goto="bob",
+            graph=Command.PARENT,
+        ),
+        Command(
+            update=[
+                ToolMessage(
+                    content="Transferred to Bob",
+                    tool_call_id="2",
+                    name="custom_transfer_to_bob",
+                )
+            ],
+            goto="bob",
+            graph=Command.PARENT,
+        ),
+    ]
+
+    # test validation (mismatch between input type and command.update type)
+    with pytest.raises(ValueError):
+
+        @dec_tool
+        def list_update_tool(tool_call_id: Annotated[str, InjectedToolCallId]):
+            """My tool"""
+            return Command(
+                update={
+                    "messages": [ToolMessage(content="foo", tool_call_id=tool_call_id)]
+                }
+            )
+
+        ToolNode([list_update_tool]).invoke(
+            [
+                AIMessage(
+                    "",
+                    tool_calls=[{"args": {}, "id": "1", "name": "list_update_tool"}],
+                )
+            ]
+        )
+
+    # test validation (missing tool message in the update for current graph)
+    with pytest.raises(ValueError):
+
+        @dec_tool
+        def no_update_tool():
+            """My tool"""
+            return Command(update=[])
+
+        ToolNode([no_update_tool]).invoke(
+            [
+                AIMessage(
+                    "",
+                    tool_calls=[{"args": {}, "id": "1", "name": "no_update_tool"}],
+                )
+            ]
+        )
+
+    # test validation (missing tool message in the update for parent graph is OK)
+    @dec_tool
+    def node_update_parent_tool():
+        """No update"""
+        return Command(update=[], graph=Command.PARENT)
+
+    assert ToolNode([node_update_parent_tool]).invoke(
+        [
+            AIMessage(
+                "",
+                tool_calls=[{"args": {}, "id": "1", "name": "node_update_parent_tool"}],
+            )
+        ]
+    ) == [Command(update=[], graph=Command.PARENT)]
+
+    # test validation (multiple tool messages)
+    with pytest.raises(ValueError):
+        for graph in (None, Command.PARENT):
+
+            @dec_tool
+            def multiple_tool_messages_tool():
+                """My tool"""
+                return Command(
+                    update=[
+                        ToolMessage(content="foo", tool_call_id=""),
+                        ToolMessage(content="bar", tool_call_id=""),
+                    ],
+                    graph=graph,
+                )
+
+            ToolNode([multiple_tool_messages_tool]).invoke(
+                [
+                    AIMessage(
+                        "",
+                        tool_calls=[
+                            {
+                                "args": {},
+                                "id": "1",
+                                "name": "multiple_tool_messages_tool",
+                            }
+                        ],
+                    )
+                ]
+            )
+
+
+@pytest.mark.skipif(
+    not IS_LANGCHAIN_CORE_030_OR_GREATER,
+    reason="Langchain core 0.3.0 or greater is required",
+)
+def test_react_agent_update_state():
+    from langchain_core.tools.base import InjectedToolCallId
+
+    class State(AgentState):
+        user_name: str
+
+    @dec_tool
+    def get_user_name(tool_call_id: Annotated[str, InjectedToolCallId]):
+        """Retrieve user name"""
+        user_name = interrupt("Please provider user name:")
+        return Command(
+            update={
+                "user_name": user_name,
+                "messages": [
+                    ToolMessage(
+                        "Successfully retrieved user name", tool_call_id=tool_call_id
+                    )
+                ],
+            }
+        )
+
+    def state_modifier(state: State):
+        user_name = state.get("user_name")
+        if user_name is None:
+            return state["messages"]
+
+        system_msg = f"User name is {user_name}"
+        return [{"role": "system", "content": system_msg}] + state["messages"]
+
+    checkpointer = MemorySaver()
+    tool_calls = [[{"args": {}, "id": "1", "name": "get_user_name"}]]
+    model = FakeToolCallingModel(tool_calls=tool_calls)
+    agent = create_react_agent(
+        model,
+        [get_user_name],
+        state_schema=State,
+        state_modifier=state_modifier,
+        checkpointer=checkpointer,
+    )
+    config = {"configurable": {"thread_id": "1"}}
+    # run until interrpupted
+    agent.invoke({"messages": [("user", "what's my name")]}, config)
+    # supply the value for the interrupt
+    response = agent.invoke(Command(resume="Archibald"), config)
+    # confirm that the state was updated
+    assert response["user_name"] == "Archibald"
+    assert len(response["messages"]) == 4
+    tool_message: ToolMessage = response["messages"][-2]
+    assert tool_message.content == "Successfully retrieved user name"
+    assert tool_message.tool_call_id == "1"
+    assert tool_message.name == "get_user_name"
+
+
 def my_function(some_val: int, some_other_val: str) -> str:
     return f"{some_val} - {some_other_val}"
 

libs/langgraph/tests/test_pregel.py

@@ -65,6 +65,7 @@ from langgraph.constants import (
     START,
 )
 from langgraph.errors import InvalidUpdateError, MultipleSubgraphsError, NodeInterrupt
+from langgraph.func import entrypoint, task
 from langgraph.graph import END, Graph, StateGraph
 from langgraph.graph.message import MessageGraph, MessagesState, add_messages
 from langgraph.managed.shared_value import SharedValue
@@ -1957,6 +1958,85 @@ def test_send_sequences() -> None:
     )
 
 
+@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC)
+def test_imp_task(request: pytest.FixtureRequest, checkpointer_name: str) -> None:
+    checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}")
+    mapper_calls = 0
+
+    @task()
+    def mapper(input: int) -> str:
+        nonlocal mapper_calls
+        mapper_calls += 1
+        time.sleep(input / 100)
+        return str(input) * 2
+
+    @entrypoint(checkpointer=checkpointer)
+    def graph(input: list[int]) -> list[str]:
+        futures = [mapper(i) for i in input]
+        mapped = [f.result() for f in futures]
+        answer = interrupt("question")
+        return [m + answer for m in mapped]
+
+    thread1 = {"configurable": {"thread_id": "1"}}
+    assert [*graph.stream([0, 1], thread1)] == [
+        {"mapper": "00"},
+        {"mapper": "11"},
+        {
+            "__interrupt__": (
+                Interrupt(
+                    value="question",
+                    resumable=True,
+                    ns=[AnyStr("graph:")],
+                    when="during",
+                ),
+            )
+        },
+    ]
+    assert mapper_calls == 2
+
+    assert graph.invoke(Command(resume="answer"), thread1) == [
+        "00answer",
+        "11answer",
+    ]
+    assert mapper_calls == 2
+
+
+@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC)
+def test_imp_stream_order(
+    request: pytest.FixtureRequest, checkpointer_name: str
+) -> None:
+    checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}")
+
+    @task()
+    def foo(state: dict) -> dict:
+        return {"a": state["a"] + "foo", "b": "bar"}
+
+    @task()
+    def bar(state: dict) -> dict:
+        return {"a": state["a"] + state["b"], "c": "bark"}
+
+    @task()
+    def baz(state: dict) -> dict:
+        return {"a": state["a"] + "baz", "c": "something else"}
+
+    @entrypoint(checkpointer=checkpointer)
+    def graph(state: dict) -> dict:
+        fut_foo = foo(state)
+        fut_bar = bar(fut_foo.result())
+        fut_baz = baz(fut_bar.result())
+        return fut_baz.result()
+
+    thread1 = {"configurable": {"thread_id": "1"}}
+    assert [c for c in graph.stream({"a": "0"}, thread1)] == [
+        {"foo": {"a": "0foo", "b": "bar"}},
+        {"bar": {"a": "0foobar", "c": "bark"}},
+        {"baz": {"a": "0foobarbaz", "c": "something else"}},
+        {"graph": {"a": "0foobarbaz", "c": "something else"}},
+    ]
+
+    assert graph.get_state(thread1).values == {"a": "0foobarbaz", "c": "something else"}
+
+
 @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC)
 def test_send_dedupe_on_resume(
     request: pytest.FixtureRequest, checkpointer_name: str
@@ -2484,7 +2564,7 @@ def test_send_react_interrupt(
             PregelTask(
                 id=AnyStr(),
                 name="foo",
-                path=("__pregel_push", ("__pregel_pull", "agent"), 2, AnyStr()),
+                path=("__pregel_push", ("__pregel_pull", "agent"), 2),
                 error=None,
                 interrupts=(),
                 state=None,
@@ -2641,7 +2721,7 @@ def test_send_react_interrupt(
             PregelTask(
                 id=AnyStr(),
                 name="foo",
-                path=("__pregel_push", ("__pregel_pull", "agent"), 2, AnyStr()),
+                path=("__pregel_push", ("__pregel_pull", "agent"), 2),
                 error=None,
                 interrupts=(),
                 state=None,
@@ -2728,7 +2808,7 @@ def test_send_react_interrupt(
             PregelTask(
                 id=AnyStr(),
                 name="foo",
-                path=("__pregel_push", (), 0, AnyStr()),
+                path=("__pregel_push", (), 0),
                 error=None,
                 interrupts=(),
                 state=None,
@@ -2953,7 +3033,7 @@ def test_send_react_interrupt_control(
             PregelTask(
                 id=AnyStr(),
                 name="foo",
-                path=("__pregel_push", ("__pregel_pull", "agent"), 2, AnyStr()),
+                path=("__pregel_push", ("__pregel_pull", "agent"), 2),
                 error=None,
                 interrupts=(),
                 state=None,
@@ -5745,6 +5825,7 @@ def test_state_graph_packets(
     @tool()
     def search_api(query: str) -> str:
         """Searches the API for the query."""
+        time.sleep(0.1)
         return f"result for {query}"
 
     tools = [search_api]
@@ -6031,9 +6112,7 @@ def test_state_graph_packets(
                     )
                 },
             ),
-            PregelTask(
-                AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2, AnyStr())
-            ),
+            PregelTask(AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2)),
         ),
         next=("tools",),
         config=(app_w_interrupt.checkpointer.get_tuple(config)).config,
@@ -6073,7 +6152,7 @@ def test_state_graph_packets(
                 ),
             ]
         },
-        tasks=(PregelTask(AnyStr(), "tools", (PUSH, (), 0, AnyStr())),),
+        tasks=(PregelTask(AnyStr(), "tools", (PUSH, (), 0)),),
         next=("tools",),
         config=app_w_interrupt.checkpointer.get_tuple(config).config,
         created_at=(app_w_interrupt.checkpointer.get_tuple(config)).checkpoint["ts"],
@@ -6202,12 +6281,8 @@ def test_state_graph_packets(
                     )
                 },
             ),
-            PregelTask(
-                AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2, AnyStr())
-            ),
-            PregelTask(
-                AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 3, AnyStr())
-            ),
+            PregelTask(AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2)),
+            PregelTask(AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 3)),
         ),
         next=("tools", "tools"),
         config=app_w_interrupt.checkpointer.get_tuple(config).config,
@@ -6354,9 +6429,7 @@ def test_state_graph_packets(
                     )
                 },
             ),
-            PregelTask(
-                AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2, AnyStr())
-            ),
+            PregelTask(AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2)),
         ),
         next=("tools",),
         config=(app_w_interrupt.checkpointer.get_tuple(config)).config,
@@ -6396,7 +6469,7 @@ def test_state_graph_packets(
                 ),
             ]
         },
-        tasks=(PregelTask(AnyStr(), "tools", (PUSH, (), 0, AnyStr())),),
+        tasks=(PregelTask(AnyStr(), "tools", (PUSH, (), 0)),),
         next=("tools",),
         config=app_w_interrupt.checkpointer.get_tuple(config).config,
         created_at=(app_w_interrupt.checkpointer.get_tuple(config)).checkpoint["ts"],
@@ -6525,12 +6598,8 @@ def test_state_graph_packets(
                     )
                 },
             ),
-            PregelTask(
-                AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2, AnyStr())
-            ),
-            PregelTask(
-                AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 3, AnyStr())
-            ),
+            PregelTask(AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2)),
+            PregelTask(AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 3)),
         ),
         next=("tools", "tools"),
         config=app_w_interrupt.checkpointer.get_tuple(config).config,
@@ -7533,7 +7602,7 @@ def test_root_graph(
                     content="result for query",
                     name="search_api",
                     tool_call_id="tool_call123",
-                    id="00000000-0000-4000-8000-000000000033",
+                    id="00000000-0000-4000-8000-000000000037",
                 )
             ]
         },
@@ -7556,7 +7625,7 @@ def test_root_graph(
                     content="result for another",
                     name="search_api",
                     tool_call_id="tool_call456",
-                    id="00000000-0000-4000-8000-000000000041",
+                    id="00000000-0000-4000-8000-000000000045",
                 )
             ]
         },
@@ -8166,7 +8235,7 @@ def test_root_graph(
         "__root__": [
             HumanMessage(
                 content="what is weather in sf",
-                id="00000000-0000-4000-8000-000000000070",
+                id="00000000-0000-4000-8000-000000000078",
             ),
             AIMessage(
                 content="",
@@ -8186,7 +8255,7 @@ def test_root_graph(
             ),
             AIMessage(content="answer", id="ai2"),
             AIMessage(
-                content="an extra message", id="00000000-0000-4000-8000-000000000092"
+                content="an extra message", id="00000000-0000-4000-8000-000000000100"
             ),
             HumanMessage(content="what is weather in la"),
         ],
@@ -12774,7 +12843,7 @@ def test_send_to_nested_graphs(
             PregelTask(
                 AnyStr(),
                 "generate_joke",
-                (PUSH, ("__pregel_pull", "__start__"), 1, AnyStr()),
+                (PUSH, ("__pregel_pull", "__start__"), 1),
                 state={
                     "configurable": {
                         "thread_id": "1",
@@ -12785,7 +12854,7 @@ def test_send_to_nested_graphs(
             PregelTask(
                 AnyStr(),
                 "generate_joke",
-                (PUSH, ("__pregel_pull", "__start__"), 2, AnyStr()),
+                (PUSH, ("__pregel_pull", "__start__"), 2),
                 state={
                     "configurable": {
                         "thread_id": "1",
@@ -12832,13 +12901,13 @@ def test_send_to_nested_graphs(
         metadata={
             "step": 1,
             "source": "loop",
-            "writes": {"edit": None},
+            "writes": None,
             "parents": {"": AnyStr()},
             "thread_id": "1",
             "checkpoint_ns": AnyStr("generate_joke:"),
             "langgraph_checkpoint_ns": AnyStr("generate_joke:"),
             "langgraph_node": "generate_joke",
-            "langgraph_path": [PUSH, ["__pregel_pull", "__start__"], 1, AnyStr()],
+            "langgraph_path": [PUSH, ["__pregel_pull", "__start__"], 1],
             "langgraph_step": 0,
             "langgraph_triggers": [PUSH],
         },
@@ -12877,13 +12946,13 @@ def test_send_to_nested_graphs(
         metadata={
             "step": 1,
             "source": "loop",
-            "writes": {"edit": None},
+            "writes": None,
             "parents": {"": AnyStr()},
             "thread_id": "1",
             "checkpoint_ns": AnyStr("generate_joke:"),
             "langgraph_checkpoint_ns": AnyStr("generate_joke:"),
             "langgraph_node": "generate_joke",
-            "langgraph_path": [PUSH, ["__pregel_pull", "__start__"], 2, AnyStr()],
+            "langgraph_path": [PUSH, ["__pregel_pull", "__start__"], 2],
             "langgraph_step": 0,
             "langgraph_triggers": [PUSH],
         },
@@ -13009,7 +13078,7 @@ def test_send_to_nested_graphs(
                 PregelTask(
                     AnyStr(),
                     "generate_joke",
-                    (PUSH, ("__pregel_pull", "__start__"), 1, AnyStr()),
+                    (PUSH, ("__pregel_pull", "__start__"), 1),
                     state={
                         "configurable": {
                             "thread_id": "1",
@@ -13021,7 +13090,7 @@ def test_send_to_nested_graphs(
                 PregelTask(
                     AnyStr(),
                     "generate_joke",
-                    (PUSH, ("__pregel_pull", "__start__"), 2, AnyStr()),
+                    (PUSH, ("__pregel_pull", "__start__"), 2),
                     state={
                         "configurable": {
                             "thread_id": "1",
@@ -14807,3 +14876,169 @@ def test_interrupt_loop(request: pytest.FixtureRequest, checkpointer_name: str):
     assert [event for event in graph.stream(Command(resume="19"), thread1)] == [
         {"node": {"age": 19}},
     ]
+
+
+def test_root_mixed_return() -> None:
+    def my_node(state: list[str]):
+        return [Command(update=["a"]), ["b"]]
+
+    graph = StateGraph(Annotated[list[str], operator.add])
+
+    graph.add_node(my_node)
+    graph.add_edge(START, "my_node")
+    graph = graph.compile()
+
+    assert graph.invoke([]) == ["a", "b"]
+
+
+def test_dict_mixed_return() -> None:
+    class State(TypedDict):
+        foo: Annotated[str, operator.add]
+
+    def my_node(state: State):
+        return [Command(update={"foo": "a"}), {"foo": "b"}]
+
+    graph = StateGraph(State)
+    graph.add_node(my_node)
+    graph.add_edge(START, "my_node")
+    graph = graph.compile()
+
+    assert graph.invoke({"foo": ""}) == {"foo": "ab"}
+
+
+@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC)
+def test_command_with_static_breakpoints(
+    request: pytest.FixtureRequest, checkpointer_name: str
+) -> None:
+    """Test that we can use Command to resume and update with static breakpoints."""
+
+    checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}")
+
+    class State(TypedDict):
+        """The graph state."""
+
+        foo: str
+
+    def node1(state: State):
+        return {
+            "foo": state["foo"] + "|node-1",
+        }
+
+    def node2(state: State):
+        return {
+            "foo": state["foo"] + "|node-2",
+        }
+
+    builder = StateGraph(State)
+    builder.add_node("node1", node1)
+    builder.add_node("node2", node2)
+    builder.add_edge(START, "node1")
+    builder.add_edge("node1", "node2")
+
+    graph = builder.compile(checkpointer=checkpointer, interrupt_before=["node1"])
+    config = {"configurable": {"thread_id": str(uuid.uuid4())}}
+
+    # Start the graph and interrupt at the first node
+    graph.invoke({"foo": "abc"}, config)
+    result = graph.invoke(Command(resume="node1"), config)
+    assert result == {"foo": "abc|node-1|node-2"}
+
+
+@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC)
+def test_multistep_plan(request: pytest.FixtureRequest, checkpointer_name: str):
+    from langchain_core.messages import AnyMessage
+
+    checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}")
+
+    class State(TypedDict, total=False):
+        plan: list[Union[str, list[str]]]
+        messages: Annotated[list[AnyMessage], add_messages]
+
+    def planner(state: State):
+        if state.get("plan") is None:
+            # create plan somehow
+            plan = ["step1", ["step2", "step3"], "step4"]
+            # pick the first step to execute next
+            first_step, *plan = plan
+            # put the rest of plan in state
+            return Command(goto=first_step, update={"plan": plan})
+        elif state["plan"]:
+            # go to the next step of the plan
+            next_step, *next_plan = state["plan"]
+            return Command(goto=next_step, update={"plan": next_plan})
+        else:
+            # the end of the plan
+            pass
+
+    def step1(state: State):
+        return Command(goto="planner", update={"messages": [("human", "step1")]})
+
+    def step2(state: State):
+        return Command(goto="planner", update={"messages": [("human", "step2")]})
+
+    def step3(state: State):
+        return Command(goto="planner", update={"messages": [("human", "step3")]})
+
+    def step4(state: State):
+        return Command(goto="planner", update={"messages": [("human", "step4")]})
+
+    builder = StateGraph(State)
+    builder.add_node(planner)
+    builder.add_node(step1)
+    builder.add_node(step2)
+    builder.add_node(step3)
+    builder.add_node(step4)
+    builder.add_edge(START, "planner")
+    graph = builder.compile(checkpointer=checkpointer)
+
+    config = {"configurable": {"thread_id": "1"}}
+
+    assert graph.invoke({"messages": [("human", "start")]}, config) == {
+        "messages": [
+            _AnyIdHumanMessage(content="start"),
+            _AnyIdHumanMessage(content="step1"),
+            _AnyIdHumanMessage(content="step2"),
+            _AnyIdHumanMessage(content="step3"),
+            _AnyIdHumanMessage(content="step4"),
+        ],
+        "plan": [],
+    }
+
+
+@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC)
+def test_command_goto_with_static_breakpoints(
+    request: pytest.FixtureRequest, checkpointer_name: str
+) -> None:
+    """Use Command goto with static breakpoints."""
+
+    checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}")
+
+    class State(TypedDict):
+        """The graph state."""
+
+        foo: Annotated[str, operator.add]
+
+    def node1(state: State):
+        return {
+            "foo": "|node-1",
+        }
+
+    def node2(state: State):
+        return {
+            "foo": "|node-2",
+        }
+
+    builder = StateGraph(State)
+    builder.add_node("node1", node1)
+    builder.add_node("node2", node2)
+    builder.add_edge(START, "node1")
+    builder.add_edge("node1", "node2")
+
+    graph = builder.compile(checkpointer=checkpointer, interrupt_before=["node1"])
+
+    config = {"configurable": {"thread_id": str(uuid.uuid4())}}
+
+    # Start the graph and interrupt at the first node
+    graph.invoke({"foo": "abc"}, config)
+    result = graph.invoke(Command(goto=["node2"]), config)
+    assert result == {"foo": "abc|node-1|node-2|node-2"}

libs/langgraph/tests/test_pregel_async.py

@@ -62,6 +62,7 @@ from langgraph.constants import (
     START,
 )
 from langgraph.errors import InvalidUpdateError, MultipleSubgraphsError, NodeInterrupt
+from langgraph.func import entrypoint, task
 from langgraph.graph import END, Graph, StateGraph
 from langgraph.graph.message import MessageGraph, MessagesState, add_messages
 from langgraph.managed.shared_value import SharedValue
@@ -2647,6 +2648,178 @@ async def test_send_sequences(checkpointer_name: str) -> None:
         ]
 
 
+@pytest.mark.skipif(
+    sys.version_info < (3, 11),
+    reason="Python 3.11+ is required for async contextvars support",
+)
+@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
+async def test_imp_task(checkpointer_name: str) -> None:
+    async with awith_checkpointer(checkpointer_name) as checkpointer:
+        mapper_calls = 0
+
+        @task()
+        async def mapper(input: int) -> str:
+            nonlocal mapper_calls
+            mapper_calls += 1
+            return str(input) * 2
+
+        @entrypoint(checkpointer=checkpointer)
+        async def graph(input: list[int]) -> list[str]:
+            futures = [mapper(i) for i in input]
+            mapped = await asyncio.gather(*futures)
+            answer = interrupt("question")
+            return [m + answer for m in mapped]
+
+        thread1 = {"configurable": {"thread_id": "1"}}
+        assert [c async for c in graph.astream([0, 1], thread1)] == [
+            {"mapper": "00"},
+            {"mapper": "11"},
+            {
+                "__interrupt__": (
+                    Interrupt(
+                        value="question",
+                        resumable=True,
+                        ns=[AnyStr("graph:")],
+                        when="during",
+                    ),
+                )
+            },
+        ]
+        assert mapper_calls == 2
+
+        assert await graph.ainvoke(Command(resume="answer"), thread1) == [
+            "00answer",
+            "11answer",
+        ]
+        assert mapper_calls == 2
+
+
+@pytest.mark.skipif(
+    sys.version_info < (3, 11),
+    reason="Python 3.11+ is required for async contextvars support",
+)
+@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
+async def test_imp_task_cancel(checkpointer_name: str) -> None:
+    async with awith_checkpointer(checkpointer_name) as checkpointer:
+        mapper_calls = 0
+        mapper_cancels = 0
+
+        @task()
+        async def mapper(input: int) -> str:
+            nonlocal mapper_calls, mapper_cancels
+            mapper_calls += 1
+            try:
+                await asyncio.sleep(1)
+            except asyncio.CancelledError:
+                mapper_cancels += 1
+                raise
+            return str(input) * 2
+
+        @entrypoint(checkpointer=checkpointer)
+        async def graph(input: list[int]) -> list[str]:
+            futures = [mapper(i) for i in input]
+            await asyncio.sleep(0.1)
+            futures.pop().cancel()  # cancel one
+            mapped = await asyncio.gather(*futures)
+            answer = interrupt("question")
+            return [m + answer for m in mapped]
+
+        thread1 = {"configurable": {"thread_id": "1"}}
+        assert [c async for c in graph.astream([0, 1], thread1)] == [
+            {"mapper": "00"},
+            {
+                "__interrupt__": (
+                    Interrupt(
+                        value="question",
+                        resumable=True,
+                        ns=[AnyStr("graph:")],
+                        when="during",
+                    ),
+                )
+            },
+        ]
+        assert mapper_calls == 2
+        assert mapper_cancels == 1
+
+        assert await graph.ainvoke(Command(resume="answer"), thread1) == [
+            "00answer",
+        ]
+        assert mapper_calls == 3
+        assert mapper_cancels == 2
+
+
+@pytest.mark.skipif(
+    sys.version_info < (3, 11),
+    reason="Python 3.11+ is required for async contextvars support",
+)
+@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
+async def test_imp_sync_from_async(checkpointer_name: str) -> None:
+    async with awith_checkpointer(checkpointer_name) as checkpointer:
+
+        @task()
+        def foo(state: dict) -> dict:
+            return {"a": state["a"] + "foo", "b": "bar"}
+
+        @task()
+        def bar(state: dict) -> dict:
+            return {"a": state["a"] + state["b"], "c": "bark"}
+
+        @task()
+        def baz(state: dict) -> dict:
+            return {"a": state["a"] + "baz", "c": "something else"}
+
+        @entrypoint(checkpointer=checkpointer)
+        def graph(state: dict) -> dict:
+            fut_foo = foo(state)
+            fut_bar = bar(fut_foo.result())
+            fut_baz = baz(fut_bar.result())
+            return fut_baz.result()
+
+        thread1 = {"configurable": {"thread_id": "1"}}
+        assert [c async for c in graph.astream({"a": "0"}, thread1)] == [
+            {"foo": {"a": "0foo", "b": "bar"}},
+            {"bar": {"a": "0foobar", "c": "bark"}},
+            {"baz": {"a": "0foobarbaz", "c": "something else"}},
+            {"graph": {"a": "0foobarbaz", "c": "something else"}},
+        ]
+
+
+@pytest.mark.skipif(
+    sys.version_info < (3, 11),
+    reason="Python 3.11+ is required for async contextvars support",
+)
+@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
+async def test_imp_stream_order(checkpointer_name: str) -> None:
+    async with awith_checkpointer(checkpointer_name) as checkpointer:
+
+        @task()
+        async def foo(state: dict) -> dict:
+            return {"a": state["a"] + "foo", "b": "bar"}
+
+        @task()
+        async def bar(state: dict) -> dict:
+            return {"a": state["a"] + state["b"], "c": "bark"}
+
+        @task()
+        async def baz(state: dict) -> dict:
+            return {"a": state["a"] + "baz", "c": "something else"}
+
+        @entrypoint(checkpointer=checkpointer)
+        async def graph(state: dict) -> dict:
+            fut_foo = foo(state)
+            fut_bar = bar(await fut_foo)
+            fut_baz = baz(await fut_bar)
+            return await fut_baz
+
+        thread1 = {"configurable": {"thread_id": "1"}}
+        assert [c async for c in graph.astream({"a": "0"}, thread1)] == [
+            {"foo": {"a": "0foo", "b": "bar"}},
+            {"bar": {"a": "0foobar", "c": "bark"}},
+            {"baz": {"a": "0foobarbaz", "c": "something else"}},
+            {"graph": {"a": "0foobarbaz", "c": "something else"}},
+        ]
+
+
 @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
 async def test_send_dedupe_on_resume(checkpointer_name: str) -> None:
     if not FF_SEND_V2:
@@ -2864,12 +3037,7 @@ async def test_send_dedupe_on_resume(checkpointer_name: str) -> None:
                     PregelTask(
                         id=AnyStr(),
                         name="2",
-                        path=(
-                            "__pregel_push",
-                            ("__pregel_pull", "1"),
-                            2,
-                            AnyStr(),
-                        ),
+                        path=("__pregel_push", ("__pregel_pull", "1"), 2),
                         error=None,
                         interrupts=(),
                         state=None,
@@ -2878,12 +3046,7 @@ async def test_send_dedupe_on_resume(checkpointer_name: str) -> None:
                     PregelTask(
                         id=AnyStr(),
                         name="2",
-                        path=(
-                            "__pregel_push",
-                            ("__pregel_pull", "1"),
-                            3,
-                            AnyStr(),
-                        ),
+                        path=("__pregel_push", ("__pregel_pull", "1"), 3),
                         error=None,
                         interrupts=(),
                         state=None,
@@ -2894,14 +3057,8 @@ async def test_send_dedupe_on_resume(checkpointer_name: str) -> None:
                         name="2",
                         path=(
                             "__pregel_push",
-                            (
-                                "__pregel_push",
-                                ("__pregel_pull", "1"),
-                                2,
-                                AnyStr(),
-                            ),
+                            ("__pregel_push", ("__pregel_pull", "1"), 2),
                             2,
-                            AnyStr(),
                         ),
                         error=None,
                         interrupts=(),
@@ -2913,14 +3070,8 @@ async def test_send_dedupe_on_resume(checkpointer_name: str) -> None:
                         name="flaky",
                         path=(
                             "__pregel_push",
-                            (
-                                "__pregel_push",
-                                ("__pregel_pull", "1"),
-                                3,
-                                AnyStr(),
-                            ),
+                            ("__pregel_push", ("__pregel_pull", "1"), 3),
                             2,
-                            AnyStr(),
                         ),
                         error=None,
                         interrupts=(Interrupt(value="Bahh", when="during"),),
@@ -3157,7 +3308,7 @@ async def test_send_react_interrupt(checkpointer_name: str) -> None:
                 PregelTask(
                     id=AnyStr(),
                     name="foo",
-                    path=("__pregel_push", ("__pregel_pull", "agent"), 2, AnyStr()),
+                    path=("__pregel_push", ("__pregel_pull", "agent"), 2),
                     error=None,
                     interrupts=(),
                     state=None,
@@ -3314,7 +3465,7 @@ async def test_send_react_interrupt(checkpointer_name: str) -> None:
                 PregelTask(
                     id=AnyStr(),
                     name="foo",
-                    path=("__pregel_push", ("__pregel_pull", "agent"), 2, AnyStr()),
+                    path=("__pregel_push", ("__pregel_pull", "agent"), 2),
                     error=None,
                     interrupts=(),
                     state=None,
@@ -3401,7 +3552,7 @@ async def test_send_react_interrupt(checkpointer_name: str) -> None:
                 PregelTask(
                     id=AnyStr(),
                     name="foo",
-                    path=("__pregel_push", (), 0, AnyStr()),
+                    path=("__pregel_push", (), 0),
                     error=None,
                     interrupts=(),
                     state=None,
@@ -3625,7 +3776,7 @@ async def test_send_react_interrupt_control(
                 PregelTask(
                     id=AnyStr(),
                     name="foo",
-                    path=("__pregel_push", ("__pregel_pull", "agent"), 2, AnyStr()),
+                    path=("__pregel_push", ("__pregel_pull", "agent"), 2),
                     error=None,
                     interrupts=(),
                     state=None,
@@ -6420,9 +6571,7 @@ async def test_state_graph_packets(checkpointer_name: str) -> None:
                         )
                     },
                 ),
-                PregelTask(
-                    AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2, AnyStr())
-                ),
+                PregelTask(AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2)),
             ),
             next=("tools",),
             config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config,
@@ -6465,7 +6614,7 @@ async def test_state_graph_packets(checkpointer_name: str) -> None:
                     ),
                 ]
             },
-            tasks=(PregelTask(AnyStr(), "tools", (PUSH, (), 0, AnyStr())),),
+            tasks=(PregelTask(AnyStr(), "tools", (PUSH, (), 0)),),
             next=("tools",),
             config=tup.config,
             created_at=tup.checkpoint["ts"],
@@ -6596,12 +6745,8 @@ async def test_state_graph_packets(checkpointer_name: str) -> None:
                         )
                     },
                 ),
-                PregelTask(
-                    AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2, AnyStr())
-                ),
-                PregelTask(
-                    AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 3, AnyStr())
-                ),
+                PregelTask(AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2)),
+                PregelTask(AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 3)),
             ),
             next=("tools", "tools"),
             config=tup.config,
@@ -6751,9 +6896,7 @@ async def test_state_graph_packets(checkpointer_name: str) -> None:
                         )
                     },
                 ),
-                PregelTask(
-                    AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2, AnyStr())
-                ),
+                PregelTask(AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2)),
             ),
             next=("tools",),
             config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config,
@@ -6796,7 +6939,7 @@ async def test_state_graph_packets(checkpointer_name: str) -> None:
                     ),
                 ]
             },
-            tasks=(PregelTask(AnyStr(), "tools", (PUSH, (), 0, AnyStr())),),
+            tasks=(PregelTask(AnyStr(), "tools", (PUSH, (), 0)),),
             next=("tools",),
             config=tup.config,
             created_at=tup.checkpoint["ts"],
@@ -6929,12 +7072,8 @@ async def test_state_graph_packets(checkpointer_name: str) -> None:
                         )
                     },
                 ),
-                PregelTask(
-                    AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2, AnyStr())
-                ),
-                PregelTask(
-                    AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 3, AnyStr())
-                ),
+                PregelTask(AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2)),
+                PregelTask(AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 3)),
             ),
             next=("tools", "tools"),
             config=tup.config,
@@ -9670,14 +9809,14 @@ async def test_stream_subgraphs_during_execution(checkpointer_name: str) -> None
             ),
             (FloatBetween(0.2, 0.4), ((), {"outer_1": {"my_key": " and parallel"}})),
             (
-                FloatBetween(0.5, 0.7),
+                FloatBetween(0.5, 0.8),
                 (
                     (AnyStr("inner:"),),
                     {"inner_2": {"my_key": " and there", "my_other_key": "got here"}},
                 ),
             ),
-            (FloatBetween(0.5, 0.7), ((), {"inner": {"my_key": "got here and there"}})),
-            (FloatBetween(0.5, 0.7), ((), {"outer_2": {"my_key": " and back again"}})),
+            (FloatBetween(0.5, 0.8), ((), {"inner": {"my_key": "got here and there"}})),
+            (FloatBetween(0.5, 0.8), ((), {"outer_2": {"my_key": " and back again"}})),
         ]
 
 
@@ -11612,7 +11751,7 @@ async def test_send_to_nested_graphs(checkpointer_name: str) -> None:
                 PregelTask(
                     AnyStr(),
                     "generate_joke",
-                    (PUSH, ("__pregel_pull", "__start__"), 1, AnyStr()),
+                    (PUSH, ("__pregel_pull", "__start__"), 1),
                     state={
                         "configurable": {
                             "thread_id": "1",
@@ -11623,7 +11762,7 @@ async def test_send_to_nested_graphs(checkpointer_name: str) -> None:
                 PregelTask(
                     AnyStr(),
                     "generate_joke",
-                    (PUSH, ("__pregel_pull", "__start__"), 2, AnyStr()),
+                    (PUSH, ("__pregel_pull", "__start__"), 2),
                     state={
                         "configurable": {
                             "thread_id": "1",
@@ -11764,7 +11903,7 @@ async def test_send_to_nested_graphs(checkpointer_name: str) -> None:
                     PregelTask(
                         AnyStr(),
                         "generate_joke",
-                        (PUSH, ("__pregel_pull", "__start__"), 1, AnyStr()),
+                        (PUSH, ("__pregel_pull", "__start__"), 1),
                         state={
                             "configurable": {
                                 "thread_id": "1",
@@ -11776,7 +11915,7 @@ async def test_send_to_nested_graphs(checkpointer_name: str) -> None:
                     PregelTask(
                         AnyStr(),
                         "generate_joke",
-                        (PUSH, ("__pregel_pull", "__start__"), 2, AnyStr()),
+                        (PUSH, ("__pregel_pull", "__start__"), 2),
                         state={
                             "configurable": {
                                 "thread_id": "1",
@@ -12482,9 +12621,19 @@ async def test_store_injected_async(checkpointer_name: str, store_name: str) ->
             )
             return {"count": 1}
 
+    def other_node(inputs: State, config: RunnableConfig, store: BaseStore):
+        assert isinstance(store, BaseStore)
+        store.put(("not", "interesting"), "key", {"val": "val"})
+        item = store.get(("not", "interesting"), "key")
+        assert item is not None
+        assert item.value == {"val": "val"}
+        return {"count": 0}
+
     builder = StateGraph(State)
     builder.add_node("node", Node())
+    builder.add_node("other_node", other_node)
     builder.add_edge("__start__", "node")
+    builder.add_edge("node", "other_node")
 
     N = 500
     M = 1
@@ -13050,3 +13199,135 @@ async def test_interrupt_loop(checkpointer_name: str):
         ] == [
             {"node": {"age": 19}},
         ]
+
+
+@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
+async def test_command_with_static_breakpoints(checkpointer_name: str) -> None:
+    """Test that we can use Command to resume and update with static breakpoints."""
+
+    class State(TypedDict):
+        """The graph state."""
+
+        foo: str
+
+    def node1(state: State):
+        return {
+            "foo": state["foo"] + "|node-1",
+        }
+
+    def node2(state: State):
+        return {
+            "foo": state["foo"] + "|node-2",
+        }
+
+    builder = StateGraph(State)
+    builder.add_node("node1", node1)
+    builder.add_node("node2", node2)
+    builder.add_edge(START, "node1")
+    builder.add_edge("node1", "node2")
+
+    async with awith_checkpointer(checkpointer_name) as checkpointer:
+        graph = builder.compile(checkpointer=checkpointer, interrupt_before=["node1"])
+        config = {"configurable": {"thread_id": str(uuid.uuid4())}}
+
+        # Start the graph and interrupt at the first node
+        await graph.ainvoke({"foo": "abc"}, config)
+        result = await graph.ainvoke(Command(update={"foo": "def"}), config)
+        assert result == {"foo": "def|node-1|node-2"}
+
+
+@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
+async def test_multistep_plan(checkpointer_name: str):
+    from langchain_core.messages import AnyMessage
+
+    class State(TypedDict, total=False):
+        plan: list[Union[str, list[str]]]
+        messages: Annotated[list[AnyMessage], add_messages]
+
+    def planner(state: State):
+        if state.get("plan") is None:
+            # create plan somehow
+            plan = ["step1", ["step2", "step3"], "step4"]
+            # pick the first step to execute next
+            first_step, *plan = plan
+            # put the rest of plan in state
+            return Command(goto=first_step, update={"plan": plan})
+        elif state["plan"]:
+            # go to the next step of the plan
+            next_step, *next_plan = state["plan"]
+            return Command(goto=next_step, update={"plan": next_plan})
+        else:
+            # the end of the plan
+            pass
+
+    def step1(state: State):
+        return Command(goto="planner", update={"messages": [("human", "step1")]})
+
+    def step2(state: State):
+        return Command(goto="planner", update={"messages": [("human", "step2")]})
+
+    def step3(state: State):
+        return Command(goto="planner", update={"messages": [("human", "step3")]})
+
+    def step4(state: State):
+        return Command(goto="planner", update={"messages": [("human", "step4")]})
+
+    builder = StateGraph(State)
+    builder.add_node(planner)
+    builder.add_node(step1)
+    builder.add_node(step2)
+    builder.add_node(step3)
+    builder.add_node(step4)
+    builder.add_edge(START, "planner")
+
+    async with awith_checkpointer(checkpointer_name) as checkpointer:
+        graph = builder.compile(checkpointer=checkpointer)
+
+        config = {"configurable": {"thread_id": "1"}}
+
+        assert await graph.ainvoke({"messages": [("human", "start")]}, config) == {
+            "messages": [
+                _AnyIdHumanMessage(content="start"),
+                _AnyIdHumanMessage(content="step1"),
+                _AnyIdHumanMessage(content="step2"),
+                _AnyIdHumanMessage(content="step3"),
+                _AnyIdHumanMessage(content="step4"),
+            ],
+            "plan": [],
+        }
+
+
+@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
+async def test_command_goto_with_static_breakpoints(checkpointer_name: str) -> None:
+    """Use Command goto with static breakpoints."""
+
+    class State(TypedDict):
+        """The graph state."""
+
+        foo: Annotated[str, operator.add]
+
+    def node1(state: State):
+        return {
+            "foo": "|node-1",
+        }
+
+    def node2(state: State):
+        return {
+            "foo": "|node-2",
+        }
+
+    builder = StateGraph(State)
+    builder.add_node("node1", node1)
+    builder.add_node("node2", node2)
+    builder.add_edge(START, "node1")
+    builder.add_edge("node1", "node2")
+
+    async with awith_checkpointer(checkpointer_name) as checkpointer:
+        graph = builder.compile(checkpointer=checkpointer, interrupt_before=["node1"])
+
+        config = {"configurable": {"thread_id": str(uuid.uuid4())}}
+
+        # Start the graph and interrupt at the first node
+        await graph.ainvoke({"foo": "abc"}, config)
+        result = await graph.ainvoke(Command(goto=["node2"]), config)
+        assert result == {"foo": "abc|node-1|node-2|node-2"}

libs/langgraph/tests/test_utils.py

@@ -21,7 +21,11 @@ from typing_extensions import Annotated, NotRequired, Required
 
 from langgraph.graph import END, StateGraph
 from langgraph.graph.graph import CompiledGraph
-from langgraph.utils.fields import _is_optional_type, get_field_default
+from langgraph.utils.fields import (
+    _is_optional_type,
+    get_enhanced_type_hints,
+    get_field_default,
+)
 from langgraph.utils.runnable import is_async_callable, is_async_generator
 
 pytestmark = pytest.mark.anyio
@@ -227,3 +231,57 @@ def test_is_required():
     assert get_field_default("val_12", gcannos["val_12"], MyGrandChildDict) is None
     assert get_field_default("val_9", gcannos["val_9"], MyGrandChildDict) is None
     assert get_field_default("val_13", gcannos["val_13"], MyGrandChildDict) == ...
+
+
+def test_enhanced_type_hints() -> None:
+    from dataclasses import dataclass
+    from typing import Annotated
+
+    from pydantic import BaseModel, Field
+
+    class MyTypedDict(TypedDict):
+        val_1: str
+        val_2: int = 42
+        val_3: str = "default"
+
+    hints = list(get_enhanced_type_hints(MyTypedDict))
+    assert len(hints) == 3
+    assert hints[0] == ("val_1", str, None, None)
+    assert hints[1] == ("val_2", int, 42, None)
+    assert hints[2] == ("val_3", str, "default", None)
+
+    @dataclass
+    class MyDataclass:
+        val_1: str
+        val_2: int = 42
+        val_3: str = "default"
+
+    hints = list(get_enhanced_type_hints(MyDataclass))
+    assert len(hints) == 3
+    assert hints[0] == ("val_1", str, None, None)
+    assert hints[1] == ("val_2", int, 42, None)
+    assert hints[2] == ("val_3", str, "default", None)
+
+    class MyPydanticModel(BaseModel):
+        val_1: str
+        val_2: int = 42
+        val_3: str = Field(default="default", description="A description")
+
+    hints = list(get_enhanced_type_hints(MyPydanticModel))
+    assert len(hints) == 3
+    assert hints[0] == ("val_1", str, None, None)
+    assert hints[1] == ("val_2", int, 42, None)
+    assert hints[2] == ("val_3", str, "default", "A description")
+
+    class MyPydanticModelWithAnnotated(BaseModel):
+        val_1: Annotated[str, Field(description="A description")]
+        val_2: Annotated[int, Field(default=42)]
+        val_3: Annotated[
+            str, Field(default="default", description="Another description")
+        ]
+
+    hints = list(get_enhanced_type_hints(MyPydanticModelWithAnnotated))
+    assert len(hints) == 3
+    assert hints[0] == ("val_1", str, None, "A description")
+    assert hints[1] == ("val_2", int, 42, None)
+    assert hints[2] == ("val_3", str, "default", "Another description")

libs/scheduler-kafka/tests/test_subgraph.py

@@ -191,6 +191,7 @@ async def test_subgraph_w_interrupt(
                             "__pregel_delegate": False,
                             "__pregel_read": None,
                             "__pregel_send": None,
+                            "__pregel_call": None,
                             "__pregel_ensure_latest": True,
                             "__pregel_dedupe_tasks": True,
                             "__pregel_resuming": False,
@@ -257,6 +258,7 @@ async def test_subgraph_w_interrupt(
                             "__pregel_delegate": False,
                             "__pregel_read": None,
                             "__pregel_send": None,
+                            "__pregel_call": None,
                             "__pregel_ensure_latest": True,
                             "__pregel_dedupe_tasks": True,
                             "__pregel_resuming": False,
@@ -353,6 +355,7 @@ async def test_subgraph_w_interrupt(
                             "__pregel_delegate": False,
                             "__pregel_read": None,
                             "__pregel_send": None,
+                            "__pregel_call": None,
                             "__pregel_ensure_latest": True,
                             "__pregel_dedupe_tasks": True,
                             "__pregel_resuming": False,
@@ -459,6 +462,7 @@ async def test_subgraph_w_interrupt(
                             "__pregel_delegate": False,
                             "__pregel_read": None,
                             "__pregel_send": None,
+                            "__pregel_call": None,
                             "__pregel_ensure_latest": True,
                             "__pregel_dedupe_tasks": True,
                             "__pregel_resuming": True,
@@ -520,6 +524,7 @@ async def test_subgraph_w_interrupt(
                             "__pregel_delegate": False,
                             "__pregel_read": None,
                             "__pregel_send": None,
+                            "__pregel_call": None,
                             "__pregel_ensure_latest": True,
                             "__pregel_dedupe_tasks": True,
                             "__pregel_resuming": True,
@@ -637,6 +642,7 @@ async def test_subgraph_w_interrupt(
                             "__pregel_delegate": False,
                             "__pregel_read": None,
                             "__pregel_send": None,
+                            "__pregel_call": None,
                             "__pregel_ensure_latest": True,
                             "__pregel_dedupe_tasks": True,
                             "__pregel_resuming": True,

libs/scheduler-kafka/tests/test_subgraph_sync.py

@@ -190,6 +190,7 @@ def test_subgraph_w_interrupt(
                             "__pregel_delegate": False,
                             "__pregel_read": None,
                             "__pregel_send": None,
+                            "__pregel_call": None,
                             "__pregel_ensure_latest": True,
                             "__pregel_dedupe_tasks": True,
                             "__pregel_resuming": False,
@@ -256,6 +257,7 @@ def test_subgraph_w_interrupt(
                             "__pregel_delegate": False,
                             "__pregel_read": None,
                             "__pregel_send": None,
+                            "__pregel_call": None,
                             "__pregel_ensure_latest": True,
                             "__pregel_store": None,
                             "__pregel_dedupe_tasks": True,
@@ -352,6 +354,7 @@ def test_subgraph_w_interrupt(
                             "__pregel_delegate": False,
                             "__pregel_read": None,
                             "__pregel_send": None,
+                            "__pregel_call": None,
                             "__pregel_ensure_latest": True,
                             "__pregel_dedupe_tasks": True,
                             "__pregel_store": None,
@@ -457,6 +460,7 @@ def test_subgraph_w_interrupt(
                             "__pregel_delegate": False,
                             "__pregel_read": None,
                             "__pregel_send": None,
+                            "__pregel_call": None,
                             "__pregel_ensure_latest": True,
                             "__pregel_dedupe_tasks": True,
                             "__pregel_store": None,
@@ -518,6 +522,7 @@ def test_subgraph_w_interrupt(
                             "__pregel_delegate": False,
                             "__pregel_read": None,
                             "__pregel_send": None,
+                            "__pregel_call": None,
                             "__pregel_ensure_latest": True,
                             "__pregel_dedupe_tasks": True,
                             "__pregel_store": None,
@@ -635,6 +640,7 @@ def test_subgraph_w_interrupt(
                             "__pregel_delegate": False,
                             "__pregel_read": None,
                             "__pregel_send": None,
+                            "__pregel_call": None,
                             "__pregel_ensure_latest": True,
                             "__pregel_dedupe_tasks": True,
                             "__pregel_resuming": True,

libs/sdk-js/package.json

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

poetry.lock

@@ -2933,13 +2933,13 @@ tenacity = ">=8.1.0,<8.4.0 || >8.4.0,<9.0.0"
 
 [[package]]
 name = "langchain-core"
-version = "0.3.21"
+version = "0.3.23"
 description = "Building applications with LLMs through composability"
 optional = false
 python-versions = "<4.0,>=3.9"
 files = [
-    {file = "langchain_core-0.3.21-py3-none-any.whl", hash = "sha256:7e723dff80946a1198976c6876fea8326dc82566ef9bcb5f8d9188f738733665"},
-    {file = "langchain_core-0.3.21.tar.gz", hash = "sha256:561b52b258ffa50a9fb11d7a1940ebfd915654d1ec95b35e81dfd5ee84143411"},
+    {file = "langchain_core-0.3.23-py3-none-any.whl", hash = "sha256:550c0b996990830fa6515a71a1192a8a0343367999afc36d4ede14222941e420"},
+    {file = "langchain_core-0.3.23.tar.gz", hash = "sha256:f9e175e3b82063cc3b160c2ca2b155832e1c6f915312e1204828f97d4aabf6e1"},
 ]
 
 [package.dependencies]
@@ -3035,7 +3035,7 @@ langchain-core = ">=0.3.0,<0.4.0"
 
 [[package]]
 name = "langgraph"
-version = "0.2.54"
+version = "0.2.57"
 description = "Building stateful, multi-actor applications with LLMs"
 optional = false
 python-versions = ">=3.9.0,<4.0"
@@ -3043,7 +3043,7 @@ files = []
 develop = true
 
 [package.dependencies]
-langchain-core = ">=0.2.43,<0.4.0,!=0.3.0,!=0.3.1,!=0.3.2,!=0.3.3,!=0.3.4,!=0.3.5,!=0.3.6,!=0.3.7,!=0.3.8,!=0.3.9,!=0.3.10,!=0.3.11,!=0.3.12,!=0.3.13,!=0.3.14"
+langchain-core = ">=0.2.43,<0.4.0,!=0.3.0,!=0.3.1,!=0.3.2,!=0.3.3,!=0.3.4,!=0.3.5,!=0.3.6,!=0.3.7,!=0.3.8,!=0.3.9,!=0.3.10,!=0.3.11,!=0.3.12,!=0.3.13,!=0.3.14,!=0.3.15,!=0.3.16,!=0.3.17,!=0.3.18,!=0.3.19,!=0.3.20,!=0.3.21,!=0.3.22"
 langgraph-checkpoint = "^2.0.4"
 langgraph-sdk = "^0.1.42"
 
@@ -3070,7 +3070,7 @@ url = "libs/checkpoint"
 
 [[package]]
 name = "langgraph-checkpoint-postgres"
-version = "2.0.7"
+version = "2.0.8"
 description = "Library with a Postgres implementation of LangGraph checkpoint saver."
 optional = false
 python-versions = "^3.9.0,<4.0"
@@ -3106,7 +3106,7 @@ url = "libs/checkpoint-sqlite"
 
 [[package]]
 name = "langgraph-sdk"
-version = "0.1.42"
+version = "0.1.43"
 description = "SDK for interacting with LangGraph API"
 optional = false
 python-versions = "^3.9.0,<4.0"
@@ -3585,6 +3585,7 @@ optional = false
 python-versions = ">=3.6"
 files = [
     {file = "mkdocs-redirects-1.2.1.tar.gz", hash = "sha256:9420066d70e2a6bb357adf86e67023dcdca1857f97f07c7fe450f8f1fb42f861"},
+    {file = "mkdocs_redirects-1.2.1-py3-none-any.whl", hash = "sha256:497089f9e0219e7389304cffefccdfa1cac5ff9509f2cb706f4c9b221726dffb"},
 ]
 
 [package.dependencies]
@@ -5096,7 +5097,6 @@ description = "Pure-Python implementation of ASN.1 types and DER/BER/CER codecs
 optional = false
 python-versions = ">=3.8"
 files = [
-    {file = "pyasn1-0.6.1-py3-none-any.whl", hash = "sha256:0d632f46f2ba09143da3a8afe9e33fb6f92fa2320ab7e886e2d0f7672af84629"},
     {file = "pyasn1-0.6.1.tar.gz", hash = "sha256:6f580d2bdd84365380830acf45550f2511469f673cb4a5ae3857a3170128b034"},
 ]
 
@@ -5107,7 +5107,6 @@ description = "A collection of ASN.1-based protocols modules"
 optional = false
 python-versions = ">=3.8"
 files = [
-    {file = "pyasn1_modules-0.4.1-py3-none-any.whl", hash = "sha256:49bfa96b45a292b711e986f222502c1c9a5e1f4e568fc30e2574a6c7d07838fd"},
     {file = "pyasn1_modules-0.4.1.tar.gz", hash = "sha256:c28e2dbf9c06ad61c71a075c7e0f9fd0f1b0bb2d2ad4377f240d33ac2ab60a7c"},
 ]
 
@@ -6167,11 +6166,6 @@ files = [
     {file = "scikit_learn-1.5.2-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:f60021ec1574e56632be2a36b946f8143bf4e5e6af4a06d85281adc22938e0dd"},
     {file = "scikit_learn-1.5.2-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:394397841449853c2290a32050382edaec3da89e35b3e03d6cc966aebc6a8ae6"},
     {file = "scikit_learn-1.5.2-cp312-cp312-win_amd64.whl", hash = "sha256:57cc1786cfd6bd118220a92ede80270132aa353647684efa385a74244a41e3b1"},
-    {file = "scikit_learn-1.5.2-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:e9a702e2de732bbb20d3bad29ebd77fc05a6b427dc49964300340e4c9328b3f5"},
-    {file = "scikit_learn-1.5.2-cp313-cp313-macosx_12_0_arm64.whl", hash = "sha256:b0768ad641981f5d3a198430a1d31c3e044ed2e8a6f22166b4d546a5116d7908"},
-    {file = "scikit_learn-1.5.2-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:178ddd0a5cb0044464fc1bfc4cca5b1833bfc7bb022d70b05db8530da4bb3dd3"},
-    {file = "scikit_learn-1.5.2-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:f7284ade780084d94505632241bf78c44ab3b6f1e8ccab3d2af58e0e950f9c12"},
-    {file = "scikit_learn-1.5.2-cp313-cp313-win_amd64.whl", hash = "sha256:b7b0f9a0b1040830d38c39b91b3a44e1b643f4b36e36567b80b7c6bd2202a27f"},
     {file = "scikit_learn-1.5.2-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:757c7d514ddb00ae249832fe87100d9c73c6ea91423802872d9e74970a0e40b9"},
     {file = "scikit_learn-1.5.2-cp39-cp39-macosx_12_0_arm64.whl", hash = "sha256:52788f48b5d8bca5c0736c175fa6bdaab2ef00a8f536cda698db61bd89c551c1"},
     {file = "scikit_learn-1.5.2-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:643964678f4b5fbdc95cbf8aec638acc7aa70f5f79ee2cdad1eec3df4ba6ead8"},
@@ -6962,6 +6956,7 @@ description = "Automatically mock your HTTP interactions to simplify and speed u
 optional = false
 python-versions = ">=3.8"
 files = [
+    {file = "vcrpy-6.0.1-py2.py3-none-any.whl", hash = "sha256:621c3fb2d6bd8aa9f87532c688e4575bcbbde0c0afeb5ebdb7e14cac409edfdd"},
     {file = "vcrpy-6.0.1.tar.gz", hash = "sha256:9e023fee7f892baa0bbda2f7da7c8ac51165c1c6e38ff8688683a12a4bde9278"},
 ]