lib0.2.12

* Remove current_tasks from checkpoint interface

- Not used, now clear that it can be supported with put_writes(SCHEDULE)

* Add comment

examples/memory/manage-conversation-history.ipynb

@@ -268,8 +268,8 @@
     "\n",
     "\n",
     "def filter_messages(messages: list):\n",
-    "    # This is very simple helper function which only ever uses the last two messages\n",
-    "    return messages[-2:]\n",
+    "    # This is very simple helper function which only ever uses the last message\n",
+    "    return messages[-1:]\n",
     "\n",
     "\n",
     "# Define the function that calls the model\n",
@@ -372,9 +372,9 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 3 (ipykernel)",
+   "display_name": "langgraph-example-dev",
    "language": "python",
-   "name": "python3"
+   "name": "langgraph-example-dev"
   },
   "language_info": {
    "codemirror_mode": {
@@ -386,7 +386,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.11.1"
+   "version": "3.11.9"
   }
  },
  "nbformat": 4,

examples/rag/langgraph_agentic_rag.ipynb

@@ -20,7 +20,10 @@
    "id": "969fb438",
    "metadata": {},
    "outputs": [],
-   "source": ["%%capture --no-stderr\n%pip install -U --quiet langchain-community tiktoken langchain-openai langchainhub chromadb langchain langgraph langchain-text-splitters"]
+   "source": [
+    "%%capture --no-stderr\n",
+    "%pip install -U --quiet langchain-community tiktoken langchain-openai langchainhub chromadb langchain langgraph langchain-text-splitters"
+   ]
   },
   {
    "cell_type": "code",
@@ -28,7 +31,22 @@
    "id": "e4958a8c",
    "metadata": {},
    "outputs": [],
-   "source": ["import getpass\nimport os\n\n\ndef _set_env(key: str):\n    if key not in os.environ:\n        os.environ[key] = getpass.getpass(f\"{key}:\")\n\n\n_set_env(\"OPENAI_API_KEY\")\n\n# (Optional) For tracing\nos.environ[\"LANGCHAIN_TRACING_V2\"] = \"true\"\n_set_env(\"LANGCHAIN_API_KEY\")"]
+   "source": [
+    "import getpass\n",
+    "import os\n",
+    "\n",
+    "\n",
+    "def _set_env(key: str):\n",
+    "    if key not in os.environ:\n",
+    "        os.environ[key] = getpass.getpass(f\"{key}:\")\n",
+    "\n",
+    "\n",
+    "_set_env(\"OPENAI_API_KEY\")\n",
+    "\n",
+    "# (Optional) For tracing\n",
+    "os.environ[\"LANGCHAIN_TRACING_V2\"] = \"true\"\n",
+    "_set_env(\"LANGCHAIN_API_KEY\")"
+   ]
   },
   {
    "cell_type": "markdown",
@@ -46,7 +64,34 @@
    "id": "e50c9efe-4abe-42fa-b35a-05eeeede9ec6",
    "metadata": {},
    "outputs": [],
-   "source": ["from langchain_community.document_loaders import WebBaseLoader\nfrom langchain_community.vectorstores import Chroma\nfrom langchain_openai import OpenAIEmbeddings\nfrom langchain_text_splitters import RecursiveCharacterTextSplitter\n\nurls = [\n    \"https://lilianweng.github.io/posts/2023-06-23-agent/\",\n    \"https://lilianweng.github.io/posts/2023-03-15-prompt-engineering/\",\n    \"https://lilianweng.github.io/posts/2023-10-25-adv-attack-llm/\",\n]\n\ndocs = [WebBaseLoader(url).load() for url in urls]\ndocs_list = [item for sublist in docs for item in sublist]\n\ntext_splitter = RecursiveCharacterTextSplitter.from_tiktoken_encoder(\n    chunk_size=100, chunk_overlap=50\n)\ndoc_splits = text_splitter.split_documents(docs_list)\n\n# Add to vectorDB\nvectorstore = Chroma.from_documents(\n    documents=doc_splits,\n    collection_name=\"rag-chroma\",\n    embedding=OpenAIEmbeddings(),\n)\nretriever = vectorstore.as_retriever()"]
+   "source": [
+    "from langchain_community.document_loaders import WebBaseLoader\n",
+    "from langchain_community.vectorstores import Chroma\n",
+    "from langchain_openai import OpenAIEmbeddings\n",
+    "from langchain_text_splitters import RecursiveCharacterTextSplitter\n",
+    "\n",
+    "urls = [\n",
+    "    \"https://lilianweng.github.io/posts/2023-06-23-agent/\",\n",
+    "    \"https://lilianweng.github.io/posts/2023-03-15-prompt-engineering/\",\n",
+    "    \"https://lilianweng.github.io/posts/2023-10-25-adv-attack-llm/\",\n",
+    "]\n",
+    "\n",
+    "docs = [WebBaseLoader(url).load() for url in urls]\n",
+    "docs_list = [item for sublist in docs for item in sublist]\n",
+    "\n",
+    "text_splitter = RecursiveCharacterTextSplitter.from_tiktoken_encoder(\n",
+    "    chunk_size=100, chunk_overlap=50\n",
+    ")\n",
+    "doc_splits = text_splitter.split_documents(docs_list)\n",
+    "\n",
+    "# Add to vectorDB\n",
+    "vectorstore = Chroma.from_documents(\n",
+    "    documents=doc_splits,\n",
+    "    collection_name=\"rag-chroma\",\n",
+    "    embedding=OpenAIEmbeddings(),\n",
+    ")\n",
+    "retriever = vectorstore.as_retriever()"
+   ]
   },
   {
    "cell_type": "markdown",
@@ -62,7 +107,17 @@
    "id": "0b97bdd8-d7e3-444d-ac96-5ef4725f9048",
    "metadata": {},
    "outputs": [],
-   "source": ["from langchain.tools.retriever import create_retriever_tool\n\nretriever_tool = create_retriever_tool(\n    retriever,\n    \"retrieve_blog_posts\",\n    \"Search and return information about Lilian Weng blog posts on LLM agents, prompt engineering, and adversarial attacks on LLMs.\",\n)\n\ntools = [retriever_tool]"]
+   "source": [
+    "from langchain.tools.retriever import create_retriever_tool\n",
+    "\n",
+    "retriever_tool = create_retriever_tool(\n",
+    "    retriever,\n",
+    "    \"retrieve_blog_posts\",\n",
+    "    \"Search and return information about Lilian Weng blog posts on LLM agents, prompt engineering, and adversarial attacks on LLMs.\",\n",
+    ")\n",
+    "\n",
+    "tools = [retriever_tool]"
+   ]
   },
   {
    "cell_type": "markdown",
@@ -86,7 +141,19 @@
    "id": "0e378706-47d5-425a-8ba0-57b9acffbd0c",
    "metadata": {},
    "outputs": [],
-   "source": ["from typing import Annotated, Sequence, TypedDict\n\nfrom langchain_core.messages import BaseMessage\n\nfrom langgraph.graph.message import add_messages\n\n\nclass AgentState(TypedDict):\n    # The add_messages function defines how an update should be processed\n    # Default is to replace. add_messages says \"append\"\n    messages: Annotated[Sequence[BaseMessage], add_messages]"]
+   "source": [
+    "from typing import Annotated, Sequence, TypedDict\n",
+    "\n",
+    "from langchain_core.messages import BaseMessage\n",
+    "\n",
+    "from langgraph.graph.message import add_messages\n",
+    "\n",
+    "\n",
+    "class AgentState(TypedDict):\n",
+    "    # The add_messages function defines how an update should be processed\n",
+    "    # Default is to replace. add_messages says \"append\"\n",
+    "    messages: Annotated[Sequence[BaseMessage], add_messages]"
+   ]
   },
   {
    "attachments": {
@@ -129,7 +196,173 @@
      ]
     }
    ],
-   "source": ["from typing import Annotated, Literal, Sequence, TypedDict\n\nfrom langchain import hub\nfrom langchain_core.messages import BaseMessage, HumanMessage\nfrom langchain_core.output_parsers import StrOutputParser\nfrom langchain_core.prompts import PromptTemplate\nfrom langchain_core.pydantic_v1 import BaseModel, Field\nfrom langchain_openai import ChatOpenAI\n\nfrom langgraph.prebuilt import tools_condition\n\n### Edges\n\n\ndef grade_documents(state) -> Literal[\"generate\", \"rewrite\"]:\n    \"\"\"\n    Determines whether the retrieved documents are relevant to the question.\n\n    Args:\n        state (messages): The current state\n\n    Returns:\n        str: A decision for whether the documents are relevant or not\n    \"\"\"\n\n    print(\"---CHECK RELEVANCE---\")\n\n    # Data model\n    class grade(BaseModel):\n        \"\"\"Binary score for relevance check.\"\"\"\n\n        binary_score: str = Field(description=\"Relevance score 'yes' or 'no'\")\n\n    # LLM\n    model = ChatOpenAI(temperature=0, model=\"gpt-4-0125-preview\", streaming=True)\n\n    # LLM with tool and validation\n    llm_with_tool = model.with_structured_output(grade)\n\n    # Prompt\n    prompt = PromptTemplate(\n        template=\"\"\"You are a grader assessing relevance of a retrieved document to a user question. \\n \n        Here is the retrieved document: \\n\\n {context} \\n\\n\n        Here is the user question: {question} \\n\n        If the document contains keyword(s) or semantic meaning related to the user question, grade it as relevant. \\n\n        Give a binary score 'yes' or 'no' score to indicate whether the document is relevant to the question.\"\"\",\n        input_variables=[\"context\", \"question\"],\n    )\n\n    # Chain\n    chain = prompt | llm_with_tool\n\n    messages = state[\"messages\"]\n    last_message = messages[-1]\n\n    question = messages[0].content\n    docs = last_message.content\n\n    scored_result = chain.invoke({\"question\": question, \"context\": docs})\n\n    score = scored_result.binary_score\n\n    if score == \"yes\":\n        print(\"---DECISION: DOCS RELEVANT---\")\n        return \"generate\"\n\n    else:\n        print(\"---DECISION: DOCS NOT RELEVANT---\")\n        print(score)\n        return \"rewrite\"\n\n\n### Nodes\n\n\ndef agent(state):\n    \"\"\"\n    Invokes the agent model to generate a response based on the current state. Given\n    the question, it will decide to retrieve using the retriever tool, or simply end.\n\n    Args:\n        state (messages): The current state\n\n    Returns:\n        dict: The updated state with the agent response appended to messages\n    \"\"\"\n    print(\"---CALL AGENT---\")\n    messages = state[\"messages\"]\n    model = ChatOpenAI(temperature=0, streaming=True, model=\"gpt-4-turbo\")\n    model = model.bind_tools(tools)\n    response = model.invoke(messages)\n    # We return a list, because this will get added to the existing list\n    return {\"messages\": [response]}\n\n\ndef rewrite(state):\n    \"\"\"\n    Transform the query to produce a better question.\n\n    Args:\n        state (messages): The current state\n\n    Returns:\n        dict: The updated state with re-phrased question\n    \"\"\"\n\n    print(\"---TRANSFORM QUERY---\")\n    messages = state[\"messages\"]\n    question = messages[0].content\n\n    msg = [\n        HumanMessage(\n            content=f\"\"\" \\n \n    Look at the input and try to reason about the underlying semantic intent / meaning. \\n \n    Here is the initial question:\n    \\n ------- \\n\n    {question} \n    \\n ------- \\n\n    Formulate an improved question: \"\"\",\n        )\n    ]\n\n    # Grader\n    model = ChatOpenAI(temperature=0, model=\"gpt-4-0125-preview\", streaming=True)\n    response = model.invoke(msg)\n    return {\"messages\": [response]}\n\n\ndef generate(state):\n    \"\"\"\n    Generate answer\n\n    Args:\n        state (messages): The current state\n\n    Returns:\n         dict: The updated state with re-phrased question\n    \"\"\"\n    print(\"---GENERATE---\")\n    messages = state[\"messages\"]\n    question = messages[0].content\n    last_message = messages[-1]\n\n    question = messages[0].content\n    docs = last_message.content\n\n    # Prompt\n    prompt = hub.pull(\"rlm/rag-prompt\")\n\n    # LLM\n    llm = ChatOpenAI(model_name=\"gpt-3.5-turbo\", temperature=0, streaming=True)\n\n    # Post-processing\n    def format_docs(docs):\n        return \"\\n\\n\".join(doc.page_content for doc in docs)\n\n    # Chain\n    rag_chain = prompt | llm | StrOutputParser()\n\n    # Run\n    response = rag_chain.invoke({\"context\": docs, \"question\": question})\n    return {\"messages\": [response]}\n\n\nprint(\"*\" * 20 + \"Prompt[rlm/rag-prompt]\" + \"*\" * 20)\nprompt = hub.pull(\"rlm/rag-prompt\").pretty_print()  # Show what the prompt looks like"]
+   "source": [
+    "from typing import Annotated, Literal, Sequence, TypedDict\n",
+    "\n",
+    "from langchain import hub\n",
+    "from langchain_core.messages import BaseMessage, HumanMessage\n",
+    "from langchain_core.output_parsers import StrOutputParser\n",
+    "from langchain_core.prompts import PromptTemplate\n",
+    "from langchain_core.pydantic_v1 import BaseModel, Field\n",
+    "from langchain_openai import ChatOpenAI\n",
+    "\n",
+    "from langgraph.prebuilt import tools_condition\n",
+    "\n",
+    "### Edges\n",
+    "\n",
+    "\n",
+    "def grade_documents(state) -> Literal[\"generate\", \"rewrite\"]:\n",
+    "    \"\"\"\n",
+    "    Determines whether the retrieved documents are relevant to the question.\n",
+    "\n",
+    "    Args:\n",
+    "        state (messages): The current state\n",
+    "\n",
+    "    Returns:\n",
+    "        str: A decision for whether the documents are relevant or not\n",
+    "    \"\"\"\n",
+    "\n",
+    "    print(\"---CHECK RELEVANCE---\")\n",
+    "\n",
+    "    # Data model\n",
+    "    class grade(BaseModel):\n",
+    "        \"\"\"Binary score for relevance check.\"\"\"\n",
+    "\n",
+    "        binary_score: str = Field(description=\"Relevance score 'yes' or 'no'\")\n",
+    "\n",
+    "    # LLM\n",
+    "    model = ChatOpenAI(temperature=0, model=\"gpt-4-0125-preview\", streaming=True)\n",
+    "\n",
+    "    # LLM with tool and validation\n",
+    "    llm_with_tool = model.with_structured_output(grade)\n",
+    "\n",
+    "    # Prompt\n",
+    "    prompt = PromptTemplate(\n",
+    "        template=\"\"\"You are a grader assessing relevance of a retrieved document to a user question. \\n \n",
+    "        Here is the retrieved document: \\n\\n {context} \\n\\n\n",
+    "        Here is the user question: {question} \\n\n",
+    "        If the document contains keyword(s) or semantic meaning related to the user question, grade it as relevant. \\n\n",
+    "        Give a binary score 'yes' or 'no' score to indicate whether the document is relevant to the question.\"\"\",\n",
+    "        input_variables=[\"context\", \"question\"],\n",
+    "    )\n",
+    "\n",
+    "    # Chain\n",
+    "    chain = prompt | llm_with_tool\n",
+    "\n",
+    "    messages = state[\"messages\"]\n",
+    "    last_message = messages[-1]\n",
+    "\n",
+    "    question = messages[0].content\n",
+    "    docs = last_message.content\n",
+    "\n",
+    "    scored_result = chain.invoke({\"question\": question, \"context\": docs})\n",
+    "\n",
+    "    score = scored_result.binary_score\n",
+    "\n",
+    "    if score == \"yes\":\n",
+    "        print(\"---DECISION: DOCS RELEVANT---\")\n",
+    "        return \"generate\"\n",
+    "\n",
+    "    else:\n",
+    "        print(\"---DECISION: DOCS NOT RELEVANT---\")\n",
+    "        print(score)\n",
+    "        return \"rewrite\"\n",
+    "\n",
+    "\n",
+    "### Nodes\n",
+    "\n",
+    "\n",
+    "def agent(state):\n",
+    "    \"\"\"\n",
+    "    Invokes the agent model to generate a response based on the current state. Given\n",
+    "    the question, it will decide to retrieve using the retriever tool, or simply end.\n",
+    "\n",
+    "    Args:\n",
+    "        state (messages): The current state\n",
+    "\n",
+    "    Returns:\n",
+    "        dict: The updated state with the agent response appended to messages\n",
+    "    \"\"\"\n",
+    "    print(\"---CALL AGENT---\")\n",
+    "    messages = state[\"messages\"]\n",
+    "    model = ChatOpenAI(temperature=0, streaming=True, model=\"gpt-4-turbo\")\n",
+    "    model = model.bind_tools(tools)\n",
+    "    response = model.invoke(messages)\n",
+    "    # We return a list, because this will get added to the existing list\n",
+    "    return {\"messages\": [response]}\n",
+    "\n",
+    "\n",
+    "def rewrite(state):\n",
+    "    \"\"\"\n",
+    "    Transform the query to produce a better question.\n",
+    "\n",
+    "    Args:\n",
+    "        state (messages): The current state\n",
+    "\n",
+    "    Returns:\n",
+    "        dict: The updated state with re-phrased question\n",
+    "    \"\"\"\n",
+    "\n",
+    "    print(\"---TRANSFORM QUERY---\")\n",
+    "    messages = state[\"messages\"]\n",
+    "    question = messages[0].content\n",
+    "\n",
+    "    msg = [\n",
+    "        HumanMessage(\n",
+    "            content=f\"\"\" \\n \n",
+    "    Look at the input and try to reason about the underlying semantic intent / meaning. \\n \n",
+    "    Here is the initial question:\n",
+    "    \\n ------- \\n\n",
+    "    {question} \n",
+    "    \\n ------- \\n\n",
+    "    Formulate an improved question: \"\"\",\n",
+    "        )\n",
+    "    ]\n",
+    "\n",
+    "    # Grader\n",
+    "    model = ChatOpenAI(temperature=0, model=\"gpt-4-0125-preview\", streaming=True)\n",
+    "    response = model.invoke(msg)\n",
+    "    return {\"messages\": [response]}\n",
+    "\n",
+    "\n",
+    "def generate(state):\n",
+    "    \"\"\"\n",
+    "    Generate answer\n",
+    "\n",
+    "    Args:\n",
+    "        state (messages): The current state\n",
+    "\n",
+    "    Returns:\n",
+    "         dict: The updated state with re-phrased question\n",
+    "    \"\"\"\n",
+    "    print(\"---GENERATE---\")\n",
+    "    messages = state[\"messages\"]\n",
+    "    question = messages[0].content\n",
+    "    last_message = messages[-1]\n",
+    "\n",
+    "    docs = last_message.content\n",
+    "\n",
+    "    # Prompt\n",
+    "    prompt = hub.pull(\"rlm/rag-prompt\")\n",
+    "\n",
+    "    # LLM\n",
+    "    llm = ChatOpenAI(model_name=\"gpt-3.5-turbo\", temperature=0, streaming=True)\n",
+    "\n",
+    "    # Post-processing\n",
+    "    def format_docs(docs):\n",
+    "        return \"\\n\\n\".join(doc.page_content for doc in docs)\n",
+    "\n",
+    "    # Chain\n",
+    "    rag_chain = prompt | llm | StrOutputParser()\n",
+    "\n",
+    "    # Run\n",
+    "    response = rag_chain.invoke({\"context\": docs, \"question\": question})\n",
+    "    return {\"messages\": [response]}\n",
+    "\n",
+    "\n",
+    "print(\"*\" * 20 + \"Prompt[rlm/rag-prompt]\" + \"*\" * 20)\n",
+    "prompt = hub.pull(\"rlm/rag-prompt\").pretty_print()  # Show what the prompt looks like"
+   ]
   },
   {
    "cell_type": "markdown",
@@ -150,7 +383,48 @@
    "id": "8718a37f-83c2-4f16-9850-e61e0f49c3d4",
    "metadata": {},
    "outputs": [],
-   "source": ["from langgraph.graph import END, StateGraph, START\nfrom langgraph.prebuilt import ToolNode\n\n# Define a new graph\nworkflow = StateGraph(AgentState)\n\n# Define the nodes we will cycle between\nworkflow.add_node(\"agent\", agent)  # agent\nretrieve = ToolNode([retriever_tool])\nworkflow.add_node(\"retrieve\", retrieve)  # retrieval\nworkflow.add_node(\"rewrite\", rewrite)  # Re-writing the question\nworkflow.add_node(\n    \"generate\", generate\n)  # Generating a response after we know the documents are relevant\n# Call agent node to decide to retrieve or not\nworkflow.add_edge(START, \"agent\")\n\n# Decide whether to retrieve\nworkflow.add_conditional_edges(\n    \"agent\",\n    # Assess agent decision\n    tools_condition,\n    {\n        # Translate the condition outputs to nodes in our graph\n        \"tools\": \"retrieve\",\n        END: END,\n    },\n)\n\n# Edges taken after the `action` node is called.\nworkflow.add_conditional_edges(\n    \"retrieve\",\n    # Assess agent decision\n    grade_documents,\n)\nworkflow.add_edge(\"generate\", END)\nworkflow.add_edge(\"rewrite\", \"agent\")\n\n# Compile\ngraph = workflow.compile()"]
+   "source": [
+    "from langgraph.graph import END, StateGraph, START\n",
+    "from langgraph.prebuilt import ToolNode\n",
+    "\n",
+    "# Define a new graph\n",
+    "workflow = StateGraph(AgentState)\n",
+    "\n",
+    "# Define the nodes we will cycle between\n",
+    "workflow.add_node(\"agent\", agent)  # agent\n",
+    "retrieve = ToolNode([retriever_tool])\n",
+    "workflow.add_node(\"retrieve\", retrieve)  # retrieval\n",
+    "workflow.add_node(\"rewrite\", rewrite)  # Re-writing the question\n",
+    "workflow.add_node(\n",
+    "    \"generate\", generate\n",
+    ")  # Generating a response after we know the documents are relevant\n",
+    "# Call agent node to decide to retrieve or not\n",
+    "workflow.add_edge(START, \"agent\")\n",
+    "\n",
+    "# Decide whether to retrieve\n",
+    "workflow.add_conditional_edges(\n",
+    "    \"agent\",\n",
+    "    # Assess agent decision\n",
+    "    tools_condition,\n",
+    "    {\n",
+    "        # Translate the condition outputs to nodes in our graph\n",
+    "        \"tools\": \"retrieve\",\n",
+    "        END: END,\n",
+    "    },\n",
+    ")\n",
+    "\n",
+    "# Edges taken after the `action` node is called.\n",
+    "workflow.add_conditional_edges(\n",
+    "    \"retrieve\",\n",
+    "    # Assess agent decision\n",
+    "    grade_documents,\n",
+    ")\n",
+    "workflow.add_edge(\"generate\", END)\n",
+    "workflow.add_edge(\"rewrite\", \"agent\")\n",
+    "\n",
+    "# Compile\n",
+    "graph = workflow.compile()"
+   ]
   },
   {
    "cell_type": "code",
@@ -169,7 +443,15 @@
      "output_type": "display_data"
     }
    ],
-   "source": ["from IPython.display import Image, display\n\ntry:\n    display(Image(graph.get_graph(xray=True).draw_mermaid_png()))\nexcept Exception:\n    # This requires some extra dependencies and is optional\n    pass"]
+   "source": [
+    "from IPython.display import Image, display\n",
+    "\n",
+    "try:\n",
+    "    display(Image(graph.get_graph(xray=True).draw_mermaid_png()))\n",
+    "except Exception:\n",
+    "    # This requires some extra dependencies and is optional\n",
+    "    pass"
+   ]
   },
   {
    "cell_type": "code",
@@ -203,7 +485,21 @@
      ]
     }
    ],
-   "source": ["import pprint\n\ninputs = {\n    \"messages\": [\n        (\"user\", \"What does Lilian Weng say about the types of agent memory?\"),\n    ]\n}\nfor output in graph.stream(inputs):\n    for key, value in output.items():\n        pprint.pprint(f\"Output from node '{key}':\")\n        pprint.pprint(\"---\")\n        pprint.pprint(value, indent=2, width=80, depth=None)\n    pprint.pprint(\"\\n---\\n\")"]
+   "source": [
+    "import pprint\n",
+    "\n",
+    "inputs = {\n",
+    "    \"messages\": [\n",
+    "        (\"user\", \"What does Lilian Weng say about the types of agent memory?\"),\n",
+    "    ]\n",
+    "}\n",
+    "for output in graph.stream(inputs):\n",
+    "    for key, value in output.items():\n",
+    "        pprint.pprint(f\"Output from node '{key}':\")\n",
+    "        pprint.pprint(\"---\")\n",
+    "        pprint.pprint(value, indent=2, width=80, depth=None)\n",
+    "    pprint.pprint(\"\\n---\\n\")"
+   ]
   },
   {
    "cell_type": "code",
@@ -211,7 +507,7 @@
    "id": "189333cc-5d34-4869-9f9b-741210e1096f",
    "metadata": {},
    "outputs": [],
-   "source": [""]
+   "source": []
   }
  ],
  "metadata": {

libs/checkpoint-postgres/README.md

@@ -44,7 +44,6 @@ with PostgresSaver.from_conn_string(DB_URI) as checkpointer:
             }
         },
         "pending_sends": [],
-        "current_tasks": {}
     }
 
     # store checkpoint
@@ -87,7 +86,6 @@ async with AsyncPostgresSaver.from_conn_string(DB_URI) as checkpointer:
             }
         },
         "pending_sends": [],
-        "current_tasks": {}
     }
 
     # store checkpoint

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

@@ -150,6 +150,7 @@ class PostgresSaver(BasePostgresSaver):
                 }
                 if value["parent_checkpoint_id"]
                 else None,
+                self._load_writes(value["pending_writes"]),
             )
 
     def get_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:
@@ -317,16 +318,6 @@ class PostgresSaver(BasePostgresSaver):
             task_id (str): Identifier for the task creating the writes.
         """
         with self._cursor(pipeline=True) as cur:
-            cur.execute(
-                self.DELETE_WRITES_SQL,
-                (
-                    config["configurable"]["thread_id"],
-                    config["configurable"]["checkpoint_ns"],
-                    config["configurable"]["checkpoint_id"],
-                    task_id,
-                    len(writes),
-                ),
-            )
             cur.executemany(
                 self.UPSERT_CHECKPOINT_WRITES_SQL,
                 self._dump_writes(

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

@@ -135,6 +135,7 @@ class AsyncPostgresSaver(BasePostgresSaver):
                 }
                 if value["parent_checkpoint_id"]
                 else None,
+                await asyncio.to_thread(self._load_writes, value["pending_writes"]),
             )
 
     async def aget_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:
@@ -273,16 +274,6 @@ class AsyncPostgresSaver(BasePostgresSaver):
             task_id (str): Identifier for the task creating the writes.
         """
         async with self._cursor(pipeline=True) as cur:
-            await cur.execute(
-                self.DELETE_WRITES_SQL,
-                (
-                    config["configurable"]["thread_id"],
-                    config["configurable"]["checkpoint_ns"],
-                    config["configurable"]["checkpoint_id"],
-                    task_id,
-                    len(writes),
-                ),
-            )
             await cur.executemany(
                 self.UPSERT_CHECKPOINT_WRITES_SQL,
                 await asyncio.to_thread(

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

@@ -6,6 +6,7 @@ from langchain_core.runnables import RunnableConfig
 from psycopg.types.json import Jsonb
 
 from langgraph.checkpoint.base import (
+    WRITES_IDX_MAP,
     BaseCheckpointSaver,
     Checkpoint,
     EmptyChannelError,
@@ -105,15 +106,6 @@ UPSERT_CHECKPOINT_WRITES_SQL = """
     ON CONFLICT (thread_id, checkpoint_ns, checkpoint_id, task_id, idx) DO NOTHING
 """
 
-DELETE_WRITES_SQL = """
-    DELETE FROM checkpoint_writes
-    WHERE thread_id = %s
-    AND checkpoint_ns = %s
-    AND checkpoint_id = %s
-    AND task_id = %s
-    AND idx >= %s
-"""
-
 
 class BasePostgresSaver(BaseCheckpointSaver):
     SELECT_SQL = SELECT_SQL
@@ -121,7 +113,6 @@ class BasePostgresSaver(BaseCheckpointSaver):
     UPSERT_CHECKPOINT_BLOBS_SQL = UPSERT_CHECKPOINT_BLOBS_SQL
     UPSERT_CHECKPOINTS_SQL = UPSERT_CHECKPOINTS_SQL
     UPSERT_CHECKPOINT_WRITES_SQL = UPSERT_CHECKPOINT_WRITES_SQL
-    DELETE_WRITES_SQL = DELETE_WRITES_SQL
 
     jsonplus_serde = JsonPlusSerializer()
 
@@ -210,7 +201,7 @@ class BasePostgresSaver(BaseCheckpointSaver):
                 checkpoint_ns,
                 checkpoint_id,
                 task_id,
-                idx,
+                WRITES_IDX_MAP.get(channel, idx),
                 channel,
                 *self.serde.dumps_typed(value),
             )

libs/checkpoint-sqlite/README.md

@@ -35,7 +35,6 @@ with SqliteSaver.from_conn_string(":memory:") as checkpointer:
             }
         },
         "pending_sends": [],
-        "current_tasks": {}
     }
 
     # store checkpoint
@@ -78,7 +77,6 @@ async with AsyncSqliteSaver.from_conn_string(":memory:") as checkpointer:
             }
         },
         "pending_sends": [],
-        "current_tasks": {}
     }
 
     # store checkpoint
@@ -89,4 +87,4 @@ async with AsyncSqliteSaver.from_conn_string(":memory:") as checkpointer:
 
     # list checkpoints
     [c async for c in checkpointer.alist(read_config)]
-```
+```

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

@@ -1,12 +1,13 @@
 import sqlite3
 import threading
-from contextlib import contextmanager
+from contextlib import closing, contextmanager
 from hashlib import md5
 from typing import Any, AsyncIterator, Dict, Iterator, Optional, Sequence, Tuple
 
 from langchain_core.runnables import RunnableConfig
 
 from langgraph.checkpoint.base import (
+    WRITES_IDX_MAP,
     BaseCheckpointSaver,
     ChannelVersions,
     Checkpoint,
@@ -318,7 +319,7 @@ class SqliteSaver(BaseCheckpointSaver):
         ORDER BY checkpoint_id DESC"""
         if limit:
             query += f" LIMIT {limit}"
-        with self.cursor(transaction=False) as cur:
+        with self.cursor(transaction=False) as cur, closing(self.conn.cursor()) as wcur:
             cur.execute(query, param_values)
             for (
                 thread_id,
@@ -329,6 +330,10 @@ class SqliteSaver(BaseCheckpointSaver):
                 checkpoint,
                 metadata,
             ) in cur:
+                wcur.execute(
+                    "SELECT task_id, channel, type, value FROM writes WHERE thread_id = ? AND checkpoint_ns = ? AND checkpoint_id = ?",
+                    (thread_id, checkpoint_ns, checkpoint_id),
+                )
                 yield CheckpointTuple(
                     {
                         "configurable": {
@@ -350,6 +355,10 @@ class SqliteSaver(BaseCheckpointSaver):
                         if parent_checkpoint_id
                         else None
                     ),
+                    [
+                        (task_id, channel, self.serde.loads_typed((type, value)))
+                        for task_id, channel, type, value in wcur
+                    ],
                 )
 
     def put(
@@ -424,25 +433,15 @@ class SqliteSaver(BaseCheckpointSaver):
             task_id (str): Identifier for the task creating the writes.
         """
         with self.lock, self.cursor() as cur:
-            cur.execute(
-                "DELETE FROM writes WHERE thread_id = ? AND checkpoint_ns = ? AND checkpoint_id = ? AND task_id = ? AND idx >= ?",
-                (
-                    str(config["configurable"]["thread_id"]),
-                    str(config["configurable"]["checkpoint_ns"]),
-                    str(config["configurable"]["checkpoint_id"]),
-                    task_id,
-                    len(writes),
-                ),
-            )
             cur.executemany(
-                "INSERT OR REPLACE INTO writes (thread_id, checkpoint_ns, checkpoint_id, task_id, idx, channel, type, value) VALUES (?, ?, ?, ?, ?, ?, ?, ?)",
+                "INSERT OR IGNORE INTO writes (thread_id, checkpoint_ns, checkpoint_id, task_id, idx, channel, type, value) VALUES (?, ?, ?, ?, ?, ?, ?, ?)",
                 [
                     (
                         str(config["configurable"]["thread_id"]),
                         str(config["configurable"]["checkpoint_ns"]),
                         str(config["configurable"]["checkpoint_id"]),
                         task_id,
-                        idx,
+                        WRITES_IDX_MAP.get(channel, idx),
                         channel,
                         *self.serde.dumps_typed(value),
                     )

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

@@ -16,6 +16,7 @@ import aiosqlite
 from langchain_core.runnables import RunnableConfig
 
 from langgraph.checkpoint.base import (
+    WRITES_IDX_MAP,
     BaseCheckpointSaver,
     ChannelVersions,
     Checkpoint,
@@ -329,14 +330,14 @@ class AsyncSqliteSaver(BaseCheckpointSaver):
             AsyncIterator[CheckpointTuple]: An asynchronous iterator of matching checkpoint tuples.
         """
         await self.setup()
-        where, param_values = search_where(config, filter, before)
+        where, params = search_where(config, filter, before)
         query = f"""SELECT thread_id, checkpoint_ns, checkpoint_id, parent_checkpoint_id, type, checkpoint, metadata
         FROM checkpoints
         {where}
         ORDER BY checkpoint_id DESC"""
         if limit:
             query += f" LIMIT {limit}"
-        async with self.conn.execute(query, param_values) as cursor:
+        async with self.conn.execute(query, params) as cur, self.conn.cursor() as wcur:
             async for (
                 thread_id,
                 checkpoint_ns,
@@ -345,7 +346,11 @@ class AsyncSqliteSaver(BaseCheckpointSaver):
                 type,
                 checkpoint,
                 metadata,
-            ) in cursor:
+            ) in cur:
+                await wcur.execute(
+                    "SELECT task_id, channel, type, value FROM writes WHERE thread_id = ? AND checkpoint_ns = ? AND checkpoint_id = ?",
+                    (thread_id, checkpoint_ns, checkpoint_id),
+                )
                 yield CheckpointTuple(
                     {
                         "configurable": {
@@ -367,6 +372,10 @@ class AsyncSqliteSaver(BaseCheckpointSaver):
                         if parent_checkpoint_id
                         else None
                     ),
+                    [
+                        (task_id, channel, self.serde.loads_typed((type, value)))
+                        async for task_id, channel, type, value in wcur
+                    ],
                 )
 
     async def aput(
@@ -433,25 +442,15 @@ class AsyncSqliteSaver(BaseCheckpointSaver):
         """
         await self.setup()
         async with self.conn.cursor() as cur:
-            await cur.execute(
-                "DELETE FROM writes WHERE thread_id = ? AND checkpoint_ns = ? AND checkpoint_id = ? AND task_id = ? AND idx >= ?",
-                (
-                    str(config["configurable"]["thread_id"]),
-                    str(config["configurable"]["checkpoint_ns"]),
-                    str(config["configurable"]["checkpoint_id"]),
-                    task_id,
-                    len(writes),
-                ),
-            )
             await cur.executemany(
-                "INSERT OR REPLACE INTO writes (thread_id, checkpoint_ns, checkpoint_id, task_id, idx, channel, type, value) VALUES (?, ?, ?, ?, ?, ?, ?, ?)",
+                "INSERT OR IGNORE INTO writes (thread_id, checkpoint_ns, checkpoint_id, task_id, idx, channel, type, value) VALUES (?, ?, ?, ?, ?, ?, ?, ?)",
                 [
                     (
                         str(config["configurable"]["thread_id"]),
                         str(config["configurable"]["checkpoint_ns"]),
                         str(config["configurable"]["checkpoint_id"]),
                         task_id,
-                        idx,
+                        WRITES_IDX_MAP.get(channel, idx),
                         channel,
                         *self.serde.dumps_typed(value),
                     )

libs/checkpoint/README.md

@@ -74,7 +74,6 @@ checkpoint = {
       }
     },
     "pending_sends": [],
-    "current_tasks": {}
 }
 
 # store checkpoint

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

@@ -22,6 +22,7 @@ from langgraph.checkpoint.base.id import uuid6
 from langgraph.checkpoint.serde.base import SerializerProtocol, maybe_add_typed_methods
 from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
 from langgraph.checkpoint.serde.types import (
+    ERROR,
     ChannelProtocol,
     SendProtocol,
 )
@@ -96,8 +97,6 @@ class Checkpoint(TypedDict):
     pending_sends: List[SendProtocol]
     """List of packets sent to nodes but not yet processed.
     Cleared by the next checkpoint."""
-    current_tasks: Dict[str, TaskInfo]
-    """Map from task ID to task info."""
 
 
 def empty_checkpoint() -> Checkpoint:
@@ -109,7 +108,6 @@ def empty_checkpoint() -> Checkpoint:
         channel_versions={},
         versions_seen={},
         pending_sends=[],
-        current_tasks={},
     )
 
 
@@ -122,7 +120,6 @@ def copy_checkpoint(checkpoint: Checkpoint) -> Checkpoint:
         channel_versions=checkpoint["channel_versions"].copy(),
         versions_seen={k: v.copy() for k, v in checkpoint["versions_seen"].items()},
         pending_sends=checkpoint.get("pending_sends", []).copy(),
-        current_tasks=checkpoint.get("current_tasks", {}).copy(),
     )
 
 
@@ -140,6 +137,8 @@ def create_checkpoint(
     else:
         values: dict[str, Any] = {}
         for k, v in channels.items():
+            if k not in checkpoint["channel_versions"]:
+                continue
             try:
                 values[k] = v.checkpoint()
             except EmptyChannelError:
@@ -152,7 +151,6 @@ def create_checkpoint(
         channel_versions=checkpoint["channel_versions"],
         versions_seen=checkpoint["versions_seen"],
         pending_sends=checkpoint.get("pending_sends", []),
-        current_tasks={},
     )
 
 
@@ -437,3 +435,14 @@ def get_checkpoint_id(config: RunnableConfig) -> Optional[str]:
     return config["configurable"].get(
         "checkpoint_id", config["configurable"].get("thread_ts")
     )
+
+
+"""
+Mapping from error type to error index.
+Regular writes just map to their index in the list of writes being saved.
+Special writes (e.g. errors) map to negative indices, to avoid those writes from
+saving regular writes.
+Each Checkpointer implementation should use this mapping in put_writes.
+"""
+WRITES_IDX_MAP = {ERROR: -1}
+# TODO To store scheduled status of tasks, add a special channel here

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

@@ -8,6 +8,7 @@ from typing import Any, AsyncIterator, Dict, Iterator, List, Optional, Tuple
 from langchain_core.runnables import RunnableConfig
 
 from langgraph.checkpoint.base import (
+    WRITES_IDX_MAP,
     BaseCheckpointSaver,
     ChannelVersions,
     Checkpoint,
@@ -52,6 +53,9 @@ class MemorySaver(
 
     # thread ID ->  checkpoint NS -> checkpoint ID -> checkpoint mapping
     storage: defaultdict[str, dict[str, dict[str, tuple[bytes, bytes, Optional[str]]]]]
+    writes: defaultdict[
+        tuple[str, str, str], dict[tuple[str, int], tuple[str, str, bytes]]
+    ]
 
     def __init__(
         self,
@@ -60,7 +64,7 @@ class MemorySaver(
     ) -> None:
         super().__init__(serde=serde)
         self.storage = defaultdict(lambda: defaultdict(dict))
-        self.writes = defaultdict(list)
+        self.writes = defaultdict(dict)
 
     def __enter__(self) -> "MemorySaver":
         return self
@@ -103,7 +107,7 @@ class MemorySaver(
         if checkpoint_id := get_checkpoint_id(config):
             if saved := self.storage[thread_id][checkpoint_ns].get(checkpoint_id):
                 checkpoint, metadata, parent_checkpoint_id = saved
-                writes = self.writes[(thread_id, checkpoint_ns, checkpoint_id)]
+                writes = self.writes[(thread_id, checkpoint_ns, checkpoint_id)].values()
                 return CheckpointTuple(
                     config=config,
                     checkpoint=self.serde.loads_typed(checkpoint),
@@ -125,7 +129,7 @@ class MemorySaver(
             if checkpoints := self.storage[thread_id][checkpoint_ns]:
                 checkpoint_id = max(checkpoints.keys())
                 checkpoint, metadata, parent_checkpoint_id = checkpoints[checkpoint_id]
-                writes = self.writes[(thread_id, checkpoint_ns, checkpoint_id)]
+                writes = self.writes[(thread_id, checkpoint_ns, checkpoint_id)].values()
                 return CheckpointTuple(
                     config={
                         "configurable": {
@@ -204,6 +208,8 @@ class MemorySaver(
                 elif limit is not None:
                     limit -= 1
 
+                writes = self.writes[(thread_id, checkpoint_ns, checkpoint_id)].values()
+
                 yield CheckpointTuple(
                     config={
                         "configurable": {
@@ -223,6 +229,9 @@ class MemorySaver(
                     }
                     if parent_checkpoint_id
                     else None,
+                    pending_writes=[
+                        (id, c, self.serde.loads_typed(v)) for id, c, v in writes
+                    ],
                 )
 
     def put(
@@ -287,11 +296,10 @@ class MemorySaver(
         thread_id = config["configurable"]["thread_id"]
         checkpoint_ns = config["configurable"]["checkpoint_ns"]
         checkpoint_id = config["configurable"]["checkpoint_id"]
-        key = (thread_id, checkpoint_ns, checkpoint_id)
-        self.writes[key] = [w for w in self.writes[key] if w[0] != task_id]
-        self.writes[key].extend(
-            [(task_id, c, self.serde.dumps_typed(v)) for c, v in writes]
-        )
+        outer_key = (thread_id, checkpoint_ns, checkpoint_id)
+        for idx, (c, v) in enumerate(writes):
+            inner_key = (task_id, WRITES_IDX_MAP.get(c, idx))
+            self.writes[outer_key][inner_key] = (task_id, c, self.serde.dumps_typed(v))
 
     async def aget_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:
         """Asynchronous version of get_tuple.

libs/checkpoint/langgraph/checkpoint/serde/types.py

@@ -12,6 +12,8 @@ from typing import (
 from langchain_core.runnables import RunnableConfig
 from typing_extensions import Self
 
+ERROR = "__error__"
+
 Value = TypeVar("Value")
 Update = TypeVar("Update")
 C = TypeVar("C")

libs/checkpoint/pyproject.toml

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

libs/cli/langgraph_cli/docker.py

@@ -12,6 +12,15 @@ DEFAULT_POSTGRES_URI = (
     "postgres://postgres:postgres@langgraph-postgres:5432/postgres?sslmode=disable"
 )
 
+REDIS = """
+    langgraph-redis:
+        image: redis:6
+        healthcheck:
+            test: redis-cli ping
+            interval: 5s
+            timeout: 1s
+            retries: 5
+"""
 
 DB = """
     langgraph-postgres:
@@ -165,18 +174,22 @@ def compose(
             interval: 5s"""
 
     compose_str = f"""{volumes}services:
+{REDIS}
 {db}
 {debugger_compose(port=debugger_port, base_url=debugger_base_url)}
     langgraph-api:
         ports:
-            - "{port}:8000\""""
+            - "{port}:8000\"
+        depends_on:
+            langgraph-redis:
+                condition: service_healthy"""
     if include_db:
         compose_str += """
-        depends_on:
             langgraph-postgres:
                 condition: service_healthy"""
     compose_str += f"""
         environment:
+            REDIS_URI: redis://langgraph-redis:6379
             POSTGRES_URI: {postgres_uri}
 """
     if capabilities.healthcheck_start_interval:

libs/cli/tests/unit_tests/test_cli.py

@@ -45,6 +45,13 @@ def test_prepare_args_and_stdin():
     langgraph-data:
         driver: local
 services:
+    langgraph-redis:
+        image: redis:6
+        healthcheck:
+            test: redis-cli ping
+            interval: 5s
+            timeout: 1s
+            retries: 5
     langgraph-postgres:
         image: postgres:16
         ports:
@@ -76,9 +83,12 @@ services:
         ports:
             - "8000:8000"
         depends_on:
+            langgraph-redis:
+                condition: service_healthy
             langgraph-postgres:
                 condition: service_healthy
         environment:
+            REDIS_URI: redis://langgraph-redis:6379
             POSTGRES_URI: {DEFAULT_POSTGRES_URI}
         healthcheck:
             test: python /api/healthcheck.py

libs/cli/tests/unit_tests/test_docker.py

@@ -20,10 +20,21 @@ def test_compose_with_no_debugger_and_custom_db():
         DEFAULT_DOCKER_CAPABILITIES, port=port, postgres_uri=custom_postgres_uri
     )
     expected_compose_str = f"""services:
+    langgraph-redis:
+        image: redis:6
+        healthcheck:
+            test: redis-cli ping
+            interval: 5s
+            timeout: 1s
+            retries: 5
     langgraph-api:
         ports:
             - "{port}:8000"
+        depends_on:
+            langgraph-redis:
+                condition: service_healthy
         environment:
+            REDIS_URI: redis://langgraph-redis:6379
             POSTGRES_URI: {custom_postgres_uri}"""
     assert clean_empty_lines(actual_compose_str) == expected_compose_str
 
@@ -37,10 +48,21 @@ def test_compose_with_no_debugger_and_custom_db_with_healthcheck():
         postgres_uri=custom_postgres_uri,
     )
     expected_compose_str = f"""services:
+    langgraph-redis:
+        image: redis:6
+        healthcheck:
+            test: redis-cli ping
+            interval: 5s
+            timeout: 1s
+            retries: 5
     langgraph-api:
         ports:
             - "{port}:8000"
+        depends_on:
+            langgraph-redis:
+                condition: service_healthy
         environment:
+            REDIS_URI: redis://langgraph-redis:6379
             POSTGRES_URI: {custom_postgres_uri}
         healthcheck:
             test: python /api/healthcheck.py
@@ -59,10 +81,21 @@ def test_compose_with_debugger_and_custom_db():
         postgres_uri=custom_postgres_uri,
     )
     expected_compose_str = f"""services:
+    langgraph-redis:
+        image: redis:6
+        healthcheck:
+            test: redis-cli ping
+            interval: 5s
+            timeout: 1s
+            retries: 5
     langgraph-api:
         ports:
             - "{port}:8000"
+        depends_on:
+            langgraph-redis:
+                condition: service_healthy
         environment:
+            REDIS_URI: redis://langgraph-redis:6379
             POSTGRES_URI: {custom_postgres_uri}"""
     assert clean_empty_lines(actual_compose_str) == expected_compose_str
 
@@ -74,6 +107,13 @@ def test_compose_with_debugger_and_default_db():
     langgraph-data:
         driver: local
 services:
+    langgraph-redis:
+        image: redis:6
+        healthcheck:
+            test: redis-cli ping
+            interval: 5s
+            timeout: 1s
+            retries: 5
     langgraph-postgres:
         image: postgres:16
         ports:
@@ -94,8 +134,11 @@ services:
         ports:
             - "{port}:8000"
         depends_on:
+            langgraph-redis:
+                condition: service_healthy
             langgraph-postgres:
                 condition: service_healthy
         environment:
+            REDIS_URI: redis://langgraph-redis:6379
             POSTGRES_URI: {DEFAULT_POSTGRES_URI}"""
     assert clean_empty_lines(actual_compose_str) == expected_compose_str

libs/langgraph/langgraph/channels/base.py

@@ -2,9 +2,9 @@ from abc import ABC, abstractmethod
 from contextlib import asynccontextmanager, contextmanager
 from typing import (
     Any,
-    AsyncGenerator,
-    Generator,
+    AsyncIterator,
     Generic,
+    Iterator,
     Optional,
     Sequence,
     TypeVar,
@@ -21,6 +21,8 @@ C = TypeVar("C")
 
 
 class BaseChannel(Generic[Value, Update, C], ABC):
+    key: str = ""
+
     @property
     @abstractmethod
     def ValueType(self) -> Any:
@@ -43,19 +45,35 @@ class BaseChannel(Generic[Value, Update, C], ABC):
     @abstractmethod
     def from_checkpoint(
         self, checkpoint: Optional[C], config: RunnableConfig
-    ) -> Generator[Self, None, None]:
+    ) -> Iterator[Self]:
         """Return a new identical channel, optionally initialized from a checkpoint.
         If the checkpoint contains complex data structures, they should be copied."""
 
+    @contextmanager
+    def from_checkpoint_named(
+        self, checkpoint: Optional[C], config: RunnableConfig
+    ) -> Iterator[Self]:
+        with self.from_checkpoint(checkpoint, config) as value:
+            value.key = self.key
+            yield value
+
     @asynccontextmanager
     async def afrom_checkpoint(
         self, checkpoint: Optional[C], config: RunnableConfig
-    ) -> AsyncGenerator[Self, None]:
+    ) -> AsyncIterator[Self]:
         """Return a new identical channel, optionally initialized from a checkpoint.
         If the checkpoint contains complex data structures, they should be copied."""
         with self.from_checkpoint(checkpoint, config) as value:
             yield value
 
+    @asynccontextmanager
+    async def afrom_checkpoint_named(
+        self, checkpoint: Optional[C], config: RunnableConfig
+    ) -> AsyncIterator[Self]:
+        async with self.afrom_checkpoint(checkpoint, config) as value:
+            value.key = self.key
+            yield value
+
     # state methods
 
     @abstractmethod

libs/langgraph/langgraph/channels/context.py

@@ -112,7 +112,9 @@ class Context(Generic[Value], BaseChannel[Value, None, None]):
 
     def update(self, values: Sequence[None]) -> bool:
         if values:
-            raise InvalidUpdateError("Context channel does not accept writes.")
+            raise InvalidUpdateError(
+                f"At key '{self.key}': Context channel does not accept writes."
+            )
         return False
 
     def get(self) -> Value:

libs/langgraph/langgraph/channels/dynamic_barrier_value.py

@@ -69,7 +69,7 @@ class DynamicBarrierValue(
         if wait_for_names := [v for v in values if isinstance(v, WaitForNames)]:
             if len(wait_for_names) > 1:
                 raise InvalidUpdateError(
-                    "Received multiple WaitForNames updates in the same step."
+                    f"At key '{self.key}': Received multiple WaitForNames updates in the same step."
                 )
             self.names = wait_for_names[0].names
             return True

libs/langgraph/langgraph/channels/ephemeral_value.py

@@ -58,7 +58,7 @@ class EphemeralValue(Generic[Value], BaseChannel[Value, Value, Value]):
                 return False
         if len(values) != 1 and self.guard:
             raise InvalidUpdateError(
-                "EphemeralValue can only receive one value per step."
+                f"At key '{self.key}': EphemeralValue(guard=True) can receive only one value per step. Use guard=False if you want to store any one of multiple values."
             )
 
         self.value = values[-1]

libs/langgraph/langgraph/channels/last_value.py

@@ -52,7 +52,9 @@ class LastValue(Generic[Value], BaseChannel[Value, Value, Value]):
         if len(values) == 0:
             return False
         if len(values) != 1:
-            raise InvalidUpdateError("LastValue can only receive one value per step.")
+            raise InvalidUpdateError(
+                f"At key '{self.key}': Can receive only one value per step. Use an Annotated key to handle multiple values."
+            )
 
         self.value = values[-1]
         return True

libs/langgraph/langgraph/channels/manager.py

@@ -1,39 +0,0 @@
-from contextlib import AsyncExitStack, ExitStack, asynccontextmanager, contextmanager
-from typing import AsyncGenerator, Generator, Mapping
-
-from langchain_core.runnables import RunnableConfig
-
-from langgraph.channels.base import BaseChannel
-from langgraph.checkpoint.base import Checkpoint
-
-
-@contextmanager
-def ChannelsManager(
-    channels: Mapping[str, BaseChannel],
-    checkpoint: Checkpoint,
-    config: RunnableConfig,
-) -> Generator[Mapping[str, BaseChannel], None, None]:
-    """Manage channels for the lifetime of a Pregel invocation (multiple steps)."""
-    with ExitStack() as stack:
-        yield {
-            k: stack.enter_context(
-                v.from_checkpoint(checkpoint["channel_values"].get(k), config)
-            )
-            for k, v in channels.items()
-        }
-
-
-@asynccontextmanager
-async def AsyncChannelsManager(
-    channels: Mapping[str, BaseChannel],
-    checkpoint: Checkpoint,
-    config: RunnableConfig,
-) -> AsyncGenerator[Mapping[str, BaseChannel], None]:
-    """Manage channels for the lifetime of a Pregel invocation (multiple steps)."""
-    async with AsyncExitStack() as stack:
-        yield {
-            k: await stack.enter_async_context(
-                v.afrom_checkpoint(checkpoint["channel_values"].get(k), config)
-            )
-            for k, v in channels.items()
-        }

libs/langgraph/langgraph/channels/named_barrier_value.py

@@ -53,7 +53,9 @@ class NamedBarrierValue(Generic[Value], BaseChannel[Value, Value, set[Value]]):
                     self.seen.add(value)
                     updated = True
             else:
-                raise InvalidUpdateError(f"Value {value} not in {self.names}")
+                raise InvalidUpdateError(
+                    f"At key '{self.key}': Value {value} not in {self.names}"
+                )
         return updated
 
     def get(self) -> Value:

libs/langgraph/langgraph/channels/untracked_value.py

@@ -49,7 +49,7 @@ class UntrackedValue(Generic[Value], BaseChannel[Value, Value, Value]):
             return False
         if len(values) != 1 and self.guard:
             raise InvalidUpdateError(
-                "UntrackedValue can only receive one value per step."
+                f"At key '{self.key}': UntrackedValue(guard=True) can receive only one value per step. Use guard=False if you want to store any one of multiple values."
             )
 
         self.value = values[-1]

libs/langgraph/langgraph/constants.py

@@ -5,6 +5,7 @@ INPUT = "__input__"
 CONFIG_KEY_SEND = "__pregel_send"
 CONFIG_KEY_READ = "__pregel_read"
 CONFIG_KEY_CHECKPOINTER = "__pregel_checkpointer"
+CONFIG_KEY_STORE = "__pregel_store"
 CONFIG_KEY_RESUMING = "__pregel_resuming"
 CONFIG_KEY_TASK_ID = "__pregel_task_id"
 INTERRUPT = "__interrupt__"
@@ -17,6 +18,7 @@ RESERVED = {
     CONFIG_KEY_SEND,
     CONFIG_KEY_READ,
     CONFIG_KEY_CHECKPOINTER,
+    CONFIG_KEY_STORE,
     CONFIG_KEY_RESUMING,
     CONFIG_KEY_TASK_ID,
     INPUT,
@@ -100,5 +102,5 @@ class Send:
 
 @dataclass
 class Interrupt:
-    when: Literal["before", "during", "after"]
-    value: Any = None
+    value: Any
+    when: Literal["during"] = "during"

libs/langgraph/langgraph/errors.py

@@ -1,4 +1,4 @@
-from typing import Any
+from typing import Any, Sequence
 
 from langgraph.checkpoint.base import EmptyChannelError
 from langgraph.constants import Interrupt
@@ -32,7 +32,7 @@ class InvalidUpdateError(Exception):
 class GraphInterrupt(Exception):
     """Raised when a subgraph is interrupted."""
 
-    def __init__(self, interrupts: list[Interrupt]) -> None:
+    def __init__(self, interrupts: Sequence[Interrupt] = ()) -> None:
         super().__init__(interrupts)
 
 
@@ -40,7 +40,7 @@ class NodeInterrupt(GraphInterrupt):
     """Raised by a node to interrupt execution."""
 
     def __init__(self, value: Any) -> None:
-        super().__init__([Interrupt("during", value)])
+        super().__init__([Interrupt(value)])
 
 
 class EmptyInputError(Exception):

libs/langgraph/langgraph/graph/graph.py

@@ -192,12 +192,14 @@ class Graph:
             raise ValueError("END cannot be a start node")
         if end_key == START:
             raise ValueError("START cannot be an end node")
-        if not self.support_multiple_edges and start_key in set(
+
+        # run this validation only for non-StateGraph graphs
+        if not hasattr(self, "channels") and start_key in set(
             start for start, _ in self.edges
         ):
             raise ValueError(
                 f"Already found path for node '{start_key}'.\n"
-                "For multiple edges, use StateGraph with an annotated state key."
+                "For multiple edges, use StateGraph with an Annotated state key."
             )
 
         self.edges.add((start_key, end_key))

libs/langgraph/langgraph/graph/state.py

@@ -1,3 +1,4 @@
+import inspect
 import logging
 import typing
 import warnings
@@ -40,10 +41,18 @@ from langgraph.graph.graph import (
     Graph,
     Send,
 )
-from langgraph.managed.base import ManagedValue, is_managed_value
+from langgraph.managed.base import (
+    ChannelKeyPlaceholder,
+    ChannelTypePlaceholder,
+    ConfiguredManagedValue,
+    ManagedValueSpec,
+    is_managed_value,
+    is_writable_managed_value,
+)
 from langgraph.pregel.read import ChannelRead, PregelNode
 from langgraph.pregel.types import All, RetryPolicy
 from langgraph.pregel.write import SKIP_WRITE, ChannelWrite, ChannelWriteEntry
+from langgraph.store.base import BaseStore
 from langgraph.utils import RunnableCallable, coerce_to_runnable
 
 logger = logging.getLogger(__name__)
@@ -121,8 +130,8 @@ class StateGraph(Graph):
 
     nodes: dict[str, StateNodeSpec]
     channels: dict[str, BaseChannel]
-    managed: dict[str, Type[ManagedValue]]
-    schemas: dict[Type[Any], dict[str, Union[BaseChannel, Type[ManagedValue]]]]
+    managed: dict[str, ManagedValueSpec]
+    schemas: dict[Type[Any], dict[str, Union[BaseChannel, ManagedValueSpec]]]
 
     def __init__(
         self,
@@ -184,10 +193,6 @@ class StateGraph(Graph):
                         )
                 else:
                     self.managed[key] = managed
-            if any(
-                isinstance(c, BinaryOperatorAggregate) for c in self.channels.values()
-            ):
-                self.support_multiple_edges = True
 
     @overload
     def add_node(
@@ -322,10 +327,13 @@ class StateGraph(Graph):
                 hints := get_type_hints(action.__call__) or get_type_hints(action)
             ):
                 if input is None:
-                    input_hint = hints[list(hints.keys())[0]]
-                    if isinstance(input_hint, type) and get_type_hints(input_hint):
-                        input = input_hint
-        except TypeError:
+                    first_parameter_name = next(
+                        iter(inspect.signature(action).parameters.keys())
+                    )
+                    if input_hint := hints.get(first_parameter_name):
+                        if isinstance(input_hint, type) and get_type_hints(input_hint):
+                            input = input_hint
+        except (TypeError, StopIteration):
             pass
         if input is not None:
             self._add_schema(input)
@@ -374,6 +382,8 @@ class StateGraph(Graph):
     def compile(
         self,
         checkpointer: Optional[BaseCheckpointSaver] = None,
+        *,
+        store: Optional[BaseStore] = None,
         interrupt_before: Optional[Union[All, Sequence[str]]] = None,
         interrupt_after: Optional[Union[All, Sequence[str]]] = None,
         debug: bool = False,
@@ -432,7 +442,11 @@ class StateGraph(Graph):
             builder=self,
             config_type=self.config_schema,
             nodes={},
-            channels={**self.channels, START: EphemeralValue(self.input)},
+            channels={
+                **self.channels,
+                **self.managed,
+                START: EphemeralValue(self.input),
+            },
             input_channels=START,
             stream_mode="updates",
             output_channels=output_channels,
@@ -442,6 +456,7 @@ class StateGraph(Graph):
             interrupt_after_nodes=interrupt_after,
             auto_validate=False,
             debug=debug,
+            store=store,
         )
 
         compiled.attach_node(START, None)
@@ -486,7 +501,7 @@ class CompiledStateGraph(CompiledGraph):
                     **{
                         k: (self.channels[k].UpdateType, None)
                         for k in self.builder.schemas[self.builder.input]
-                        if k in self.channels
+                        if isinstance(self.channels[k], BaseChannel)
                         and not isinstance(self.channels[k], Context)
                     },
                 )
@@ -511,7 +526,11 @@ class CompiledStateGraph(CompiledGraph):
                 if not isinstance(v, Context) and not is_managed_value(v)
             ]
         else:
-            output_keys = list(self.builder.channels)
+            output_keys = list(self.builder.channels) + [
+                k
+                for k, v in self.builder.managed.items()
+                if is_writable_managed_value(v)
+            ]
 
         def _get_state_key(
             input: Union[None, dict, Any], config: RunnableConfig, *, key: str
@@ -557,10 +576,7 @@ class CompiledStateGraph(CompiledGraph):
             )
         else:
             input_schema = node.input if node else self.builder.schema
-            input_values = {
-                k: v if is_managed_value(v) else k
-                for k, v in self.builder.schemas[input_schema].items()
-            }
+            input_values = {k: k for k in self.builder.schemas[input_schema]}
             is_single_input = len(input_values) == 1 and "__root__" in input_values
 
             self.channels[key] = EphemeralValue(Any, guard=False)
@@ -679,12 +695,12 @@ def _coerce_state(schema: Type[Any], input: dict[str, Any]) -> dict[str, Any]:
 
 def _get_channels(
     schema: Type[dict],
-) -> tuple[dict[str, BaseChannel], dict[str, Type[ManagedValue]]]:
+) -> tuple[dict[str, BaseChannel], dict[str, ManagedValueSpec]]:
     if not hasattr(schema, "__annotations__"):
-        return {"__root__": _get_channel(schema, allow_managed=False)}, {}
+        return {"__root__": _get_channel("__root__", schema, allow_managed=False)}, {}
 
     all_keys = {
-        name: _get_channel(typ)
+        name: _get_channel(name, typ)
         for name, typ in get_type_hints(schema, include_extras=True).items()
         if name != "__slots__"
     }
@@ -695,18 +711,23 @@ def _get_channels(
 
 
 def _get_channel(
-    annotation: Any, *, allow_managed: bool = True
-) -> Union[BaseChannel, Type[ManagedValue]]:
-    if manager := _is_field_managed_value(annotation):
+    name: str, annotation: Any, *, allow_managed: bool = True
+) -> Union[BaseChannel, ManagedValueSpec]:
+    if manager := _is_field_managed_value(name, annotation):
         if allow_managed:
             return manager
         else:
             raise ValueError(f"This {annotation} not allowed in this position")
     elif channel := _is_field_channel(annotation):
+        channel.key = name
         return channel
     elif channel := _is_field_binop(annotation):
+        channel.key = name
         return channel
-    return LastValue(annotation)
+
+    fallback = LastValue(annotation)
+    fallback.key = name
+    return fallback
 
 
 def _is_field_channel(typ: Type[Any]) -> Optional[BaseChannel]:
@@ -736,12 +757,18 @@ def _is_field_binop(typ: Type[Any]) -> Optional[BinaryOperatorAggregate]:
     return None
 
 
-def _is_field_managed_value(typ: Type[Any]) -> Optional[Type[ManagedValue]]:
+def _is_field_managed_value(name: str, typ: Type[Any]) -> Optional[ManagedValueSpec]:
     if hasattr(typ, "__metadata__"):
         meta = typ.__metadata__
         if len(meta) >= 1:
             decoration = get_origin(meta[-1]) or meta[-1]
             if is_managed_value(decoration):
+                if isinstance(decoration, ConfiguredManagedValue):
+                    for k, v in decoration.kwargs.items():
+                        if v is ChannelKeyPlaceholder:
+                            decoration.kwargs[k] = name
+                        if v is ChannelTypePlaceholder:
+                            decoration.kwargs[k] = typ.__origin__
                 return decoration
 
     return None

libs/langgraph/langgraph/managed/base.py

@@ -1,13 +1,13 @@
-import asyncio
 from abc import ABC, abstractmethod
-from contextlib import AsyncExitStack, ExitStack, asynccontextmanager, contextmanager
+from contextlib import asynccontextmanager, contextmanager
 from inspect import isclass
 from typing import (
     Any,
-    AsyncGenerator,
-    Generator,
+    AsyncIterator,
     Generic,
+    Iterator,
     NamedTuple,
+    Sequence,
     Type,
     TypeVar,
     Union,
@@ -17,6 +17,7 @@ from langchain_core.runnables import RunnableConfig
 from typing_extensions import Self, TypeGuard
 
 V = TypeVar("V")
+U = TypeVar("U")
 
 
 class ManagedValue(ABC, Generic[V]):
@@ -25,9 +26,7 @@ class ManagedValue(ABC, Generic[V]):
 
     @classmethod
     @contextmanager
-    def enter(
-        cls, config: RunnableConfig, **kwargs: Any
-    ) -> Generator[Self, None, None]:
+    def enter(cls, config: RunnableConfig, **kwargs: Any) -> Iterator[Self]:
         try:
             value = cls(config, **kwargs)
             yield value
@@ -41,9 +40,7 @@ class ManagedValue(ABC, Generic[V]):
 
     @classmethod
     @asynccontextmanager
-    async def aenter(
-        cls, config: RunnableConfig, **kwargs: Any
-    ) -> AsyncGenerator[Self, None]:
+    async def aenter(cls, config: RunnableConfig, **kwargs: Any) -> AsyncIterator[Self]:
         try:
             value = cls(config, **kwargs)
             yield value
@@ -60,6 +57,16 @@ class ManagedValue(ABC, Generic[V]):
         ...
 
 
+class WritableManagedValue(Generic[V, U], ManagedValue[V], ABC):
+    @abstractmethod
+    def update(self, writes: Sequence[U]) -> None:
+        ...
+
+    @abstractmethod
+    async def aupdate(self, writes: Sequence[U]) -> None:
+        ...
+
+
 class ConfiguredManagedValue(NamedTuple):
     cls: Type[ManagedValue]
     kwargs: dict[str, Any]
@@ -76,46 +83,23 @@ def is_managed_value(value: Any) -> TypeGuard[ManagedValueSpec]:
     )
 
 
-@contextmanager
-def ManagedValuesManager(
-    values: dict[str, ManagedValueSpec],
-    config: RunnableConfig,
-) -> Generator[ManagedValueMapping, None, None]:
-    if values:
-        with ExitStack() as stack:
-            yield {
-                key: stack.enter_context(
-                    value.cls.enter(config, **value.kwargs)
-                    if isinstance(value, ConfiguredManagedValue)
-                    else value.enter(config)
-                )
-                for key, value in values.items()
-            }
-    else:
-        yield {}
+def is_readonly_managed_value(value: Any) -> TypeGuard[Type[ManagedValue]]:
+    return (
+        isclass(value)
+        and issubclass(value, ManagedValue)
+        and not issubclass(value, WritableManagedValue)
+    ) or (
+        isinstance(value, ConfiguredManagedValue)
+        and not issubclass(value.cls, WritableManagedValue)
+    )
 
 
-@asynccontextmanager
-async def AsyncManagedValuesManager(
-    values: dict[str, ManagedValueSpec],
-    config: RunnableConfig,
-) -> AsyncGenerator[ManagedValueMapping, None]:
-    if values:
-        async with AsyncExitStack() as stack:
-            # create enter tasks with reference to spec
-            tasks = {
-                asyncio.create_task(
-                    stack.enter_async_context(
-                        value.cls.aenter(config, **value.kwargs)
-                        if isinstance(value, ConfiguredManagedValue)
-                        else value.aenter(config)
-                    )
-                ): key
-                for key, value in values.items()
-            }
-            # wait for all enter tasks
-            done, _ = await asyncio.wait(tasks, return_when=asyncio.ALL_COMPLETED)
-            # build mapping from spec to result
-            yield {tasks[task]: task.result() for task in done}
-    else:
-        yield {}
+def is_writable_managed_value(value: Any) -> TypeGuard[Type[WritableManagedValue]]:
+    return (isclass(value) and issubclass(value, WritableManagedValue)) or (
+        isinstance(value, ConfiguredManagedValue)
+        and issubclass(value.cls, WritableManagedValue)
+    )
+
+
+ChannelKeyPlaceholder = object()
+ChannelTypePlaceholder = object()

libs/langgraph/langgraph/managed/shared_value.py

@@ -0,0 +1,126 @@
+import collections.abc
+from contextlib import asynccontextmanager, contextmanager
+from typing import (
+    Any,
+    AsyncIterator,
+    Iterator,
+    Optional,
+    Sequence,
+    Type,
+)
+
+from langchain_core.runnables import RunnableConfig
+from typing_extensions import NotRequired, Required, Self
+
+from langgraph.constants import CONFIG_KEY_STORE
+from langgraph.errors import InvalidUpdateError
+from langgraph.managed.base import (
+    ChannelKeyPlaceholder,
+    ChannelTypePlaceholder,
+    ConfiguredManagedValue,
+    WritableManagedValue,
+)
+from langgraph.store.base import BaseStore
+
+V = dict[str, Any]
+
+
+Value = dict[str, V]
+Update = dict[str, Optional[V]]
+
+
+# Adapted from typing_extensions
+def _strip_extras(t):
+    """Strips Annotated, Required and NotRequired from a given type."""
+    if hasattr(t, "__origin__"):
+        return _strip_extras(t.__origin__)
+    if hasattr(t, "__origin__") and t.__origin__ in (Required, NotRequired):
+        return _strip_extras(t.__args__[0])
+
+    return t
+
+
+class SharedValue(WritableManagedValue[Value, Update]):
+    @staticmethod
+    def on(scope: str) -> ConfiguredManagedValue:
+        return ConfiguredManagedValue(
+            SharedValue,
+            {
+                "scope": scope,
+                "key": ChannelKeyPlaceholder,
+                "typ": ChannelTypePlaceholder,
+            },
+        )
+
+    @classmethod
+    @contextmanager
+    def enter(cls, config: RunnableConfig, **kwargs: Any) -> Iterator[Self]:
+        with super().enter(config, **kwargs) as value:
+            if value.store is not None:
+                saved = value.store.list([value.ns])
+                value.value = saved[value.ns] or {}
+            yield value
+
+    @classmethod
+    @asynccontextmanager
+    async def aenter(cls, config: RunnableConfig, **kwargs: Any) -> AsyncIterator[Self]:
+        async with super().aenter(config, **kwargs) as value:
+            if value.store is not None:
+                saved = await value.store.alist([value.ns])
+                value.value = saved[value.ns] or {}
+            yield value
+
+    def __init__(
+        self, config: RunnableConfig, *, typ: Type[Any], scope: str, key: str
+    ) -> None:
+        if typ := _strip_extras(typ):
+            if typ not in (
+                dict,
+                collections.abc.Mapping,
+                collections.abc.MutableMapping,
+            ):
+                raise ValueError("SharedValue must be a dict")
+        self.scope = scope
+        self.config = config
+        self.value: Value = {}
+        self.store: BaseStore = config["configurable"].get(CONFIG_KEY_STORE)
+        if self.store is None:
+            self.ns: Optional[str] = None
+        elif scope_value := config["configurable"].get(self.scope):
+            self.ns = f"scoped:{scope}:{key}:{scope_value}"
+        else:
+            raise ValueError(
+                f"Scope {scope} for shared state key not in config.configurable"
+            )
+
+    def __call__(self, step: int) -> Value:
+        return self.value.copy()
+
+    def _process_update(
+        self, values: Sequence[Update]
+    ) -> list[tuple[str, str, Optional[dict[str, Any]]]]:
+        writes = []
+        for vv in values:
+            for k, v in vv.items():
+                if v is None:
+                    if k in self.value:
+                        self.value[k] = None
+                        writes.append((self.ns, k, None))
+                elif not isinstance(v, dict):
+                    raise InvalidUpdateError("Received a non-dict value")
+                else:
+                    self.value[k] = v
+                    writes.append((self.ns, k, v))
+        return writes
+
+    def update(self, values: Sequence[Update]) -> None:
+        if self.store is None:
+            self._process_update(values)
+        else:
+            return self.store.put(self._process_update(values))
+
+    async def aupdate(self, writes: Sequence[Update]) -> None:
+        if self.store is None:
+            self._process_update(writes)
+        else:
+            return await self.store.aput(self._process_update(writes))

libs/langgraph/langgraph/pregel/__init__.py

@@ -52,11 +52,6 @@ from langgraph.channels.base import (
     BaseChannel,
 )
 from langgraph.channels.context import Context
-from langgraph.channels.last_value import LastValue
-from langgraph.channels.manager import (
-    AsyncChannelsManager,
-    ChannelsManager,
-)
 from langgraph.checkpoint.base import (
     BaseCheckpointSaver,
     copy_checkpoint,
@@ -70,33 +65,19 @@ from langgraph.constants import (
     CONFIG_KEY_SEND,
     ERROR,
     INTERRUPT,
-    Interrupt,
 )
 from langgraph.errors import GraphInterrupt, GraphRecursionError, InvalidUpdateError
-from langgraph.managed.base import (
-    AsyncManagedValuesManager,
-    ManagedValuesManager,
-    ManagedValueSpec,
-    is_managed_value,
-)
-from langgraph.pregel.algo import (
-    apply_writes,
-    local_read,
-    prepare_next_tasks,
-    should_interrupt,
-)
+from langgraph.managed.base import ManagedValueSpec
+from langgraph.pregel.algo import apply_writes, local_read, prepare_next_tasks
 from langgraph.pregel.debug import (
-    map_debug_task_results,
     print_step_checkpoint,
     print_step_tasks,
     print_step_writes,
     tasks_w_writes,
 )
-from langgraph.pregel.io import (
-    map_output_updates,
-    read_channels,
-)
+from langgraph.pregel.io import read_channels
 from langgraph.pregel.loop import AsyncPregelLoop, SyncPregelLoop
+from langgraph.pregel.manager import AsyncChannelsManager, ChannelsManager
 from langgraph.pregel.read import PregelNode
 from langgraph.pregel.retry import RetryPolicy, arun_with_retry, run_with_retry
 from langgraph.pregel.types import (
@@ -108,6 +89,7 @@ from langgraph.pregel.types import (
 from langgraph.pregel.utils import get_new_channel_versions
 from langgraph.pregel.validate import validate_graph, validate_keys
 from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry
+from langgraph.store.base import BaseStore
 
 WriteValue = Union[
     Runnable[Input, Output],
@@ -196,7 +178,9 @@ class Pregel(
 ):
     nodes: Mapping[str, PregelNode]
 
-    channels: Mapping[str, BaseChannel] = Field(default_factory=dict)
+    channels: Mapping[str, Union[BaseChannel, ManagedValueSpec]] = Field(
+        default_factory=dict
+    )
 
     auto_validate: bool = True
 
@@ -223,6 +207,9 @@ class Pregel(
     checkpointer: Optional[BaseCheckpointSaver] = None
     """Checkpointer used to save and load graph state. Defaults to None."""
 
+    store: Optional[BaseStore] = None
+    """Memory store to use for SharedValues. Defaults to None."""
+
     retry_policy: Optional[RetryPolicy] = None
     """Retry policy to use when running tasks. Set to None to disable."""
 
@@ -343,19 +330,12 @@ class Pregel(
     @property
     def stream_channels_asis(self) -> Union[str, Sequence[str]]:
         return self.stream_channels or [
-            k for k in self.channels if not isinstance(self.channels[k], Context)
+            k
+            for k in self.channels
+            if isinstance(self.channels[k], BaseChannel)
+            and not isinstance(self.channels[k], Context)
         ]
 
-    @property
-    def managed_values_dict(self) -> dict[str, ManagedValueSpec]:
-        return {
-            k: v
-            for node in self.nodes.values()
-            if isinstance(node.channels, dict)
-            for k, v in node.channels.items()
-            if is_managed_value(v)
-        }
-
     def get_state(self, config: RunnableConfig) -> StateSnapshot:
         """Get the current state of the graph."""
         if not self.checkpointer:
@@ -364,16 +344,10 @@ class Pregel(
         saved = self.checkpointer.get_tuple(config)
         checkpoint = saved.checkpoint if saved else empty_checkpoint()
         config = saved.config if saved else config
-        with ChannelsManager(
-            {
-                k: LastValue(None) if isinstance(c, Context) else c
-                for k, c in self.channels.items()
-            },
-            checkpoint,
-            config,
-        ) as channels, ManagedValuesManager(
-            self.managed_values_dict, ensure_config(config)
-        ) as managed:
+        with ChannelsManager(self.channels, checkpoint, config, skip_context=True) as (
+            channels,
+            managed,
+        ):
             next_tasks = prepare_next_tasks(
                 checkpoint,
                 self.nodes,
@@ -391,7 +365,7 @@ class Pregel(
                 saved.metadata if saved else None,
                 saved.checkpoint["ts"] if saved else None,
                 saved.parent_config if saved else None,
-                tasks_w_writes(next_tasks, saved.pending_writes),
+                tasks_w_writes(next_tasks, saved.pending_writes if saved else None),
             )
 
     async def aget_state(self, config: RunnableConfig) -> StateSnapshot:
@@ -404,15 +378,8 @@ class Pregel(
 
         config = saved.config if saved else config
         async with AsyncChannelsManager(
-            {
-                k: LastValue(None) if isinstance(c, Context) else c
-                for k, c in self.channels.items()
-            },
-            checkpoint,
-            config,
-        ) as channels, AsyncManagedValuesManager(
-            self.managed_values_dict, ensure_config(config)
-        ) as managed:
+            self.channels, checkpoint, config, skip_context=True
+        ) as (channels, managed):
             next_tasks = prepare_next_tasks(
                 checkpoint,
                 self.nodes,
@@ -429,7 +396,7 @@ class Pregel(
                 saved.metadata if saved else None,
                 saved.checkpoint["ts"] if saved else None,
                 saved.parent_config if saved else None,
-                tasks_w_writes(next_tasks, saved.pending_writes),
+                tasks_w_writes(next_tasks, saved.pending_writes if saved else None),
             )
 
     def get_state_history(
@@ -456,15 +423,8 @@ class Pregel(
             pending_writes,
         ) in self.checkpointer.list(config, before=before, limit=limit, filter=filter):
             with ChannelsManager(
-                {
-                    k: LastValue(None) if isinstance(c, Context) else c
-                    for k, c in self.channels.items()
-                },
-                checkpoint,
-                config,
-            ) as channels, ManagedValuesManager(
-                self.managed_values_dict, ensure_config(config)
-            ) as managed:
+                self.channels, checkpoint, config, skip_context=True
+            ) as (channels, managed):
                 next_tasks = prepare_next_tasks(
                     checkpoint,
                     self.nodes,
@@ -508,15 +468,8 @@ class Pregel(
             pending_writes,
         ) in self.checkpointer.alist(config, before=before, limit=limit, filter=filter):
             async with AsyncChannelsManager(
-                {
-                    k: LastValue(None) if isinstance(c, Context) else c
-                    for k, c in self.channels.items()
-                },
-                checkpoint,
-                config,
-            ) as channels, AsyncManagedValuesManager(
-                self.managed_values_dict, ensure_config(config)
-            ) as managed:
+                self.channels, checkpoint, config, skip_context=True
+            ) as (channels, managed):
                 next_tasks = prepare_next_tasks(
                     checkpoint,
                     self.nodes,
@@ -609,11 +562,10 @@ class Pregel(
         if as_node not in self.nodes:
             raise InvalidUpdateError(f"Node {as_node} does not exist")
         # update channels
-        with ChannelsManager(
-            self.channels, checkpoint, config
-        ) as channels, ManagedValuesManager(
-            self.managed_values_dict, ensure_config(config)
-        ) as managed:
+        with ChannelsManager(self.channels, checkpoint, config) as (
+            channels,
+            _,
+        ):
             # create task to run all writers of the chosen node
             writers = self.nodes[as_node].get_writers()
             if not writers:
@@ -644,35 +596,10 @@ class Pregel(
                 ),
             )
             # apply to checkpoint and save
-            apply_writes(
+            assert not apply_writes(
                 checkpoint, channels, [task], self.checkpointer.get_next_version
-            )
+            ), "Can't write to SharedValues from update_state"
             checkpoint = create_checkpoint(checkpoint, channels, step + 1)
-            # check interrupt before
-            if tasks := should_interrupt(
-                checkpoint,
-                self.interrupt_before_nodes,
-                prepare_next_tasks(
-                    checkpoint,
-                    self.nodes,
-                    channels,
-                    managed,
-                    config,
-                    step + 2,
-                    for_execution=False,
-                ),
-            ):
-                for t in tasks:
-                    self.checkpointer.put_writes(
-                        {
-                            "configurable": {
-                                **checkpoint_config["configurable"],
-                                "checkpoint_id": checkpoint["id"],
-                            }
-                        },
-                        [(INTERRUPT, Interrupt("before"))],
-                        t.id,
-                    )
             return self.checkpointer.put(
                 checkpoint_config,
                 checkpoint,
@@ -753,11 +680,10 @@ class Pregel(
         if as_node not in self.nodes:
             raise InvalidUpdateError(f"Node {as_node} does not exist")
         # update channels, acting as the chosen node
-        async with AsyncChannelsManager(
-            self.channels, checkpoint, config
-        ) as channels, AsyncManagedValuesManager(
-            self.managed_values_dict, ensure_config(config)
-        ) as managed:
+        async with AsyncChannelsManager(self.channels, checkpoint, config) as (
+            channels,
+            _,
+        ):
             # create task to run all writers of the chosen node
             writers = self.nodes[as_node].get_writers()
             if not writers:
@@ -788,39 +714,10 @@ class Pregel(
                 ),
             )
             # apply to checkpoint and save
-            apply_writes(
+            assert not apply_writes(
                 checkpoint, channels, [task], self.checkpointer.get_next_version
-            )
+            ), "Can't write to SharedValues from update_state"
             checkpoint = create_checkpoint(checkpoint, channels, step + 1)
-            # check interrupt before
-            if tasks := should_interrupt(
-                checkpoint,
-                self.interrupt_before_nodes,
-                prepare_next_tasks(
-                    checkpoint,
-                    self.nodes,
-                    channels,
-                    managed,
-                    config,
-                    step + 2,
-                    for_execution=False,
-                ),
-            ):
-                await asyncio.gather(
-                    *(
-                        self.checkpointer.aput_writes(
-                            {
-                                "configurable": {
-                                    **checkpoint_config["configurable"],
-                                    "checkpoint_id": checkpoint["id"],
-                                }
-                            },
-                            [(INTERRUPT, Interrupt("before"))],
-                            t.id,
-                        )
-                        for t in tasks
-                    )
-                )
             return await self.checkpointer.aput(
                 checkpoint_config,
                 checkpoint,
@@ -994,7 +891,14 @@ class Pregel(
             )
 
             with SyncPregelLoop(
-                input, config=config, checkpointer=checkpointer, graph=self
+                input,
+                config=config,
+                store=self.store,
+                checkpointer=checkpointer,
+                nodes=self.nodes,
+                specs=self.channels,
+                output_keys=output_keys,
+                stream_keys=self.stream_channels_asis,
             ) as loop:
                 # Similarly to Bulk Synchronous Parallel / Pregel model
                 # computation proceeds in steps, while there are channel updates
@@ -1002,7 +906,7 @@ class Pregel(
                 # channels are guaranteed to be immutable for the duration of the step,
                 # with channel updates applied only at the transition between steps
                 while loop.tick(
-                    output_keys=output_keys,
+                    input_keys=self.input_channels,
                     interrupt_before=interrupt_before,
                     interrupt_after=interrupt_after,
                     manager=run_manager,
@@ -1071,26 +975,17 @@ class Pregel(
                             else:
                                 # save task writes to checkpointer
                                 loop.put_writes(task.id, task.writes)
-                                # yield updates output for the finished task
-                                if "updates" in stream_modes:
-                                    yield from _with_mode(
-                                        "updates",
-                                        isinstance(stream_mode, list),
-                                        map_output_updates(output_keys, [task]),
-                                    )
-                                if "debug" in stream_modes:
-                                    yield from _with_mode(
-                                        "debug",
-                                        isinstance(stream_mode, list),
-                                        map_debug_task_results(
-                                            loop.step,
-                                            [task],
-                                            self.stream_channels_list,
-                                        ),
-                                    )
                         else:
                             # remove references to loop vars
                             del fut, task
+                        # emit output
+                        while loop.stream:
+                            mode, payload = loop.stream.popleft()
+                            if mode in stream_modes:
+                                if isinstance(stream_mode, list):
+                                    yield (mode, payload)
+                                else:
+                                    yield payload
                         if _should_stop_others(done):
                             break
 
@@ -1105,23 +1000,23 @@ class Pregel(
                             [w for t in loop.tasks for w in t.writes],
                             self.stream_channels_list,
                         )
-                # emit output
-                while loop.stream:
-                    mode, payload = loop.stream.popleft()
-                    if mode in stream_modes:
-                        if isinstance(stream_mode, list):
-                            yield (mode, payload)
-                        else:
-                            yield payload
-                # handle exit
-                if loop.status == "out_of_steps":
-                    raise GraphRecursionError(
-                        f"Recursion limit of {config['recursion_limit']} reached "
-                        "without hitting a stop condition. You can increase the "
-                        "limit by setting the `recursion_limit` config key."
-                    )
-                # set final channel values as run output
-                run_manager.on_chain_end(read_channels(loop.channels, output_keys))
+            # emit output
+            while loop.stream:
+                mode, payload = loop.stream.popleft()
+                if mode in stream_modes:
+                    if isinstance(stream_mode, list):
+                        yield (mode, payload)
+                    else:
+                        yield payload
+            # handle exit
+            if loop.status == "out_of_steps":
+                raise GraphRecursionError(
+                    f"Recursion limit of {config['recursion_limit']} reached "
+                    "without hitting a stop condition. You can increase the "
+                    "limit by setting the `recursion_limit` config key."
+                )
+            # set final channel values as run output
+            run_manager.on_chain_end(loop.output)
         except BaseException as e:
             run_manager.on_chain_error(e)
             raise
@@ -1244,7 +1139,14 @@ class Pregel(
                 debug=debug,
             )
             async with AsyncPregelLoop(
-                input, config=config, checkpointer=checkpointer, graph=self
+                input,
+                config=config,
+                store=self.store,
+                checkpointer=checkpointer,
+                nodes=self.nodes,
+                specs=self.channels,
+                output_keys=output_keys,
+                stream_keys=self.stream_channels_asis,
             ) as loop:
                 aioloop = asyncio.get_event_loop()
                 # Similarly to Bulk Synchronous Parallel / Pregel model
@@ -1253,7 +1155,7 @@ class Pregel(
                 # channels are guaranteed to be immutable for the duration of the step,
                 # with channel updates applied only at the transition between steps
                 while loop.tick(
-                    output_keys=output_keys,
+                    input_keys=self.input_channels,
                     interrupt_before=interrupt_before,
                     interrupt_after=interrupt_after,
                     manager=run_manager,
@@ -1323,28 +1225,17 @@ class Pregel(
                             else:
                                 # save task writes to checkpointer
                                 loop.put_writes(task.id, task.writes)
-                                # yield updates output for the finished task
-                                if "updates" in stream_modes:
-                                    for chunk in _with_mode(
-                                        "updates",
-                                        isinstance(stream_mode, list),
-                                        map_output_updates(output_keys, [task]),
-                                    ):
-                                        yield chunk
-                                if "debug" in stream_modes:
-                                    for chunk in _with_mode(
-                                        "debug",
-                                        isinstance(stream_mode, list),
-                                        map_debug_task_results(
-                                            loop.step,
-                                            [task],
-                                            self.stream_channels_list,
-                                        ),
-                                    ):
-                                        yield chunk
                         else:
                             # remove references to loop vars
                             del fut, task
+                        # emit output
+                        while loop.stream:
+                            mode, payload = loop.stream.popleft()
+                            if mode in stream_modes:
+                                if isinstance(stream_mode, list):
+                                    yield (mode, payload)
+                                else:
+                                    yield payload
                         if _should_stop_others(done):
                             break
 
@@ -1359,28 +1250,24 @@ class Pregel(
                             [w for t in loop.tasks for w in t.writes],
                             self.stream_channels_list,
                         )
-                # emit output
-                while loop.stream:
-                    mode, payload = loop.stream.popleft()
-                    if mode in stream_modes:
-                        if isinstance(stream_mode, list):
-                            yield (mode, payload)
-                        else:
-                            yield payload
-                # handle exit
-                if loop.status == "out_of_steps":
-                    raise GraphRecursionError(
-                        f"Recursion limit of {config['recursion_limit']} reached "
-                        "without hitting a stop condition. You can increase the "
-                        "limit by setting the `recursion_limit` config key."
-                    )
-
-                # set final channel values as run output
-                await run_manager.on_chain_end(
-                    read_channels(loop.channels, output_keys)
+            # emit output
+            while loop.stream:
+                mode, payload = loop.stream.popleft()
+                if mode in stream_modes:
+                    if isinstance(stream_mode, list):
+                        yield (mode, payload)
+                    else:
+                        yield payload
+            # handle exit
+            if loop.status == "out_of_steps":
+                raise GraphRecursionError(
+                    f"Recursion limit of {config['recursion_limit']} reached "
+                    "without hitting a stop condition. You can increase the "
+                    "limit by setting the `recursion_limit` config key."
                 )
+            # set final channel values as run output
+            await run_manager.on_chain_end(loop.output)
         except BaseException as e:
-            # TODO use on_chain_end if exc is GraphInterrupt
             await asyncio.shield(run_manager.on_chain_error(e))
             raise
 

libs/langgraph/langgraph/pregel/algo.py

@@ -25,7 +25,6 @@ from langchain_core.runnables.config import (
 
 from langgraph.channels.base import BaseChannel
 from langgraph.channels.context import Context
-from langgraph.channels.manager import ChannelsManager
 from langgraph.checkpoint.base import (
     BaseCheckpointSaver,
     Checkpoint,
@@ -46,9 +45,10 @@ from langgraph.constants import (
     Send,
 )
 from langgraph.errors import EmptyChannelError, InvalidUpdateError
-from langgraph.managed.base import ManagedValueMapping, is_managed_value
+from langgraph.managed.base import ManagedValueMapping
 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.pregel.types import All, PregelExecutableTask, PregelTask
 
@@ -105,11 +105,10 @@ def local_read(
     if fresh:
         new_checkpoint = create_checkpoint(copy_checkpoint(checkpoint), channels, -1)
         context_channels = {k: v for k, v in channels.items() if isinstance(v, Context)}
-        with ChannelsManager(
-            {k: v for k, v in channels.items() if k not in context_channels},
-            new_checkpoint,
-            config,
-        ) as channels:
+        with ChannelsManager(channels, new_checkpoint, config, skip_context=True) as (
+            channels,
+            _,
+        ):
             all_channels = {**channels, **context_channels}
             apply_writes(new_checkpoint, all_channels, [task], None)
             return read_channels(all_channels, select)
@@ -121,6 +120,7 @@ def local_write(
     commit: Callable[[Sequence[tuple[str, Any]]], None],
     processes: Mapping[str, PregelNode],
     channels: Mapping[str, BaseChannel],
+    managed: ManagedValueMapping,
     writes: Sequence[tuple[str, Any]],
 ) -> None:
     for chan, value in writes:
@@ -131,7 +131,7 @@ def local_write(
                 )
             if value.node not in processes:
                 raise InvalidUpdateError(f"Invalid node name {value.node} in packet")
-        elif chan not in channels:
+        elif chan not in channels and chan not in managed:
             logger.warning(f"Skipping write for channel '{chan}' which has no readers")
     commit(writes)
 
@@ -145,7 +145,7 @@ def apply_writes(
     channels: Mapping[str, BaseChannel],
     tasks: Sequence[WritesProtocol],
     get_next_version: Optional[Callable[[int, BaseChannel], int]],
-) -> None:
+) -> dict[str, list[Any]]:
     # update seen versions
     for task in tasks:
         checkpoint["versions_seen"].setdefault(task.name, {}).update(
@@ -161,6 +161,7 @@ def apply_writes(
         max_version = max(checkpoint["channel_versions"].values())
     else:
         max_version = None
+
     # Consume all channels that were read
     for chan in {
         chan for task in tasks for chan in task.triggers if chan not in RESERVED
@@ -177,12 +178,15 @@ def apply_writes(
 
     # Group writes by channel
     pending_writes_by_channel: dict[str, list[Any]] = defaultdict(list)
+    pending_writes_by_managed: dict[str, list[Any]] = defaultdict(list)
     for task in tasks:
         for chan, val in task.writes:
             if chan == TASKS:
                 checkpoint["pending_sends"].append(val)
-            else:
+            elif chan in channels:
                 pending_writes_by_channel[chan].append(val)
+            else:
+                pending_writes_by_managed[chan].append(val)
 
     # Find the highest version of all channels
     if checkpoint["channel_versions"]:
@@ -194,12 +198,7 @@ def apply_writes(
     updated_channels: set[str] = set()
     for chan, vals in pending_writes_by_channel.items():
         if chan in channels:
-            try:
-                updated = channels[chan].update(vals)
-            except InvalidUpdateError as e:
-                raise InvalidUpdateError(
-                    f"Invalid update for channel {chan} with values {vals}"
-                ) from e
+            updated = channels[chan].update(vals)
             if updated and get_next_version is not None:
                 checkpoint["channel_versions"][chan] = get_next_version(
                     max_version, channels[chan]
@@ -214,6 +213,9 @@ def apply_writes(
                     max_version, channels[chan]
                 )
 
+    # Return managed values writes to be applied externally
+    return pending_writes_by_managed
+
 
 @overload
 def prepare_next_tasks(
@@ -314,7 +316,11 @@ def prepare_next_tasks(
                                 CONFIG_KEY_TASK_ID: task_id,
                                 # deque.extend is thread-safe
                                 CONFIG_KEY_SEND: partial(
-                                    local_write, writes.extend, processes, channels
+                                    local_write,
+                                    writes.extend,
+                                    processes,
+                                    channels,
+                                    managed,
                                 ),
                                 CONFIG_KEY_READ: partial(
                                     local_read,
@@ -405,7 +411,11 @@ def prepare_next_tasks(
                                     CONFIG_KEY_TASK_ID: task_id,
                                     # deque.extend is thread-safe
                                     CONFIG_KEY_SEND: partial(
-                                        local_write, writes.extend, processes, channels
+                                        local_write,
+                                        writes.extend,
+                                        processes,
+                                        channels,
+                                        managed,
                                     ),
                                     CONFIG_KEY_READ: partial(
                                         local_read,
@@ -454,16 +464,10 @@ def _proc_input(
                     chan,
                     catch=chan not in proc.triggers,
                 )
+                if chan in channels
+                else managed[k](step)
                 for k, chan in proc.channels.items()
-                if isinstance(chan, str)
             }
-
-            managed_values = {}
-            for key, chan in proc.channels.items():
-                if is_managed_value(chan):
-                    managed_values[key] = managed[key](step)
-
-            val.update(managed_values)
         except EmptyChannelError:
             return
     elif isinstance(proc.channels, list):

libs/langgraph/langgraph/pregel/debug.py

@@ -1,5 +1,6 @@
 import json
 from collections import defaultdict
+from dataclasses import asdict
 from datetime import datetime, timezone
 from pprint import pformat
 from typing import Any, Iterator, Literal, Mapping, Optional, Sequence, TypedDict, Union
@@ -25,6 +26,8 @@ class TaskPayload(TypedDict):
 class TaskResultPayload(TypedDict):
     id: str
     name: str
+    error: Optional[str]
+    interrupts: list[dict]
     result: list[tuple[str, Any]]
 
 
@@ -97,11 +100,14 @@ def map_debug_tasks(
 
 def map_debug_task_results(
     step: int,
-    tasks: list[PregelExecutableTask],
-    stream_channels_list: Sequence[str],
+    tasks: list[tuple[PregelExecutableTask, Sequence[tuple[str, Any]]]],
+    stream_keys: Union[str, Sequence[str]],
 ) -> Iterator[DebugOutputTaskResult]:
+    stream_channels_list = (
+        [stream_keys] if isinstance(stream_keys, str) else stream_keys
+    )
     ts = datetime.now(timezone.utc).isoformat()
-    for name, _, _, writes, config, _, _, _ in tasks:
+    for (name, _, _, _, config, _, _, _), writes in tasks:
         if config is not None and TAG_HIDDEN in config.get("tags", []):
             continue
 
@@ -116,7 +122,9 @@ def map_debug_task_results(
             "payload": {
                 "id": str(uuid5(TASK_NAMESPACE, json.dumps((name, step, metadata)))),
                 "name": name,
+                "error": next((w[1] for w in writes if w[0] == ERROR), None),
                 "result": [w for w in writes if w[0] in stream_channels_list],
+                "interrupts": [asdict(w[1]) for w in writes if w[0] == INTERRUPT],
             },
         }
 
@@ -150,7 +158,7 @@ def map_debug_checkpoint(
                 else {
                     "id": t.id,
                     "name": t.name,
-                    "interrupts": t.interrupts,
+                    "interrupts": tuple(asdict(i) for i in t.interrupts),
                 }
                 for t in tasks_w_writes(tasks, pending_writes)
             ],

libs/langgraph/langgraph/pregel/io.py

@@ -3,7 +3,7 @@ from typing import Any, Iterator, Mapping, Optional, Sequence, TypeVar, Union
 from langchain_core.runnables.utils import AddableDict
 
 from langgraph.channels.base import BaseChannel, EmptyChannelError
-from langgraph.constants import TAG_HIDDEN
+from langgraph.constants import ERROR, INTERRUPT, TAG_HIDDEN
 from langgraph.pregel.log import logger
 from langgraph.pregel.types import PregelExecutableTask
 
@@ -95,19 +95,22 @@ class AddableUpdatesDict(AddableDict):
 
 def map_output_updates(
     output_channels: Union[str, Sequence[str]],
-    tasks: list[PregelExecutableTask],
+    tasks: list[tuple[PregelExecutableTask, Sequence[tuple[str, Any]]]],
 ) -> Iterator[dict[str, Union[Any, dict[str, Any]]]]:
     """Map pending writes (a sequence of tuples (channel, value)) to output chunk."""
     output_tasks = [
-        t for t in tasks if not t.config or TAG_HIDDEN not in t.config.get("tags")
+        (t, ww)
+        for t, ww in tasks
+        if (not t.config or TAG_HIDDEN not in t.config.get("tags"))
+        and all(k not in (ERROR, INTERRUPT) for k, _ in ww)
     ]
     if not output_tasks:
         return
     if isinstance(output_channels, str):
         updated = [
             (task.name, value)
-            for task in output_tasks
-            for chan, value in task.writes
+            for task, writes in output_tasks
+            for chan, value in writes
             if chan == output_channels
         ]
     else:
@@ -116,10 +119,10 @@ def map_output_updates(
                 task.name,
                 {chan: value for chan, value in task.writes if chan in output_channels},
             )
-            for task in output_tasks
-            if any(chan in output_channels for chan, _ in task.writes)
+            for task, writes in output_tasks
+            if any(chan in output_channels for chan, _ in writes)
         ]
-    grouped = {t.name: [] for t in output_tasks}
+    grouped = {t.name: [] for t, _ in output_tasks}
     for node, value in updated:
         grouped[node].append(value)
     for node, value in grouped.items():

libs/langgraph/langgraph/pregel/loop.py

@@ -4,7 +4,6 @@ from collections import deque
 from contextlib import AsyncExitStack, ExitStack
 from types import TracebackType
 from typing import (
-    TYPE_CHECKING,
     Any,
     AsyncContextManager,
     Callable,
@@ -18,6 +17,7 @@ from typing import (
     Type,
     TypeVar,
     Union,
+    cast,
 )
 
 from langchain_core.callbacks import AsyncParentRunManager, ParentRunManager
@@ -25,10 +25,6 @@ from langchain_core.runnables import RunnableConfig
 from typing_extensions import Self
 
 from langgraph.channels.base import BaseChannel
-from langgraph.channels.manager import (
-    AsyncChannelsManager,
-    ChannelsManager,
-)
 from langgraph.checkpoint.base import (
     BaseCheckpointSaver,
     Checkpoint,
@@ -45,13 +41,12 @@ from langgraph.constants import (
     ERROR,
     INPUT,
     INTERRUPT,
-    Interrupt,
 )
 from langgraph.errors import EmptyInputError, GraphInterrupt
 from langgraph.managed.base import (
-    AsyncManagedValuesManager,
     ManagedValueMapping,
-    ManagedValuesManager,
+    ManagedValueSpec,
+    WritableManagedValue,
 )
 from langgraph.pregel.algo import (
     PregelTaskWrites,
@@ -60,19 +55,29 @@ from langgraph.pregel.algo import (
     prepare_next_tasks,
     should_interrupt,
 )
-from langgraph.pregel.debug import map_debug_checkpoint, map_debug_tasks
+from langgraph.pregel.debug import (
+    map_debug_checkpoint,
+    map_debug_task_results,
+    map_debug_tasks,
+)
 from langgraph.pregel.executor import (
     AsyncBackgroundExecutor,
     BackgroundExecutor,
     Submit,
 )
-from langgraph.pregel.io import map_input, map_output_updates, map_output_values, single
+from langgraph.pregel.io import (
+    map_input,
+    map_output_updates,
+    map_output_values,
+    read_channels,
+    single,
+)
+from langgraph.pregel.manager import AsyncChannelsManager, ChannelsManager
+from langgraph.pregel.read import PregelNode
 from langgraph.pregel.types import PregelExecutableTask
 from langgraph.pregel.utils import get_new_channel_versions
-
-if TYPE_CHECKING:
-    from langgraph.pregel import Pregel
-
+from langgraph.store.base import BaseStore
+from langgraph.store.batch import AsyncBatchedStore
 
 V = TypeVar("V")
 INPUT_DONE = object()
@@ -83,7 +88,14 @@ EMPTY_SEQ = ()
 class PregelLoop:
     input: Optional[Any]
     config: RunnableConfig
+    store: Optional[BaseStore]
     checkpointer: Optional[BaseCheckpointSaver]
+    nodes: Mapping[str, PregelNode]
+    specs: Mapping[str, Union[BaseChannel, ManagedValueSpec]]
+    output_keys: Union[str, Sequence[str]]
+    stream_keys: Union[str, Sequence[str]]
+    is_nested: bool
+
     checkpointer_get_next_version: Callable[[Optional[V]], V]
     checkpointer_put_writes: Optional[
         Callable[[RunnableConfig, Sequence[tuple[str, Any]], str], Any]
@@ -99,8 +111,6 @@ class PregelLoop:
             Any,
         ]
     ]
-    graph: "Pregel"
-
     submit: Submit
     channels: Mapping[str, BaseChannel]
     managed: ManagedValueMapping
@@ -118,7 +128,7 @@ class PregelLoop:
     ]
     tasks: Sequence[PregelExecutableTask]
     stream: deque[Tuple[str, Any]]
-    is_nested: bool
+    output: Union[None, dict[str, Any], Any] = None
 
     # public
 
@@ -127,24 +137,28 @@ class PregelLoop:
         input: Optional[Any],
         *,
         config: RunnableConfig,
+        store: Optional[BaseStore],
         checkpointer: Optional[BaseCheckpointSaver],
-        graph: "Pregel",
+        nodes: Mapping[str, PregelNode],
+        specs: Mapping[str, Union[BaseChannel, ManagedValueSpec]],
+        output_keys: Union[str, Sequence[str]],
+        stream_keys: Union[str, Sequence[str]],
     ) -> None:
         self.stream = deque()
         self.input = input
         self.config = config
+        self.store = store
         self.checkpointer = checkpointer
-        self.graph = graph
-        # TODO if managed values no longer needs graph we can replace with
-        # managed_specs, channel_specs
+        self.nodes = nodes
+        self.specs = specs
+        self.output_keys = output_keys
+        self.stream_keys = stream_keys
         self.is_nested = CONFIG_KEY_READ in self.config.get("configurable", {})
 
-    def mark_tasks_scheduled(self, tasks: Sequence[PregelExecutableTask]) -> None:
-        """Mark tasks as scheduled, to be used by queue-based executors."""
-        raise NotImplementedError
-
     def put_writes(self, task_id: str, writes: Sequence[tuple[str, Any]]) -> None:
         """Put writes for a task, to be read by the next tick."""
+        if not writes:
+            return
         self.checkpoint_pending_writes.extend((task_id, k, v) for k, v in writes)
         if self.checkpointer_put_writes is not None:
             self.submit(
@@ -162,11 +176,22 @@ class PregelLoop:
                 writes,
                 task_id,
             )
+        if task := next((t for t in self.tasks if t.id == task_id), None):
+            self.stream.extend(
+                ("updates", v)
+                for v in map_output_updates(self.output_keys, [(task, writes)])
+            )
+            self.stream.extend(
+                ("debug", v)
+                for v in map_debug_task_results(
+                    self.step, [(task, writes)], self.stream_keys
+                )
+            )
 
     def tick(
         self,
         *,
-        output_keys: Union[str, Sequence[str]] = EMPTY_SEQ,
+        input_keys: Union[str, Sequence[str]],
         interrupt_after: Sequence[str] = EMPTY_SEQ,
         interrupt_before: Sequence[str] = EMPTY_SEQ,
         manager: Union[None, AsyncParentRunManager, ParentRunManager] = None,
@@ -178,20 +203,23 @@ class PregelLoop:
             raise RuntimeError("Cannot tick when status is no longer 'pending'")
 
         if self.input not in (INPUT_DONE, INPUT_RESUMING):
-            self._first()
+            self._first(input_keys=input_keys)
         elif all(task.writes for task in self.tasks):
             writes = [w for t in self.tasks for w in t.writes]
             # all tasks have finished
-            apply_writes(
+            mv_writes = apply_writes(
                 self.checkpoint,
                 self.channels,
                 self.tasks,
                 self.checkpointer_get_next_version,
             )
+            # apply writes to managed values
+            for key, values in mv_writes.items():
+                self._update_mv(key, values)
             # produce values output
             self.stream.extend(
                 ("values", v)
-                for v in map_output_values(output_keys, writes, self.channels)
+                for v in map_output_values(self.output_keys, writes, self.channels)
             )
             # clear pending writes
             self.checkpoint_pending_writes.clear()
@@ -200,20 +228,17 @@ class PregelLoop:
                 {
                     "source": "loop",
                     "writes": single(
-                        map_output_updates(output_keys, self.tasks)
-                        if self.graph.stream_mode == "updates"
-                        else map_output_values(output_keys, writes, self.channels)
+                        map_output_updates(
+                            self.output_keys, [(t, t.writes) for t in self.tasks]
+                        )
                     ),
                 }
             )
             # after execution, check if we should interrupt
-            if tasks := should_interrupt(self.checkpoint, interrupt_after, self.tasks):
+            if should_interrupt(self.checkpoint, interrupt_after, self.tasks):
                 self.status = "interrupt_after"
-                interrupts = [(t.id, Interrupt("after")) for t in tasks]
-                for tid, interrupt in interrupts:
-                    self.put_writes(tid, [(INTERRUPT, interrupt)])
                 if self.is_nested:
-                    raise GraphInterrupt([i[1] for i in interrupts])
+                    raise GraphInterrupt()
                 else:
                     return False
         else:
@@ -227,7 +252,7 @@ class PregelLoop:
         # prepare next tasks
         self.tasks = prepare_next_tasks(
             self.checkpoint,
-            self.graph.nodes,
+            self.nodes,
             self.channels,
             self.managed,
             self.config,
@@ -246,7 +271,7 @@ class PregelLoop:
                     self.step - 1,  # printing checkpoint for previous step
                     self.checkpoint_config,
                     self.channels,
-                    self.graph.stream_channels_asis,
+                    self.stream_keys,
                     self.checkpoint_metadata,
                     self.checkpoint,
                     self.tasks,
@@ -270,20 +295,17 @@ class PregelLoop:
         # if all tasks have finished, re-tick
         if all(task.writes for task in self.tasks):
             return self.tick(
-                output_keys=output_keys,
+                input_keys=input_keys,
                 interrupt_after=interrupt_after,
                 interrupt_before=interrupt_before,
                 manager=manager,
             )
 
         # before execution, check if we should interrupt
-        if tasks := should_interrupt(self.checkpoint, interrupt_before, self.tasks):
+        if should_interrupt(self.checkpoint, interrupt_before, self.tasks):
             self.status = "interrupt_before"
-            interrupts = [(t.id, Interrupt("before")) for t in tasks]
-            for tid, interrupt in interrupts:
-                self.put_writes(tid, [(INTERRUPT, interrupt)])
             if self.is_nested:
-                raise GraphInterrupt([i[1] for i in interrupts])
+                raise GraphInterrupt()
             else:
                 return False
 
@@ -294,7 +316,7 @@ class PregelLoop:
 
     # private
 
-    def _first(self) -> None:
+    def _first(self, *, input_keys: Union[str, Sequence[str]]) -> None:
         # resuming from previous checkpoint requires
         # - finding a previous checkpoint
         # - receiving None input (outer graph) or RESUMING flag (subgraph)
@@ -311,11 +333,11 @@ class PregelLoop:
                     version = self.checkpoint["channel_versions"][k]
                     self.checkpoint["versions_seen"][INTERRUPT][k] = version
         # map inputs to channel updates
-        elif input_writes := deque(map_input(self.graph.input_channels, self.input)):
+        elif input_writes := deque(map_input(input_keys, self.input)):
             # discard any unfinished tasks from previous checkpoint
             discard_tasks = prepare_next_tasks(
                 self.checkpoint,
-                self.graph.nodes,
+                self.nodes,
                 self.channels,
                 self.managed,
                 self.config,
@@ -324,16 +346,16 @@ class PregelLoop:
                 manager=None,
             )
             # apply input writes
-            apply_writes(
+            assert not apply_writes(
                 self.checkpoint,
                 self.channels,
                 discard_tasks + [PregelTaskWrites(INPUT, input_writes, [])],
                 self.checkpointer_get_next_version,
-            )
+            ), "Can't write to SharedValues in graph input"
             # save input checkpoint
             self._put_checkpoint({"source": "input", "writes": self.input})
         else:
-            raise EmptyInputError(f"Received no input for {self.graph.input_channels}")
+            raise EmptyInputError(f"Received no input for {input_keys}")
         # done with input
         self.input = INPUT_RESUMING if is_resuming else INPUT_DONE
 
@@ -395,13 +417,21 @@ class PregelLoop:
         # increment step
         self.step += 1
 
+    def _update_mv(self, key: str, values: Sequence[Any]) -> None:
+        raise NotImplementedError
+
     def _suppress_interrupt(
         self,
         exc_type: Optional[Type[BaseException]],
         exc_value: Optional[BaseException],
         traceback: Optional[TracebackType],
     ) -> Optional[bool]:
-        if isinstance(exc_value, GraphInterrupt) and not self.is_nested:
+        suppress = isinstance(exc_value, GraphInterrupt) and not self.is_nested
+        if suppress or exc_type is None:
+            # save final output
+            self.output = read_channels(self.channels, self.output_keys)
+        if suppress:
+            # suppress interrupt
             return True
 
 
@@ -411,10 +441,23 @@ class SyncPregelLoop(PregelLoop, ContextManager):
         input: Optional[Any],
         *,
         config: RunnableConfig,
+        store: Optional[BaseStore],
         checkpointer: Optional[BaseCheckpointSaver],
-        graph: "Pregel",
+        nodes: Mapping[str, PregelNode],
+        specs: Mapping[str, Union[BaseChannel, ManagedValueSpec]],
+        output_keys: Union[str, Sequence[str]] = EMPTY_SEQ,
+        stream_keys: Union[str, Sequence[str]] = EMPTY_SEQ,
     ) -> None:
-        super().__init__(input, config=config, checkpointer=checkpointer, graph=graph)
+        super().__init__(
+            input,
+            config=config,
+            checkpointer=checkpointer,
+            store=store,
+            nodes=nodes,
+            specs=specs,
+            output_keys=output_keys,
+            stream_keys=stream_keys,
+        )
         self.stack = ExitStack()
         if checkpointer:
             self.checkpointer_get_next_version = checkpointer.get_next_version
@@ -438,6 +481,9 @@ class SyncPregelLoop(PregelLoop, ContextManager):
         finally:
             self.checkpointer.put(config, checkpoint, metadata, new_versions)
 
+    def _update_mv(self, key: str, values: Sequence[Any]) -> None:
+        return self.submit(cast(WritableManagedValue, self.managed[key]).update, values)
+
     # context manager
 
     def __enter__(self) -> Self:
@@ -457,11 +503,8 @@ class SyncPregelLoop(PregelLoop, ContextManager):
         self.checkpoint_pending_writes = saved.pending_writes or []
 
         self.submit = self.stack.enter_context(BackgroundExecutor(self.config))
-        self.channels = self.stack.enter_context(
-            ChannelsManager(self.graph.channels, self.checkpoint, self.config)
-        )
-        self.managed = self.stack.enter_context(
-            ManagedValuesManager(self.graph.managed_values_dict, self.config)
+        self.channels, self.managed = self.stack.enter_context(
+            ChannelsManager(self.specs, self.checkpoint, self.config, self.store)
         )
         self.stack.push(self._suppress_interrupt)
         self.status = "pending"
@@ -478,7 +521,6 @@ class SyncPregelLoop(PregelLoop, ContextManager):
         traceback: Optional[TracebackType],
     ) -> Optional[bool]:
         # unwind stack
-        del self.graph
         return self.stack.__exit__(exc_type, exc_value, traceback)
 
 
@@ -488,10 +530,24 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager):
         input: Optional[Any],
         *,
         config: RunnableConfig,
+        store: Optional[BaseStore],
         checkpointer: Optional[BaseCheckpointSaver],
-        graph: "Pregel",
+        nodes: Mapping[str, PregelNode],
+        specs: Mapping[str, Union[BaseChannel, ManagedValueSpec]],
+        output_keys: Union[str, Sequence[str]] = EMPTY_SEQ,
+        stream_keys: Union[str, Sequence[str]] = EMPTY_SEQ,
     ) -> None:
-        super().__init__(input, config=config, checkpointer=checkpointer, graph=graph)
+        super().__init__(
+            input,
+            config=config,
+            checkpointer=checkpointer,
+            store=store,
+            nodes=nodes,
+            specs=specs,
+            output_keys=output_keys,
+            stream_keys=stream_keys,
+        )
+        self.store = AsyncBatchedStore(self.store) if self.store else None
         self.stack = AsyncExitStack()
         if checkpointer:
             self.checkpointer_get_next_version = checkpointer.get_next_version
@@ -515,6 +571,11 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager):
         finally:
             await self.checkpointer.aput(config, checkpoint, metadata, new_versions)
 
+    def _update_mv(self, key: str, values: Sequence[Any]) -> None:
+        return self.submit(
+            cast(WritableManagedValue, self.managed[key]).aupdate, values
+        )
+
     # context manager
 
     async def __aenter__(self) -> Self:
@@ -536,11 +597,8 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager):
         self.checkpoint_pending_writes = saved.pending_writes or []
 
         self.submit = await self.stack.enter_async_context(AsyncBackgroundExecutor())
-        self.channels = await self.stack.enter_async_context(
-            AsyncChannelsManager(self.graph.channels, self.checkpoint, self.config)
-        )
-        self.managed = await self.stack.enter_async_context(
-            AsyncManagedValuesManager(self.graph.managed_values_dict, self.config)
+        self.channels, self.managed = await self.stack.enter_async_context(
+            AsyncChannelsManager(self.specs, self.checkpoint, self.config, self.store)
         )
         self.stack.push(self._suppress_interrupt)
         self.status = "pending"
@@ -558,7 +616,6 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager):
         traceback: Optional[TracebackType],
     ) -> Optional[bool]:
         # unwind stack
-        del self.graph
         return await asyncio.shield(
             self.stack.__aexit__(exc_type, exc_value, traceback)
         )

libs/langgraph/langgraph/pregel/manager.py

@@ -0,0 +1,106 @@
+import asyncio
+from contextlib import AsyncExitStack, ExitStack, asynccontextmanager, contextmanager
+from typing import AsyncIterator, Iterator, Mapping, Optional, Union
+
+from langchain_core.runnables import RunnableConfig, patch_config
+
+from langgraph.channels.base import BaseChannel
+from langgraph.channels.context import Context
+from langgraph.channels.last_value import LastValue
+from langgraph.checkpoint.base import Checkpoint
+from langgraph.constants import CONFIG_KEY_STORE
+from langgraph.managed.base import (
+    ConfiguredManagedValue,
+    ManagedValueMapping,
+    ManagedValueSpec,
+)
+from langgraph.store.base import BaseStore
+
+
+@contextmanager
+def ChannelsManager(
+    specs: Mapping[str, Union[BaseChannel, ManagedValueSpec]],
+    checkpoint: Checkpoint,
+    config: RunnableConfig,
+    store: Optional[BaseStore] = None,
+    *,
+    skip_context: bool = False,
+) -> Iterator[tuple[Mapping[str, BaseChannel], ManagedValueMapping]]:
+    """Manage channels for the lifetime of a Pregel invocation (multiple steps)."""
+    config_for_managed = patch_config(config, configurable={CONFIG_KEY_STORE: store})
+    channel_specs: Mapping[str, BaseChannel] = {}
+    managed_specs: Mapping[str, ManagedValueSpec] = {}
+    for k, v in specs.items():
+        if skip_context and isinstance(v, Context):
+            channel_specs[k] = LastValue(None)
+        elif isinstance(v, BaseChannel):
+            channel_specs[k] = v
+        else:
+            managed_specs[k] = v
+    with ExitStack() as stack:
+        yield (
+            {
+                k: stack.enter_context(
+                    v.from_checkpoint_named(checkpoint["channel_values"].get(k), config)
+                )
+                for k, v in channel_specs.items()
+            },
+            {
+                key: stack.enter_context(
+                    value.cls.enter(config_for_managed, **value.kwargs)
+                    if isinstance(value, ConfiguredManagedValue)
+                    else value.enter(config_for_managed)
+                )
+                for key, value in managed_specs.items()
+            },
+        )
+
+
+@asynccontextmanager
+async def AsyncChannelsManager(
+    specs: Mapping[str, Union[BaseChannel, ManagedValueSpec]],
+    checkpoint: Checkpoint,
+    config: RunnableConfig,
+    store: Optional[BaseStore] = None,
+    *,
+    skip_context: bool = False,
+) -> AsyncIterator[Mapping[str, BaseChannel]]:
+    """Manage channels for the lifetime of a Pregel invocation (multiple steps)."""
+    config_for_managed = patch_config(config, configurable={CONFIG_KEY_STORE: store})
+    channel_specs: Mapping[str, BaseChannel] = {}
+    managed_specs: Mapping[str, ManagedValueSpec] = {}
+    for k, v in specs.items():
+        if skip_context and isinstance(v, Context):
+            channel_specs[k] = LastValue(None)
+        elif isinstance(v, BaseChannel):
+            channel_specs[k] = v
+        else:
+            managed_specs[k] = v
+    async with AsyncExitStack() as stack:
+        # managed: create enter tasks with reference to spec, await them
+        if tasks := {
+            asyncio.create_task(
+                stack.enter_async_context(
+                    value.cls.aenter(config_for_managed, **value.kwargs)
+                    if isinstance(value, ConfiguredManagedValue)
+                    else value.aenter(config_for_managed)
+                )
+            ): key
+            for key, value in managed_specs.items()
+        }:
+            done, _ = await asyncio.wait(tasks, return_when=asyncio.ALL_COMPLETED)
+        else:
+            done = set()
+        yield (
+            # channels: enter each channel with checkpoint
+            {
+                k: await stack.enter_async_context(
+                    v.afrom_checkpoint_named(
+                        checkpoint["channel_values"].get(k), config
+                    )
+                )
+                for k, v in channel_specs.items()
+            },
+            # managed: build mapping from spec to result
+            {tasks[task]: task.result() for task in done},
+        )

libs/langgraph/langgraph/pregel/read.py

@@ -15,7 +15,6 @@ from langchain_core.runnables.config import merge_configs
 from langchain_core.runnables.utils import ConfigurableFieldSpec
 
 from langgraph.constants import CONFIG_KEY_READ
-from langgraph.managed.base import ManagedValueSpec
 from langgraph.pregel.retry import RetryPolicy
 from langgraph.pregel.write import ChannelWrite
 from langgraph.utils import RunnableCallable
@@ -101,7 +100,7 @@ DEFAULT_BOUND: RunnablePassthrough = RunnablePassthrough()
 
 
 class PregelNode(RunnableBindingBase):
-    channels: Union[list[str], Mapping[str, Union[str, ManagedValueSpec]]]
+    channels: Union[list[str], Mapping[str, str]]
 
     triggers: list[str] = Field(default_factory=list)
 

libs/langgraph/langgraph/pregel/retry.py

@@ -4,6 +4,7 @@ import random
 import time
 from typing import Optional
 
+from langgraph.errors import GraphInterrupt
 from langgraph.pregel.types import PregelExecutableTask, RetryPolicy
 
 logger = logging.getLogger(__name__)
@@ -25,6 +26,9 @@ def run_with_retry(
             task.proc.invoke(task.input, task.config)
             # if successful, end
             break
+        except GraphInterrupt:
+            # if interrupted, end
+            raise
         except Exception as exc:
             if retry_policy is None:
                 raise
@@ -75,6 +79,9 @@ async def arun_with_retry(
                 await task.proc.ainvoke(task.input, task.config)
             # if successful, end
             break
+        except GraphInterrupt:
+            # if interrupted, end
+            raise
         except Exception as exc:
             if retry_policy is None:
                 raise

libs/langgraph/langgraph/store/base.py

@@ -0,0 +1,21 @@
+from typing import Any, List, Optional
+
+V = dict[str, Any]
+
+
+class BaseStore:
+    def list(self, prefixes: List[str]) -> dict[str, dict[str, V]]:
+        # list[namespace] -> dict[namespace, list[value]]
+        raise NotImplementedError
+
+    def put(self, writes: List[tuple[str, str, Optional[V]]]) -> None:
+        # list[(namespace, key, value | none)] -> None
+        raise NotImplementedError
+
+    async def alist(self, prefixes: List[str]) -> dict[str, dict[str, V]]:
+        # list[namespace] -> dict[namespace, list[value]]
+        raise NotImplementedError
+
+    async def aput(self, writes: List[tuple[str, str, Optional[V]]]) -> None:
+        # list[(namespace, key, value | none)] -> None
+        raise NotImplementedError

libs/langgraph/langgraph/store/batch.py

@@ -0,0 +1,65 @@
+import asyncio
+from typing import NamedTuple, Optional, Union
+
+from langgraph.store.base import BaseStore, V
+
+
+class ListOp(NamedTuple):
+    prefixes: list[str]
+
+
+class PutOp(NamedTuple):
+    writes: list[tuple[str, str, Optional[V]]]
+
+
+class AsyncBatchedStore(BaseStore):
+    def __init__(self, store: BaseStore) -> None:
+        self.store = store
+        self.aqueue: dict[asyncio.Future, Union[ListOp, PutOp]] = {}
+        self.task = asyncio.create_task(_run(self.aqueue, self.store))
+
+    def __del__(self) -> None:
+        self.task.cancel()
+
+    async def alist(self, prefixes: list[str]) -> dict[str, dict[str, V]]:
+        fut = asyncio.get_running_loop().create_future()
+        self.aqueue[fut] = ListOp(prefixes)
+        return await fut
+
+    async def aput(self, writes: list[tuple[str, str, Optional[V]]]) -> None:
+        fut = asyncio.get_running_loop().create_future()
+        self.aqueue[fut] = PutOp(writes)
+        return await fut
+
+
+async def _run(
+    aqueue: dict[asyncio.Future, Union[ListOp, PutOp]], store: BaseStore
+) -> None:
+    while True:
+        await asyncio.sleep(0)
+        if not aqueue:
+            continue
+        # this could use a lock, if we want thread safety
+        taken = aqueue.copy()
+        aqueue.clear()
+        # action each operation
+        lists = {f: o for f, o in taken.items() if isinstance(o, ListOp)}
+        if lists:
+            try:
+                results = await store.alist(
+                    [p for op in lists.values() for p in op.prefixes]
+                )
+                for fut, op in lists.items():
+                    fut.set_result({k: results.get(k) for k in op.prefixes})
+            except Exception as e:
+                for fut in lists:
+                    fut.set_exception(e)
+        puts = {f: o for f, o in taken.items() if isinstance(o, PutOp)}
+        if puts:
+            try:
+                await store.aput([w for op in puts.values() for w in op.writes])
+                for fut in puts:
+                    fut.set_result(None)
+            except Exception as e:
+                for fut in puts:
+                    fut.set_exception(e)

libs/langgraph/langgraph/store/memory.py

@@ -0,0 +1,25 @@
+from collections import defaultdict
+from typing import List, Optional
+
+from langgraph.store.base import BaseStore, V
+
+
+class MemoryStore(BaseStore):
+    def __init__(self) -> None:
+        self.data: dict[str, dict[str, V]] = defaultdict(dict)
+
+    def list(self, prefixes: List[str]) -> dict[str, dict[str, V]]:
+        return {prefix: self.data[prefix] for prefix in prefixes}
+
+    async def alist(self, prefixes: List[str]) -> dict[str, dict[str, V]]:
+        return self.list(prefixes)
+
+    def put(self, writes: List[tuple[str, str, Optional[V]]]) -> None:
+        for namespace, key, value in writes:
+            if value is None:
+                self.data[namespace].pop(key, None)
+            else:
+                self.data[namespace][key] = value
+
+    async def aput(self, writes: List[tuple[str, str, Optional[V]]]) -> None:
+        return self.put(writes)

libs/langgraph/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "langgraph"
-version = "0.2.6"
+version = "0.2.12"
 description = "Building stateful, multi-actor applications with LLMs"
 authors = []
 license = "MIT"
@@ -65,7 +65,7 @@ omit = ["tests/*"]
 [tool.pytest-watcher]
 now = true
 delay = 0.1
-runner_args = ["--ff", "-v", "-n", "auto", "--dist", "worksteal", "--snapshot-update", "--tb", "short"]
+runner_args = ["--ff", "-v", "-x", "-n", "auto", "--dist", "worksteal", "--snapshot-update", "--tb", "short"]
 patterns = ["*.py"]
 
 [build-system]

libs/langgraph/tests/any_str.py

@@ -1,3 +1,6 @@
+from typing import Any, Sequence
+
+
 class AnyStr(str):
     def __init__(self) -> None:
         super().__init__()
@@ -9,6 +12,17 @@ class AnyStr(str):
         return hash(str(self))
 
 
+class AnyVersion:
+    def __init__(self) -> None:
+        super().__init__()
+
+    def __eq__(self, other: object) -> bool:
+        return isinstance(other, (str, int, float))
+
+    def __hash__(self) -> int:
+        return hash(str(self))
+
+
 class ExceptionLike:
     def __init__(self, exc: Exception) -> None:
         self.exc = exc
@@ -22,3 +36,24 @@ class ExceptionLike:
 
     def __hash__(self) -> int:
         return hash((self.exc.__class__, str(self.exc)))
+
+    def __repr__(self) -> str:
+        return str(self.exc)
+
+
+class UnsortedSequence:
+    def __init__(self, *values: Any) -> None:
+        self.seq = values
+
+    def __eq__(self, value: object) -> bool:
+        return (
+            isinstance(value, Sequence)
+            and len(self.seq) == len(value)
+            and all(a in value for a in self.seq)
+        )
+
+    def __hash__(self) -> int:
+        return hash(frozenset(self.seq))
+
+    def __repr__(self) -> str:
+        return repr(self.seq)

libs/langgraph/tests/fake_tracer.py

@@ -0,0 +1,91 @@
+from typing import Any, Optional
+from uuid import UUID
+
+from langchain_core.messages.base import BaseMessage
+from langchain_core.outputs.chat_generation import ChatGeneration
+from langchain_core.outputs.llm_result import LLMResult
+from langchain_core.tracers import BaseTracer, Run
+
+
+class FakeTracer(BaseTracer):
+    """Fake tracer that records LangChain execution.
+    It replaces run ids with deterministic UUIDs for snapshotting."""
+
+    def __init__(self) -> None:
+        """Initialize the tracer."""
+        super().__init__()
+        self.runs: list[Run] = []
+        self.uuids_map: dict[UUID, UUID] = {}
+        self.uuids_generator = (
+            UUID(f"00000000-0000-4000-8000-{i:012}", version=4) for i in range(10000)
+        )
+
+    def _replace_uuid(self, uuid: UUID) -> UUID:
+        if uuid not in self.uuids_map:
+            self.uuids_map[uuid] = next(self.uuids_generator)
+        return self.uuids_map[uuid]
+
+    def _replace_message_id(self, maybe_message: Any) -> Any:
+        if isinstance(maybe_message, BaseMessage):
+            maybe_message.id = str(next(self.uuids_generator))
+        if isinstance(maybe_message, ChatGeneration):
+            maybe_message.message.id = str(next(self.uuids_generator))
+        if isinstance(maybe_message, LLMResult):
+            for i, gen_list in enumerate(maybe_message.generations):
+                for j, gen in enumerate(gen_list):
+                    maybe_message.generations[i][j] = self._replace_message_id(gen)
+        if isinstance(maybe_message, dict):
+            for k, v in maybe_message.items():
+                maybe_message[k] = self._replace_message_id(v)
+        if isinstance(maybe_message, list):
+            for i, v in enumerate(maybe_message):
+                maybe_message[i] = self._replace_message_id(v)
+
+        return maybe_message
+
+    def _copy_run(self, run: Run) -> Run:
+        if run.dotted_order:
+            levels = run.dotted_order.split(".")
+            processed_levels = []
+            for level in levels:
+                timestamp, run_id = level.split("Z")
+                new_run_id = self._replace_uuid(UUID(run_id))
+                processed_level = f"{timestamp}Z{new_run_id}"
+                processed_levels.append(processed_level)
+            new_dotted_order = ".".join(processed_levels)
+        else:
+            new_dotted_order = None
+        return run.copy(
+            update={
+                "id": self._replace_uuid(run.id),
+                "parent_run_id": (
+                    self.uuids_map[run.parent_run_id] if run.parent_run_id else None
+                ),
+                "child_runs": [self._copy_run(child) for child in run.child_runs],
+                "trace_id": self._replace_uuid(run.trace_id) if run.trace_id else None,
+                "dotted_order": new_dotted_order,
+                "inputs": self._replace_message_id(run.inputs),
+                "outputs": self._replace_message_id(run.outputs),
+            }
+        )
+
+    def _persist_run(self, run: Run) -> None:
+        """Persist a run."""
+
+        self.runs.append(self._copy_run(run))
+
+    def flattened_runs(self) -> list[Run]:
+        q = [] + self.runs
+        result = []
+        while q:
+            parent = q.pop()
+            result.append(parent)
+            if parent.child_runs:
+                q.extend(parent.child_runs)
+        return result
+
+    @property
+    def run_ids(self) -> list[Optional[UUID]]:
+        runs = self.flattened_runs()
+        uuids_map = {v: k for k, v in self.uuids_map.items()}
+        return [uuids_map.get(r.id) for r in runs]

libs/langgraph/tests/test_algo.py

@@ -1,7 +1,6 @@
-from langgraph.channels.manager import ChannelsManager
 from langgraph.checkpoint.base import empty_checkpoint
-from langgraph.managed.base import ManagedValuesManager
 from langgraph.pregel.algo import prepare_next_tasks
+from langgraph.pregel.manager import ChannelsManager
 
 
 def test_prepare_next_tasks() -> None:
@@ -9,9 +8,7 @@ def test_prepare_next_tasks() -> None:
     processes = {}
     checkpoint = empty_checkpoint()
 
-    with ManagedValuesManager({}, config) as managed, ChannelsManager(
-        {}, checkpoint, config
-    ) as channels:
+    with ChannelsManager({}, checkpoint, config) as (channels, managed):
         assert (
             prepare_next_tasks(
                 checkpoint, processes, channels, managed, config, 0, for_execution=False

libs/langgraph/tests/test_pregel.py

@@ -58,6 +58,7 @@ from langgraph.graph import END, Graph
 from langgraph.graph.graph import START
 from langgraph.graph.message import MessageGraph, add_messages
 from langgraph.graph.state import StateGraph
+from langgraph.managed.shared_value import SharedValue
 from langgraph.prebuilt.chat_agent_executor import (
     create_tool_calling_executor,
 )
@@ -65,7 +66,9 @@ from langgraph.prebuilt.tool_node import ToolNode
 from langgraph.pregel import Channel, GraphRecursionError, Pregel, StateSnapshot
 from langgraph.pregel.retry import RetryPolicy
 from langgraph.pregel.types import PregelTask
-from tests.any_str import AnyStr, ExceptionLike
+from langgraph.store.memory import MemoryStore
+from tests.any_str import AnyStr, AnyVersion, ExceptionLike, UnsortedSequence
+from tests.fake_tracer import FakeTracer
 from tests.memory_assert import (
     MemorySaverAssertCheckpointMetadata,
     MemorySaverAssertImmutable,
@@ -195,6 +198,21 @@ def test_graph_validation() -> None:
     with pytest.raises(ValueError, match="Invalid reducer"):
         StateGraph(BadReducerState)
 
+    def node_b(state: State) -> State:
+        return {"hello": "world"}
+
+    builder = StateGraph(State)
+    builder.add_node("a", node_b)
+    builder.add_node("b", node_b)
+    builder.add_node("c", node_b)
+    builder.set_entry_point("a")
+    builder.add_edge("a", "b")
+    builder.add_edge("a", "c")
+    graph = builder.compile()
+
+    with pytest.raises(InvalidUpdateError, match="At key 'hello'"):
+        graph.invoke({"hello": "there"})
+
 
 def test_checkpoint_errors() -> None:
     class FaultyGetCheckpointer(MemorySaver):
@@ -655,7 +673,7 @@ def test_invoke_two_processes_in_out_interrupt(
                     "checkpoint_id": AnyStr(),
                 }
             },
-            metadata={"source": "loop", "step": 6, "writes": 5},
+            metadata={"source": "loop", "step": 6, "writes": {"two": 5}},
             created_at=AnyStr(),
             parent_config=history[1].config,
         ),
@@ -670,7 +688,7 @@ def test_invoke_two_processes_in_out_interrupt(
                     "checkpoint_id": AnyStr(),
                 }
             },
-            metadata={"source": "loop", "step": 5, "writes": None},
+            metadata={"source": "loop", "step": 5, "writes": {"one": None}},
             created_at=AnyStr(),
             parent_config=history[2].config,
         ),
@@ -700,7 +718,7 @@ def test_invoke_two_processes_in_out_interrupt(
                     "checkpoint_id": AnyStr(),
                 }
             },
-            metadata={"source": "loop", "step": 3, "writes": None},
+            metadata={"source": "loop", "step": 3, "writes": {"one": None}},
             created_at=AnyStr(),
             parent_config=history[4].config,
         ),
@@ -730,7 +748,7 @@ def test_invoke_two_processes_in_out_interrupt(
                     "checkpoint_id": AnyStr(),
                 }
             },
-            metadata={"source": "loop", "step": 1, "writes": 4},
+            metadata={"source": "loop", "step": 1, "writes": {"two": 4}},
             created_at=AnyStr(),
             parent_config=history[6].config,
         ),
@@ -745,7 +763,7 @@ def test_invoke_two_processes_in_out_interrupt(
                     "checkpoint_id": AnyStr(),
                 }
             },
-            metadata={"source": "loop", "step": 0, "writes": None},
+            metadata={"source": "loop", "step": 0, "writes": {"one": None}},
             created_at=AnyStr(),
             parent_config=history[7].config,
         ),
@@ -766,7 +784,7 @@ def test_invoke_two_processes_in_out_interrupt(
         ),
     ]
 
-    # forking from any previous checkpoint w/out forking should do nothing
+    # re-running from any previous checkpoint w/out forking should do nothing
     assert [c for c in app.stream(None, history[0].config, stream_mode="updates")] == []
     assert [c for c in app.stream(None, history[1].config, stream_mode="updates")] == []
     assert [c for c in app.stream(None, history[2].config, stream_mode="updates")] == []
@@ -1030,6 +1048,8 @@ def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None:
                 "id": "2687f72c-e3a8-5f6f-9afa-047cbf24e923",
                 "name": "one",
                 "result": [("inbox", 3)],
+                "error": None,
+                "interrupts": [],
             },
         },
         {
@@ -1040,6 +1060,8 @@ def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None:
                 "id": "18f52f6a-828d-58a1-a501-53cc0c7af33e",
                 "name": "two",
                 "result": [("output", 13)],
+                "error": None,
+                "interrupts": [],
             },
         },
         {
@@ -1061,6 +1083,8 @@ def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None:
                 "id": "871d6e74-7bb3-565f-a4fe-cef4b8f19b62",
                 "name": "two",
                 "result": [("output", 4)],
+                "error": None,
+                "interrupts": [],
             },
         },
     ]
@@ -1188,6 +1212,21 @@ def test_invoke_two_processes_two_in_two_out_invalid(mocker: MockerFixture) -> N
         # LastValue channels can only be updated once per iteration
         app.invoke(2)
 
+    class State(TypedDict):
+        hello: str
+
+    def my_node(input: State) -> State:
+        return {"hello": "world"}
+
+    builder = StateGraph(State)
+    builder.add_node("one", my_node)
+    builder.add_node("two", my_node)
+    builder.set_conditional_entry_point(lambda _: ["one", "two"])
+
+    graph = builder.compile()
+    with pytest.raises(InvalidUpdateError, match="At key 'hello'"):
+        graph.invoke({"hello": "there"}, debug=True)
+
 
 def test_invoke_two_processes_two_in_two_out_valid(mocker: MockerFixture) -> None:
     add_one = mocker.Mock(side_effect=lambda x: x + 1)
@@ -1372,6 +1411,144 @@ def test_pending_writes_resume(
     # both the pending write and the new write were applied, 1 + 2 + 3 = 6
     assert graph.invoke(None, thread1) == {"value": 6}
 
+    # check all final checkpoints
+    checkpoints = [c for c in checkpointer.list(thread1)]
+    # we should have 3
+    assert len(checkpoints) == 3
+    # the last one not too interesting for this test
+    assert checkpoints[0] == CheckpointTuple(
+        config={
+            "configurable": {
+                "thread_id": "1",
+                "checkpoint_ns": "",
+                "checkpoint_id": AnyStr(),
+            }
+        },
+        checkpoint={
+            "v": 1,
+            "id": AnyStr(),
+            "ts": AnyStr(),
+            "pending_sends": [],
+            "versions_seen": {
+                "one": {
+                    "start:one": AnyVersion(),
+                },
+                "two": {
+                    "start:two": AnyVersion(),
+                },
+                "__input__": {},
+                "__start__": {
+                    "__start__": AnyVersion(),
+                },
+                "__interrupt__": {
+                    "value": AnyVersion(),
+                    "__start__": AnyVersion(),
+                    "start:one": AnyVersion(),
+                    "start:two": AnyVersion(),
+                },
+            },
+            "channel_versions": {
+                "one": AnyVersion(),
+                "two": AnyVersion(),
+                "value": AnyVersion(),
+                "__start__": AnyVersion(),
+                "start:one": AnyVersion(),
+                "start:two": AnyVersion(),
+            },
+            "channel_values": {"one": "one", "two": "two", "value": 6},
+        },
+        metadata={
+            "step": 1,
+            "source": "loop",
+            "writes": {"one": {"value": 2}, "two": {"value": 3}},
+        },
+        parent_config={
+            "configurable": {
+                "thread_id": "1",
+                "checkpoint_ns": "",
+                "checkpoint_id": checkpoints[1].config["configurable"]["checkpoint_id"],
+            }
+        },
+        pending_writes=[],
+    )
+    # the previous one we assert that pending writes contains both
+    # - original error
+    # - successful writes from resuming after preventing error
+    assert checkpoints[1] == CheckpointTuple(
+        config={
+            "configurable": {
+                "thread_id": "1",
+                "checkpoint_ns": "",
+                "checkpoint_id": AnyStr(),
+            }
+        },
+        checkpoint={
+            "v": 1,
+            "id": AnyStr(),
+            "ts": AnyStr(),
+            "pending_sends": [],
+            "versions_seen": {
+                "__input__": {},
+                "__start__": {
+                    "__start__": AnyVersion(),
+                },
+            },
+            "channel_versions": {
+                "value": AnyVersion(),
+                "__start__": AnyVersion(),
+                "start:one": AnyVersion(),
+                "start:two": AnyVersion(),
+            },
+            "channel_values": {
+                "value": 1,
+                "start:one": "__start__",
+                "start:two": "__start__",
+            },
+        },
+        metadata={"step": 0, "source": "loop", "writes": None},
+        parent_config={
+            "configurable": {
+                "thread_id": "1",
+                "checkpoint_ns": "",
+                "checkpoint_id": checkpoints[2].config["configurable"]["checkpoint_id"],
+            }
+        },
+        pending_writes=UnsortedSequence(
+            (AnyStr(), "one", "one"),
+            (AnyStr(), "value", 2),
+            (AnyStr(), "__error__", ExceptionLike(ConnectionError("I'm not good"))),
+            (AnyStr(), "two", "two"),
+            (AnyStr(), "value", 3),
+        ),
+    )
+    assert checkpoints[2] == CheckpointTuple(
+        config={
+            "configurable": {
+                "thread_id": "1",
+                "checkpoint_ns": "",
+                "checkpoint_id": AnyStr(),
+            }
+        },
+        checkpoint={
+            "v": 1,
+            "id": AnyStr(),
+            "ts": AnyStr(),
+            "pending_sends": [],
+            "versions_seen": {"__input__": {}},
+            "channel_versions": {
+                "__start__": AnyVersion(),
+            },
+            "channel_values": {"__start__": {"value": 1}},
+        },
+        metadata={"step": -1, "source": "input", "writes": {"value": 1}},
+        parent_config=None,
+        pending_writes=UnsortedSequence(
+            (AnyStr(), "value", 1),
+            (AnyStr(), "start:one", "__start__"),
+            (AnyStr(), "start:two", "__start__"),
+        ),
+    )
+
 
 def test_cond_edge_after_send() -> None:
     class Node:
@@ -1741,7 +1918,6 @@ def test_channel_enter_exit_timing(mocker: MockerFixture) -> None:
     assert cleanup.call_count == 0
     for i, chunk in enumerate(app.stream(2)):
         assert setup.call_count == 1, "Expected setup to be called once"
-        assert cleanup.call_count == 0, "Expected cleanup to not be called yet"
         if i == 0:
             assert chunk == {"inbox": [3]}
         elif i == 1:
@@ -1993,12 +2169,14 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None:
             "step": 0,
             "writes": {
                 "agent": {
-                    "input": "what is weather in sf",
-                    "agent_outcome": AgentAction(
-                        tool="search_api",
-                        tool_input="query",
-                        log="tool:search_api:query",
-                    ),
+                    "agent": {
+                        "input": "what is weather in sf",
+                        "agent_outcome": AgentAction(
+                            tool="search_api",
+                            tool_input="query",
+                            log="tool:search_api:query",
+                        ),
+                    }
                 },
             },
         },
@@ -2195,7 +2373,7 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None:
                 ),
             },
         },
-        tasks=(PregelTask(AnyStr(), "tools", interrupts=(Interrupt("before"),)),),
+        tasks=(PregelTask(AnyStr(), "tools"),),
         next=("tools",),
         config=app_w_interrupt.checkpointer.get_tuple(config).config,
         created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"],
@@ -2204,12 +2382,14 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None:
             "step": 0,
             "writes": {
                 "agent": {
-                    "input": "what is weather in sf",
-                    "agent_outcome": AgentAction(
-                        tool="search_api",
-                        tool_input="query",
-                        log="tool:search_api:query",
-                    ),
+                    "agent": {
+                        "input": "what is weather in sf",
+                        "agent_outcome": AgentAction(
+                            tool="search_api",
+                            tool_input="query",
+                            log="tool:search_api:query",
+                        ),
+                    }
                 }
             },
         },
@@ -2239,7 +2419,7 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None:
                 "input": "what is weather in sf",
             },
         },
-        tasks=(PregelTask(AnyStr(), "tools", interrupts=(Interrupt("before"),)),),
+        tasks=(PregelTask(AnyStr(), "tools"),),
         next=("tools",),
         config=app_w_interrupt.checkpointer.get_tuple(config).config,
         created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"],
@@ -2400,7 +2580,7 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None:
                 ),
             },
         },
-        tasks=(PregelTask(AnyStr(), "tools", interrupts=(Interrupt("before"),)),),
+        tasks=(PregelTask(AnyStr(), "tools"),),
         next=("tools",),
         config=app_w_interrupt.checkpointer.get_tuple(config).config,
         created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"],
@@ -2409,12 +2589,14 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None:
             "step": 0,
             "writes": {
                 "agent": {
-                    "input": "what is weather in sf",
-                    "agent_outcome": AgentAction(
-                        tool="search_api",
-                        tool_input="query",
-                        log="tool:search_api:query",
-                    ),
+                    "agent": {
+                        "input": "what is weather in sf",
+                        "agent_outcome": AgentAction(
+                            tool="search_api",
+                            tool_input="query",
+                            log="tool:search_api:query",
+                        ),
+                    }
                 }
             },
         },
@@ -2999,7 +3181,7 @@ def test_conditional_state_graph(
             ),
             "intermediate_steps": [],
         },
-        tasks=(PregelTask(AnyStr(), "tools", interrupts=(Interrupt("before"),)),),
+        tasks=(PregelTask(AnyStr(), "tools"),),
         next=("tools",),
         config=app_w_interrupt.checkpointer.get_tuple(config).config,
         created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"],
@@ -3039,7 +3221,7 @@ def test_conditional_state_graph(
             ),
             "intermediate_steps": [],
         },
-        tasks=(PregelTask(AnyStr(), "tools", interrupts=(Interrupt("before"),)),),
+        tasks=(PregelTask(AnyStr(), "tools"),),
         next=("tools",),
         config=app_w_interrupt.checkpointer.get_tuple(config).config,
         created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"],
@@ -3148,7 +3330,7 @@ def test_conditional_state_graph(
         values={
             "intermediate_steps": [],
         },
-        tasks=(PregelTask(AnyStr(), "agent", interrupts=(Interrupt("before"),)),),
+        tasks=(PregelTask(AnyStr(), "agent"),),
         next=("agent",),
         config=app_w_interrupt.checkpointer.get_tuple(config).config,
         created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"],
@@ -3173,7 +3355,7 @@ def test_conditional_state_graph(
             ),
             "intermediate_steps": [],
         },
-        tasks=(PregelTask(AnyStr(), "tools", interrupts=(Interrupt("before"),)),),
+        tasks=(PregelTask(AnyStr(), "tools"),),
         next=("tools",),
         config=app_w_interrupt.checkpointer.get_tuple(config).config,
         created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"],
@@ -3226,7 +3408,7 @@ def test_conditional_state_graph(
                 )
             ],
         },
-        tasks=(PregelTask(AnyStr(), "agent", interrupts=(Interrupt("before"),)),),
+        tasks=(PregelTask(AnyStr(), "agent"),),
         next=("agent",),
         config=app_w_interrupt.checkpointer.get_tuple(config).config,
         created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"],
@@ -4828,13 +5010,7 @@ def test_message_graph(
                 id="ai1",
             ),
         ],
-        tasks=(
-            PregelTask(
-                AnyStr(),
-                "tools",
-                interrupts=(Interrupt("before"),),
-            ),
-        ),
+        tasks=(PregelTask(AnyStr(), "tools"),),
         next=("tools",),
         config=app_w_interrupt.checkpointer.get_tuple(config).config,
         created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"],
@@ -4879,7 +5055,7 @@ def test_message_graph(
                 ],
             ),
         ],
-        tasks=(PregelTask(AnyStr(), "tools", interrupts=(Interrupt("before"),)),),
+        tasks=(PregelTask(AnyStr(), "tools"),),
         next=("tools",),
         config=app_w_interrupt.checkpointer.get_tuple(config).config,
         created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"],
@@ -4961,7 +5137,7 @@ def test_message_graph(
                 id="ai2",
             ),
         ],
-        tasks=(PregelTask(AnyStr(), "tools", interrupts=(Interrupt("before"),)),),
+        tasks=(PregelTask(AnyStr(), "tools"),),
         next=("tools",),
         config=app_w_interrupt.checkpointer.get_tuple(config).config,
         created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"],
@@ -5553,13 +5729,7 @@ def test_root_graph(
                 id="ai1",
             ),
         ],
-        tasks=(
-            PregelTask(
-                AnyStr(),
-                "tools",
-                interrupts=(Interrupt("before"),),
-            ),
-        ),
+        tasks=(PregelTask(AnyStr(), "tools"),),
         next=("tools",),
         config=app_w_interrupt.checkpointer.get_tuple(config).config,
         created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"],
@@ -5604,7 +5774,7 @@ def test_root_graph(
                 ],
             ),
         ],
-        tasks=(PregelTask(AnyStr(), "tools", interrupts=(Interrupt("before"),)),),
+        tasks=(PregelTask(AnyStr(), "tools"),),
         next=("tools",),
         config=app_w_interrupt.checkpointer.get_tuple(config).config,
         created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"],
@@ -5686,7 +5856,7 @@ def test_root_graph(
                 id="ai2",
             ),
         ],
-        tasks=(PregelTask(AnyStr(), "tools", interrupts=(Interrupt("before"),)),),
+        tasks=(PregelTask(AnyStr(), "tools"),),
         next=("tools",),
         config=app_w_interrupt.checkpointer.get_tuple(config).config,
         created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"],
@@ -6008,6 +6178,8 @@ def test_in_one_fan_out_out_one_graph_state() -> None:
                     "id": "592f3430-c17c-5d1c-831f-fecebb2c05bf",
                     "name": "rewrite_query",
                     "result": [("query", "query: what is weather in sf")],
+                    "error": None,
+                    "interrupts": [],
                 },
             },
         ),
@@ -6062,6 +6234,8 @@ def test_in_one_fan_out_out_one_graph_state() -> None:
                     "id": "96965ed0-2c10-52a1-86eb-081ba6de73b2",
                     "name": "retriever_two",
                     "result": [("docs", ["doc3", "doc4"])],
+                    "error": None,
+                    "interrupts": [],
                 },
             },
         ),
@@ -6079,6 +6253,8 @@ def test_in_one_fan_out_out_one_graph_state() -> None:
                     "id": "7db5e9d8-e132-5079-ab99-ced15e67d48b",
                     "name": "retriever_one",
                     "result": [("docs", ["doc1", "doc2"])],
+                    "error": None,
+                    "interrupts": [],
                 },
             },
         ),
@@ -6118,6 +6294,8 @@ def test_in_one_fan_out_out_one_graph_state() -> None:
                     "id": "8959fb57-d0f5-5725-9ac4-ec1c554fb0a0",
                     "name": "qa",
                     "result": [("answer", "doc1,doc2,doc3,doc4")],
+                    "error": None,
+                    "interrupts": [],
                 },
             },
         ),
@@ -6137,20 +6315,34 @@ def test_dynamic_interrupt(snapshot: SnapshotAssertion) -> None:
         my_key: Annotated[str, operator.add]
         market: str
 
+    tool_two_node_count = 0
+
     def tool_two_node(s: State) -> State:
+        nonlocal tool_two_node_count
+        tool_two_node_count += 1
         if s["market"] == "DE":
             raise NodeInterrupt("Just because...")
         return {"my_key": " all good"}
 
     tool_two_graph = StateGraph(State)
-    tool_two_graph.add_node("tool_two", tool_two_node)
+    tool_two_graph.add_node("tool_two", tool_two_node, retry=RetryPolicy())
     tool_two_graph.add_edge(START, "tool_two")
     tool_two = tool_two_graph.compile()
 
-    assert tool_two.invoke({"my_key": "value", "market": "DE"}) == {
+    tracer = FakeTracer()
+    assert tool_two.invoke(
+        {"my_key": "value", "market": "DE"}, {"callbacks": [tracer]}
+    ) == {
         "my_key": "value",
         "market": "DE",
     }
+    assert tool_two_node_count == 1, "interrupts aren't retried"
+    assert len(tracer.runs) == 1
+    run = tracer.runs[0]
+    assert run.end_time is not None
+    assert run.error is None
+    assert run.outputs == {"market": "DE", "my_key": "value"}
+
     assert tool_two.invoke({"my_key": "value", "market": "US"}) == {
         "my_key": "value all good",
         "market": "US",
@@ -6188,7 +6380,7 @@ def test_dynamic_interrupt(snapshot: SnapshotAssertion) -> None:
                 PregelTask(
                     AnyStr(),
                     "tool_two",
-                    interrupts=(Interrupt("during", "Just because..."),),
+                    interrupts=(Interrupt("Just because..."),),
                 ),
             ),
             config=tool_two.checkpointer.get_tuple(thread1).config,
@@ -6202,10 +6394,32 @@ def test_start_branch_then(snapshot: SnapshotAssertion) -> None:
     class State(TypedDict):
         my_key: Annotated[str, operator.add]
         market: str
+        shared: Annotated[dict[str, dict[str, Any]], SharedValue.on("assistant_id")]
+
+    def assert_shared_value(data: State, config: RunnableConfig) -> State:
+        assert "shared" in data
+        if thread_id := config["configurable"].get("thread_id"):
+            if thread_id == "1":
+                # this is the first thread, so should not see a value
+                assert data["shared"] == {}
+                return {"shared": {"1": {"hello": "world"}}}
+            elif thread_id == "2":
+                # this should get value saved by thread 1
+                assert data["shared"] == {"1": {"hello": "world"}}
+            elif thread_id == "3":
+                # this is a different assistant, so should not see previous value
+                assert data["shared"] == {}
+        return {}
+
+    def tool_two_slow(data: State, config: RunnableConfig) -> State:
+        return {"my_key": " slow", **assert_shared_value(data, config)}
+
+    def tool_two_fast(data: State, config: RunnableConfig) -> State:
+        return {"my_key": " fast", **assert_shared_value(data, config)}
 
     tool_two_graph = StateGraph(State)
-    tool_two_graph.add_node("tool_two_slow", lambda s: {"my_key": " slow"})
-    tool_two_graph.add_node("tool_two_fast", lambda s: {"my_key": " fast"})
+    tool_two_graph.add_node("tool_two_slow", tool_two_slow)
+    tool_two_graph.add_node("tool_two_fast", tool_two_fast)
     tool_two_graph.set_conditional_entry_point(
         lambda s: "tool_two_slow" if s["market"] == "DE" else "tool_two_fast", then=END
     )
@@ -6223,14 +6437,16 @@ def test_start_branch_then(snapshot: SnapshotAssertion) -> None:
 
     with SqliteSaver.from_conn_string(":memory:") as saver:
         tool_two = tool_two_graph.compile(
-            checkpointer=saver, interrupt_before=["tool_two_fast", "tool_two_slow"]
+            store=MemoryStore(),
+            checkpointer=saver,
+            interrupt_before=["tool_two_fast", "tool_two_slow"],
         )
 
         # missing thread_id
         with pytest.raises(ValueError, match="thread_id"):
             tool_two.invoke({"my_key": "value", "market": "DE"})
 
-        thread1 = {"configurable": {"thread_id": "1"}}
+        thread1 = {"configurable": {"thread_id": "1", "assistant_id": "a"}}
         # stop when about to enter node
         assert tool_two.invoke({"my_key": "value ⛰️", "market": "DE"}, thread1) == {
             "my_key": "value ⛰️",
@@ -6250,13 +6466,7 @@ def test_start_branch_then(snapshot: SnapshotAssertion) -> None:
         ]
         assert tool_two.get_state(thread1) == StateSnapshot(
             values={"my_key": "value ⛰️", "market": "DE"},
-            tasks=(
-                PregelTask(
-                    AnyStr(),
-                    "tool_two_slow",
-                    interrupts=(Interrupt("before"),),
-                ),
-            ),
+            tasks=(PregelTask(AnyStr(), "tool_two_slow"),),
             next=("tool_two_slow",),
             config=tool_two.checkpointer.get_tuple(thread1).config,
             created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"],
@@ -6282,7 +6492,7 @@ def test_start_branch_then(snapshot: SnapshotAssertion) -> None:
             parent_config=[*tool_two.checkpointer.list(thread1, limit=2)][-1].config,
         )
 
-        thread2 = {"configurable": {"thread_id": "2"}}
+        thread2 = {"configurable": {"thread_id": "2", "assistant_id": "a"}}
         # stop when about to enter node
         assert tool_two.invoke({"my_key": "value", "market": "US"}, thread2) == {
             "my_key": "value",
@@ -6290,13 +6500,7 @@ def test_start_branch_then(snapshot: SnapshotAssertion) -> None:
         }
         assert tool_two.get_state(thread2) == StateSnapshot(
             values={"my_key": "value", "market": "US"},
-            tasks=(
-                PregelTask(
-                    AnyStr(),
-                    "tool_two_fast",
-                    interrupts=(Interrupt("before"),),
-                ),
-            ),
+            tasks=(PregelTask(AnyStr(), "tool_two_fast"),),
             next=("tool_two_fast",),
             config=tool_two.checkpointer.get_tuple(thread2).config,
             created_at=tool_two.checkpointer.get_tuple(thread2).checkpoint["ts"],
@@ -6322,7 +6526,7 @@ def test_start_branch_then(snapshot: SnapshotAssertion) -> None:
             parent_config=[*tool_two.checkpointer.list(thread2, limit=2)][-1].config,
         )
 
-        thread3 = {"configurable": {"thread_id": "3"}}
+        thread3 = {"configurable": {"thread_id": "3", "assistant_id": "b"}}
         # stop when about to enter node
         assert tool_two.invoke({"my_key": "value", "market": "US"}, thread3) == {
             "my_key": "value",
@@ -6330,13 +6534,7 @@ def test_start_branch_then(snapshot: SnapshotAssertion) -> None:
         }
         assert tool_two.get_state(thread3) == StateSnapshot(
             values={"my_key": "value", "market": "US"},
-            tasks=(
-                PregelTask(
-                    AnyStr(),
-                    "tool_two_fast",
-                    interrupts=(Interrupt("before"),),
-                ),
-            ),
+            tasks=(PregelTask(AnyStr(), "tool_two_fast"),),
             next=("tool_two_fast",),
             config=tool_two.checkpointer.get_tuple(thread3).config,
             created_at=tool_two.checkpointer.get_tuple(thread3).checkpoint["ts"],
@@ -6347,13 +6545,7 @@ def test_start_branch_then(snapshot: SnapshotAssertion) -> None:
         tool_two.update_state(thread3, {"my_key": "key"})  # appends to my_key
         assert tool_two.get_state(thread3) == StateSnapshot(
             values={"my_key": "valuekey", "market": "US"},
-            tasks=(
-                PregelTask(
-                    AnyStr(),
-                    "tool_two_fast",
-                    interrupts=(Interrupt("before"),),
-                ),
-            ),
+            tasks=(PregelTask(AnyStr(), "tool_two_fast"),),
             next=("tool_two_fast",),
             config=tool_two.checkpointer.get_tuple(thread3).config,
             created_at=tool_two.checkpointer.get_tuple(thread3).checkpoint["ts"],
@@ -6497,6 +6689,8 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None:
                     "id": "7b7b0713-e958-5d07-803c-c9910a7cc162",
                     "name": "prepare",
                     "result": [("my_key", " prepared")],
+                    "error": None,
+                    "interrupts": [],
                 },
             },
             {
@@ -6549,6 +6743,8 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None:
                     "id": "dd9f2fa5-ccfa-5d12-81ec-942563056a08",
                     "name": "tool_two_slow",
                     "result": [("my_key", " slow")],
+                    "error": None,
+                    "interrupts": [],
                 },
             },
             {
@@ -6599,6 +6795,8 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None:
                     "id": "9b590c54-15ef-54b1-83a7-140d27b0bc52",
                     "name": "finish",
                     "result": [("my_key", " finished")],
+                    "error": None,
+                    "interrupts": [],
                 },
             },
             {
@@ -6648,13 +6846,7 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None:
         }
         assert tool_two.get_state(thread1) == StateSnapshot(
             values={"my_key": "value prepared", "market": "DE"},
-            tasks=(
-                PregelTask(
-                    AnyStr(),
-                    "tool_two_slow",
-                    interrupts=(Interrupt("before"),),
-                ),
-            ),
+            tasks=(PregelTask(AnyStr(), "tool_two_slow"),),
             next=("tool_two_slow",),
             config=tool_two.checkpointer.get_tuple(thread1).config,
             created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"],
@@ -6692,13 +6884,7 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None:
         }
         assert tool_two.get_state(thread2) == StateSnapshot(
             values={"my_key": "value prepared", "market": "US"},
-            tasks=(
-                PregelTask(
-                    AnyStr(),
-                    "tool_two_fast",
-                    interrupts=(Interrupt("before"),),
-                ),
-            ),
+            tasks=(PregelTask(AnyStr(), "tool_two_fast"),),
             next=("tool_two_fast",),
             config=tool_two.checkpointer.get_tuple(thread2).config,
             created_at=tool_two.checkpointer.get_tuple(thread2).checkpoint["ts"],
@@ -6745,7 +6931,7 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None:
                 "my_key": "value prepared slow",
                 "market": "DE",
             },
-            tasks=(PregelTask(AnyStr(), "finish", interrupts=(Interrupt("before"),)),),
+            tasks=(PregelTask(AnyStr(), "finish"),),
             next=("finish",),
             config=tool_two.checkpointer.get_tuple(thread1).config,
             created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"],
@@ -6764,7 +6950,7 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None:
                 "my_key": "value prepared slower",
                 "market": "DE",
             },
-            tasks=(PregelTask(AnyStr(), "finish", interrupts=(Interrupt("before"),)),),
+            tasks=(PregelTask(AnyStr(), "finish"),),
             next=("finish",),
             config=tool_two.checkpointer.get_tuple(thread1).config,
             created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"],
@@ -7021,7 +7207,7 @@ def test_in_one_fan_out_state_graph_waiting_edge(snapshot: SnapshotAssertion) ->
             "query": "analyzed: query: what is weather in sf",
             "docs": ["doc1", "doc2", "doc3", "doc4", "doc5"],
         },
-        tasks=(PregelTask(AnyStr(), "qa", interrupts=(Interrupt("before"),)),),
+        tasks=(PregelTask(AnyStr(), "qa"),),
         next=("qa",),
         config=app_w_interrupt.checkpointer.get_tuple(config).config,
         created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"],

libs/langgraph/tests/test_pregel_async.py

@@ -52,6 +52,7 @@ from langgraph.errors import InvalidUpdateError, NodeInterrupt
 from langgraph.graph import END, Graph, StateGraph
 from langgraph.graph.graph import START
 from langgraph.graph.message import MessageGraph, add_messages
+from langgraph.managed.shared_value import SharedValue
 from langgraph.prebuilt.chat_agent_executor import (
     create_tool_calling_executor,
 )
@@ -60,7 +61,9 @@ from langgraph.prebuilt.tool_node import ToolNode
 from langgraph.pregel import Channel, GraphRecursionError, Pregel, StateSnapshot
 from langgraph.pregel.retry import RetryPolicy
 from langgraph.pregel.types import PregelTask
-from tests.any_str import AnyStr, ExceptionLike
+from langgraph.store.memory import MemoryStore
+from tests.any_str import AnyStr, AnyVersion, ExceptionLike, UnsortedSequence
+from tests.fake_tracer import FakeTracer
 from tests.memory_assert import (
     MemorySaverAssertCheckpointMetadata,
     MemorySaverAssertImmutable,
@@ -204,6 +207,97 @@ async def test_node_cancellation_on_other_node_exception() -> None:
     assert inner_task_cancelled
 
 
+@pytest.mark.parametrize(
+    "checkpointer_name",
+    ["memory", "sqlite_aio", "postgres_aio", "postgres_aio_pipe"],
+)
+async def test_dynamic_interrupt(
+    checkpointer_name: str, snapshot: SnapshotAssertion, request: pytest.FixtureRequest
+) -> None:
+    checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}")
+
+    class State(TypedDict):
+        my_key: Annotated[str, operator.add]
+        market: str
+
+    tool_two_node_count = 0
+
+    async def tool_two_node(s: State) -> State:
+        nonlocal tool_two_node_count
+        tool_two_node_count += 1
+        if s["market"] == "DE":
+            raise NodeInterrupt("Just because...")
+        return {"my_key": " all good"}
+
+    tool_two_graph = StateGraph(State)
+    tool_two_graph.add_node("tool_two", tool_two_node, retry=RetryPolicy())
+    tool_two_graph.add_edge(START, "tool_two")
+    tool_two = tool_two_graph.compile()
+
+    tracer = FakeTracer()
+    assert await tool_two.ainvoke(
+        {"my_key": "value", "market": "DE"}, {"callbacks": [tracer]}
+    ) == {
+        "my_key": "value",
+        "market": "DE",
+    }
+    assert tool_two_node_count == 1, "interrupts aren't retried"
+    assert len(tracer.runs) == 1
+    run = tracer.runs[0]
+    assert run.end_time is not None
+    assert run.error is None
+    assert run.outputs == {"market": "DE", "my_key": "value"}
+
+    assert await tool_two.ainvoke({"my_key": "value", "market": "US"}) == {
+        "my_key": "value all good",
+        "market": "US",
+    }
+
+    tool_two = tool_two_graph.compile(checkpointer=checkpointer)
+
+    # missing thread_id
+    with pytest.raises(ValueError, match="thread_id"):
+        await tool_two.ainvoke({"my_key": "value", "market": "DE"})
+
+    thread1 = {"configurable": {"thread_id": "1"}}
+    # stop when about to enter node
+    assert await tool_two.ainvoke({"my_key": "value ⛰️", "market": "DE"}, thread1) == {
+        "my_key": "value ⛰️",
+        "market": "DE",
+    }
+    assert [c.metadata async for c in tool_two.checkpointer.alist(thread1)] == [
+        {
+            "source": "loop",
+            "step": 0,
+            "writes": None,
+        },
+        {
+            "source": "input",
+            "step": -1,
+            "writes": {"my_key": "value ⛰️", "market": "DE"},
+        },
+    ]
+    tup = await tool_two.checkpointer.aget_tuple(thread1)
+    assert await tool_two.aget_state(thread1) == StateSnapshot(
+        values={"my_key": "value ⛰️", "market": "DE"},
+        next=("tool_two",),
+        tasks=(
+            PregelTask(
+                AnyStr(),
+                "tool_two",
+                interrupts=(Interrupt("Just because..."),),
+            ),
+        ),
+        config=tup.config,
+        created_at=tup.checkpoint["ts"],
+        metadata={"source": "loop", "step": 0, "writes": None},
+        parent_config=[c async for c in tool_two.checkpointer.alist(thread1, limit=2)][
+            -1
+        ].config,
+    )
+    # TODO use aget_state_history
+
+
 @pytest.mark.parametrize(
     "checkpointer_name",
     ["memory", "sqlite_aio", "postgres_aio", "postgres_aio_pipe"],
@@ -798,7 +892,7 @@ async def test_invoke_two_processes_in_out_interrupt(
                     "checkpoint_id": AnyStr(),
                 }
             },
-            metadata={"source": "loop", "step": 6, "writes": 5},
+            metadata={"source": "loop", "step": 6, "writes": {"two": 5}},
             created_at=AnyStr(),
             parent_config=history[1].config,
         ),
@@ -813,7 +907,7 @@ async def test_invoke_two_processes_in_out_interrupt(
                     "checkpoint_id": AnyStr(),
                 }
             },
-            metadata={"source": "loop", "step": 5, "writes": None},
+            metadata={"source": "loop", "step": 5, "writes": {"one": None}},
             created_at=AnyStr(),
             parent_config=history[2].config,
         ),
@@ -843,7 +937,7 @@ async def test_invoke_two_processes_in_out_interrupt(
                     "checkpoint_id": AnyStr(),
                 }
             },
-            metadata={"source": "loop", "step": 3, "writes": None},
+            metadata={"source": "loop", "step": 3, "writes": {"one": None}},
             created_at=AnyStr(),
             parent_config=history[4].config,
         ),
@@ -873,7 +967,7 @@ async def test_invoke_two_processes_in_out_interrupt(
                     "checkpoint_id": AnyStr(),
                 }
             },
-            metadata={"source": "loop", "step": 1, "writes": 4},
+            metadata={"source": "loop", "step": 1, "writes": {"two": 4}},
             created_at=AnyStr(),
             parent_config=history[6].config,
         ),
@@ -888,7 +982,7 @@ async def test_invoke_two_processes_in_out_interrupt(
                     "checkpoint_id": AnyStr(),
                 }
             },
-            metadata={"source": "loop", "step": 0, "writes": None},
+            metadata={"source": "loop", "step": 0, "writes": {"one": None}},
             created_at=AnyStr(),
             parent_config=history[7].config,
         ),
@@ -1190,6 +1284,8 @@ async def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None:
                 "id": "2687f72c-e3a8-5f6f-9afa-047cbf24e923",
                 "name": "one",
                 "result": [("inbox", 3)],
+                "error": None,
+                "interrupts": [],
             },
         },
         {
@@ -1200,6 +1296,8 @@ async def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None:
                 "id": "18f52f6a-828d-58a1-a501-53cc0c7af33e",
                 "name": "two",
                 "result": [("output", 13)],
+                "error": None,
+                "interrupts": [],
             },
         },
         {
@@ -1221,6 +1319,8 @@ async def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None:
                 "id": "871d6e74-7bb3-565f-a4fe-cef4b8f19b62",
                 "name": "two",
                 "result": [("output", 4)],
+                "error": None,
+                "interrupts": [],
             },
         },
     ]
@@ -1515,6 +1615,10 @@ async def test_pending_writes_resume(
     error_write = next(w for w in checkpoint.pending_writes if w[1] == ERROR)
     assert error_write[0] != non_error_writes[0][0]
 
+    # TODO arguably this shouldn't even run the failed task again,
+    # and should require empty update_state (ie new checkpoint_id)
+    # in order to try again
+
     # resume execution
     with pytest.raises(ValueError, match="I'm not good"):
         await graph.ainvoke(None, thread1)
@@ -1533,6 +1637,144 @@ async def test_pending_writes_resume(
     # both the pending write and the new write were applied, 1 + 2 + 3 = 6
     assert await graph.ainvoke(None, thread1) == {"value": 6}
 
+    # check all final checkpoints
+    checkpoints = [c async for c in checkpointer.alist(thread1)]
+    # we should have 3
+    assert len(checkpoints) == 3
+    # the last one not too interesting for this test
+    assert checkpoints[0] == CheckpointTuple(
+        config={
+            "configurable": {
+                "thread_id": "1",
+                "checkpoint_ns": "",
+                "checkpoint_id": AnyStr(),
+            }
+        },
+        checkpoint={
+            "v": 1,
+            "id": AnyStr(),
+            "ts": AnyStr(),
+            "pending_sends": [],
+            "versions_seen": {
+                "one": {
+                    "start:one": AnyVersion(),
+                },
+                "two": {
+                    "start:two": AnyVersion(),
+                },
+                "__input__": {},
+                "__start__": {
+                    "__start__": AnyVersion(),
+                },
+                "__interrupt__": {
+                    "value": AnyVersion(),
+                    "__start__": AnyVersion(),
+                    "start:one": AnyVersion(),
+                    "start:two": AnyVersion(),
+                },
+            },
+            "channel_versions": {
+                "one": AnyVersion(),
+                "two": AnyVersion(),
+                "value": AnyVersion(),
+                "__start__": AnyVersion(),
+                "start:one": AnyVersion(),
+                "start:two": AnyVersion(),
+            },
+            "channel_values": {"one": "one", "two": "two", "value": 6},
+        },
+        metadata={
+            "step": 1,
+            "source": "loop",
+            "writes": {"one": {"value": 2}, "two": {"value": 3}},
+        },
+        parent_config={
+            "configurable": {
+                "thread_id": "1",
+                "checkpoint_ns": "",
+                "checkpoint_id": checkpoints[1].config["configurable"]["checkpoint_id"],
+            }
+        },
+        pending_writes=[],
+    )
+    # the previous one we assert that pending writes contains both
+    # - original error
+    # - successful writes from resuming after preventing error
+    assert checkpoints[1] == CheckpointTuple(
+        config={
+            "configurable": {
+                "thread_id": "1",
+                "checkpoint_ns": "",
+                "checkpoint_id": AnyStr(),
+            }
+        },
+        checkpoint={
+            "v": 1,
+            "id": AnyStr(),
+            "ts": AnyStr(),
+            "pending_sends": [],
+            "versions_seen": {
+                "__input__": {},
+                "__start__": {
+                    "__start__": AnyVersion(),
+                },
+            },
+            "channel_versions": {
+                "value": AnyVersion(),
+                "__start__": AnyVersion(),
+                "start:one": AnyVersion(),
+                "start:two": AnyVersion(),
+            },
+            "channel_values": {
+                "value": 1,
+                "start:one": "__start__",
+                "start:two": "__start__",
+            },
+        },
+        metadata={"step": 0, "source": "loop", "writes": None},
+        parent_config={
+            "configurable": {
+                "thread_id": "1",
+                "checkpoint_ns": "",
+                "checkpoint_id": checkpoints[2].config["configurable"]["checkpoint_id"],
+            }
+        },
+        pending_writes=UnsortedSequence(
+            (AnyStr(), "one", "one"),
+            (AnyStr(), "value", 2),
+            (AnyStr(), "__error__", ExceptionLike(ValueError("I'm not good"))),
+            (AnyStr(), "two", "two"),
+            (AnyStr(), "value", 3),
+        ),
+    )
+    assert checkpoints[2] == CheckpointTuple(
+        config={
+            "configurable": {
+                "thread_id": "1",
+                "checkpoint_ns": "",
+                "checkpoint_id": AnyStr(),
+            }
+        },
+        checkpoint={
+            "v": 1,
+            "id": AnyStr(),
+            "ts": AnyStr(),
+            "pending_sends": [],
+            "versions_seen": {"__input__": {}},
+            "channel_versions": {
+                "__start__": AnyVersion(),
+            },
+            "channel_values": {"__start__": {"value": 1}},
+        },
+        metadata={"step": -1, "source": "input", "writes": {"value": 1}},
+        parent_config=None,
+        pending_writes=UnsortedSequence(
+            (AnyStr(), "value", 1),
+            (AnyStr(), "start:one", "__start__"),
+            (AnyStr(), "start:two", "__start__"),
+        ),
+    )
+
 
 async def test_cond_edge_after_send() -> None:
     class Node:
@@ -1890,7 +2132,6 @@ async def test_channel_enter_exit_timing(mocker: MockerFixture) -> None:
         assert setup_sync.call_count == 0, "Sync context manager should not be used"
         assert cleanup_sync.call_count == 0, "Sync context manager should not be used"
         assert setup_async.call_count == 1, "Expected setup to be called once"
-        assert cleanup_async.call_count == 0, "Expected cleanup to not be called yet"
         if i == 0:
             assert chunk == {"inbox": [3]}
         elif i == 1:
@@ -2212,12 +2453,14 @@ async def test_conditional_graph() -> None:
             "step": 0,
             "writes": {
                 "agent": {
-                    "input": "what is weather in sf",
-                    "agent_outcome": AgentAction(
-                        tool="search_api",
-                        tool_input="query",
-                        log="tool:search_api:query",
-                    ),
+                    "agent": {
+                        "input": "what is weather in sf",
+                        "agent_outcome": AgentAction(
+                            tool="search_api",
+                            tool_input="query",
+                            log="tool:search_api:query",
+                        ),
+                    }
                 }
             },
         },
@@ -2421,13 +2664,7 @@ async def test_conditional_graph() -> None:
                 ),
             },
         },
-        tasks=(
-            PregelTask(
-                AnyStr(),
-                "tools",
-                interrupts=(Interrupt("before"),),
-            ),
-        ),
+        tasks=(PregelTask(AnyStr(), "tools"),),
         next=("tools",),
         config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config,
         created_at=(await app_w_interrupt.checkpointer.aget_tuple(config)).checkpoint[
@@ -2438,12 +2675,14 @@ async def test_conditional_graph() -> None:
             "step": 0,
             "writes": {
                 "agent": {
-                    "input": "what is weather in sf",
-                    "agent_outcome": AgentAction(
-                        tool="search_api",
-                        tool_input="query",
-                        log="tool:search_api:query",
-                    ),
+                    "agent": {
+                        "input": "what is weather in sf",
+                        "agent_outcome": AgentAction(
+                            tool="search_api",
+                            tool_input="query",
+                            log="tool:search_api:query",
+                        ),
+                    }
                 }
             },
         },
@@ -2475,13 +2714,7 @@ async def test_conditional_graph() -> None:
                 "input": "what is weather in sf",
             },
         },
-        tasks=(
-            PregelTask(
-                AnyStr(),
-                "tools",
-                interrupts=(Interrupt("before"),),
-            ),
-        ),
+        tasks=(PregelTask(AnyStr(), "tools"),),
         next=("tools",),
         config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config,
         created_at=(await app_w_interrupt.checkpointer.aget_tuple(config)).checkpoint[
@@ -2653,13 +2886,7 @@ async def test_conditional_graph() -> None:
                 ),
             },
         },
-        tasks=(
-            PregelTask(
-                AnyStr(),
-                "tools",
-                interrupts=(Interrupt("before"),),
-            ),
-        ),
+        tasks=(PregelTask(AnyStr(), "tools"),),
         next=("tools",),
         config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config,
         created_at=(await app_w_interrupt.checkpointer.aget_tuple(config)).checkpoint[
@@ -2670,12 +2897,14 @@ async def test_conditional_graph() -> None:
             "step": 0,
             "writes": {
                 "agent": {
-                    "input": "what is weather in sf",
-                    "agent_outcome": AgentAction(
-                        tool="search_api",
-                        tool_input="query",
-                        log="tool:search_api:query",
-                    ),
+                    "agent": {
+                        "input": "what is weather in sf",
+                        "agent_outcome": AgentAction(
+                            tool="search_api",
+                            tool_input="query",
+                            log="tool:search_api:query",
+                        ),
+                    }
                 }
             },
         },
@@ -3225,13 +3454,7 @@ async def test_conditional_graph_state(mocker: MockerFixture) -> None:
             ),
             "intermediate_steps": [],
         },
-        tasks=(
-            PregelTask(
-                AnyStr(),
-                "tools",
-                interrupts=(Interrupt("before"),),
-            ),
-        ),
+        tasks=(PregelTask(AnyStr(), "tools"),),
         next=("tools",),
         config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config,
         created_at=(await app_w_interrupt.checkpointer.aget_tuple(config)).checkpoint[
@@ -3275,7 +3498,7 @@ async def test_conditional_graph_state(mocker: MockerFixture) -> None:
             ),
             "intermediate_steps": [],
         },
-        tasks=(PregelTask(AnyStr(), "tools", interrupts=(Interrupt("before"),)),),
+        tasks=(PregelTask(AnyStr(), "tools"),),
         next=("tools",),
         config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config,
         created_at=(await app_w_interrupt.checkpointer.aget_tuple(config)).checkpoint[
@@ -4650,6 +4873,8 @@ async def test_in_one_fan_out_out_one_graph_state() -> None:
                     "id": "592f3430-c17c-5d1c-831f-fecebb2c05bf",
                     "name": "rewrite_query",
                     "result": [("query", "query: what is weather in sf")],
+                    "error": None,
+                    "interrupts": [],
                 },
             },
         ),
@@ -4704,6 +4929,8 @@ async def test_in_one_fan_out_out_one_graph_state() -> None:
                     "id": "96965ed0-2c10-52a1-86eb-081ba6de73b2",
                     "name": "retriever_two",
                     "result": [("docs", ["doc3", "doc4"])],
+                    "error": None,
+                    "interrupts": [],
                 },
             },
         ),
@@ -4721,6 +4948,8 @@ async def test_in_one_fan_out_out_one_graph_state() -> None:
                     "id": "7db5e9d8-e132-5079-ab99-ced15e67d48b",
                     "name": "retriever_one",
                     "result": [("docs", ["doc1", "doc2"])],
+                    "error": None,
+                    "interrupts": [],
                 },
             },
         ),
@@ -4760,6 +4989,8 @@ async def test_in_one_fan_out_out_one_graph_state() -> None:
                     "id": "8959fb57-d0f5-5725-9ac4-ec1c554fb0a0",
                     "name": "qa",
                     "result": [("answer", "doc1,doc2,doc3,doc4")],
+                    "error": None,
+                    "interrupts": [],
                 },
             },
         ),
@@ -4778,10 +5009,33 @@ async def test_start_branch_then() -> None:
     class State(TypedDict):
         my_key: Annotated[str, operator.add]
         market: str
+        shared: Annotated[dict[str, dict[str, Any]], SharedValue.on("assistant_id")]
+        other: Annotated[dict[str, dict[str, Any]], SharedValue.on("assistant_id")]
+
+    def assert_shared_value(data: State, config: RunnableConfig) -> State:
+        assert "shared" in data
+        if thread_id := config["configurable"].get("thread_id"):
+            if thread_id == "1":
+                # this is the first thread, so should not see a value
+                assert data["shared"] == {}
+                return {"shared": {"1": {"hello": "world"}}, "other": {"2": {1: 2}}}
+            elif thread_id == "2":
+                # this should get value saved by thread 1
+                assert data["shared"] == {"1": {"hello": "world"}}
+            elif thread_id == "3":
+                # this is a different assistant, so should not see previous value
+                assert data["shared"] == {}
+        return {}
+
+    def tool_two_slow(data: State, config: RunnableConfig) -> State:
+        return {"my_key": " slow", **assert_shared_value(data, config)}
+
+    def tool_two_fast(data: State, config: RunnableConfig) -> State:
+        return {"my_key": " fast", **assert_shared_value(data, config)}
 
     tool_two_graph = StateGraph(State)
-    tool_two_graph.add_node("tool_two_slow", lambda s, config: {"my_key": " slow"})
-    tool_two_graph.add_node("tool_two_fast", lambda s: {"my_key": " fast"})
+    tool_two_graph.add_node("tool_two_slow", tool_two_slow)
+    tool_two_graph.add_node("tool_two_fast", tool_two_fast)
     tool_two_graph.set_conditional_entry_point(
         lambda s: "tool_two_slow" if s["market"] == "DE" else "tool_two_fast", then=END
     )
@@ -4798,14 +5052,16 @@ async def test_start_branch_then() -> None:
 
     async with AsyncSqliteSaver.from_conn_string(":memory:") as saver:
         tool_two = tool_two_graph.compile(
-            checkpointer=saver, interrupt_before=["tool_two_fast", "tool_two_slow"]
+            store=MemoryStore(),
+            checkpointer=saver,
+            interrupt_before=["tool_two_fast", "tool_two_slow"],
         )
 
         # missing thread_id
         with pytest.raises(ValueError, match="thread_id"):
             await tool_two.ainvoke({"my_key": "value", "market": "DE"})
 
-        thread1 = {"configurable": {"thread_id": "1"}}
+        thread1 = {"configurable": {"thread_id": "1", "assistant_id": "a"}}
         # stop when about to enter node
         assert await tool_two.ainvoke({"my_key": "value", "market": "DE"}, thread1) == {
             "my_key": "value",
@@ -4825,13 +5081,7 @@ async def test_start_branch_then() -> None:
         ]
         assert await tool_two.aget_state(thread1) == StateSnapshot(
             values={"my_key": "value", "market": "DE"},
-            tasks=(
-                PregelTask(
-                    AnyStr(),
-                    "tool_two_slow",
-                    interrupts=(Interrupt("before"),),
-                ),
-            ),
+            tasks=(PregelTask(AnyStr(), "tool_two_slow"),),
             next=("tool_two_slow",),
             config=(await tool_two.checkpointer.aget_tuple(thread1)).config,
             created_at=(await tool_two.checkpointer.aget_tuple(thread1)).checkpoint[
@@ -4865,7 +5115,7 @@ async def test_start_branch_then() -> None:
             ][-1].config,
         )
 
-        thread2 = {"configurable": {"thread_id": "2"}}
+        thread2 = {"configurable": {"thread_id": "2", "assistant_id": "a"}}
         # stop when about to enter node
         assert await tool_two.ainvoke({"my_key": "value", "market": "US"}, thread2) == {
             "my_key": "value",
@@ -4873,13 +5123,7 @@ async def test_start_branch_then() -> None:
         }
         assert await tool_two.aget_state(thread2) == StateSnapshot(
             values={"my_key": "value", "market": "US"},
-            tasks=(
-                PregelTask(
-                    AnyStr(),
-                    "tool_two_fast",
-                    interrupts=(Interrupt("before"),),
-                ),
-            ),
+            tasks=(PregelTask(AnyStr(), "tool_two_fast"),),
             next=("tool_two_fast",),
             config=(await tool_two.checkpointer.aget_tuple(thread2)).config,
             created_at=(await tool_two.checkpointer.aget_tuple(thread2)).checkpoint[
@@ -4913,7 +5157,7 @@ async def test_start_branch_then() -> None:
             ][-1].config,
         )
 
-        thread3 = {"configurable": {"thread_id": "3"}}
+        thread3 = {"configurable": {"thread_id": "3", "assistant_id": "b"}}
         # stop when about to enter node
         assert await tool_two.ainvoke({"my_key": "value", "market": "US"}, thread3) == {
             "my_key": "value",
@@ -4921,13 +5165,7 @@ async def test_start_branch_then() -> None:
         }
         assert await tool_two.aget_state(thread3) == StateSnapshot(
             values={"my_key": "value", "market": "US"},
-            tasks=(
-                PregelTask(
-                    AnyStr(),
-                    "tool_two_fast",
-                    interrupts=(Interrupt("before"),),
-                ),
-            ),
+            tasks=(PregelTask(AnyStr(), "tool_two_fast"),),
             next=("tool_two_fast",),
             config=(await tool_two.checkpointer.aget_tuple(thread3)).config,
             created_at=(await tool_two.checkpointer.aget_tuple(thread3)).checkpoint[
@@ -4942,13 +5180,7 @@ async def test_start_branch_then() -> None:
         await tool_two.aupdate_state(thread3, {"my_key": "key"})  # appends to my_key
         assert await tool_two.aget_state(thread3) == StateSnapshot(
             values={"my_key": "valuekey", "market": "US"},
-            tasks=(
-                PregelTask(
-                    AnyStr(),
-                    "tool_two_fast",
-                    interrupts=(Interrupt("before"),),
-                ),
-            ),
+            tasks=(PregelTask(AnyStr(), "tool_two_fast"),),
             next=("tool_two_fast",),
             config=(await tool_two.checkpointer.aget_tuple(thread3)).config,
             created_at=(await tool_two.checkpointer.aget_tuple(thread3)).checkpoint[
@@ -5099,6 +5331,8 @@ async def test_branch_then() -> None:
                     "id": "7b7b0713-e958-5d07-803c-c9910a7cc162",
                     "name": "prepare",
                     "result": [("my_key", " prepared")],
+                    "error": None,
+                    "interrupts": [],
                 },
             },
             {
@@ -5151,6 +5385,8 @@ async def test_branch_then() -> None:
                     "id": "dd9f2fa5-ccfa-5d12-81ec-942563056a08",
                     "name": "tool_two_slow",
                     "result": [("my_key", " slow")],
+                    "error": None,
+                    "interrupts": [],
                 },
             },
             {
@@ -5201,6 +5437,8 @@ async def test_branch_then() -> None:
                     "id": "9b590c54-15ef-54b1-83a7-140d27b0bc52",
                     "name": "finish",
                     "result": [("my_key", " finished")],
+                    "error": None,
+                    "interrupts": [],
                 },
             },
             {
@@ -5244,19 +5482,125 @@ async def test_branch_then() -> None:
 
         thread1 = {"configurable": {"thread_id": "1"}}
         # stop when about to enter node
-        assert await tool_two.ainvoke({"my_key": "value", "market": "DE"}, thread1) == {
-            "my_key": "value prepared",
-            "market": "DE",
-        }
+        assert [
+            c
+            async for c in tool_two.astream(
+                {"my_key": "value", "market": "DE"}, thread1, stream_mode="debug"
+            )
+        ] == [
+            {
+                "type": "checkpoint",
+                "timestamp": AnyStr(),
+                "step": -1,
+                "payload": {
+                    "config": {
+                        "tags": [],
+                        "metadata": {"thread_id": "1"},
+                        "callbacks": None,
+                        "recursion_limit": 25,
+                        "configurable": {
+                            "thread_id": "1",
+                            "checkpoint_ns": "",
+                            "checkpoint_id": AnyStr(),
+                        },
+                    },
+                    "values": {"my_key": ""},
+                    "metadata": {
+                        "source": "input",
+                        "step": -1,
+                        "writes": {"my_key": "value", "market": "DE"},
+                    },
+                    "next": ["__start__"],
+                    "tasks": [{"id": AnyStr(), "name": "__start__", "interrupts": ()}],
+                },
+            },
+            {
+                "type": "checkpoint",
+                "timestamp": AnyStr(),
+                "step": 0,
+                "payload": {
+                    "config": {
+                        "tags": [],
+                        "metadata": {"thread_id": "1"},
+                        "callbacks": None,
+                        "recursion_limit": 25,
+                        "configurable": {
+                            "thread_id": "1",
+                            "checkpoint_ns": "",
+                            "checkpoint_id": AnyStr(),
+                        },
+                    },
+                    "values": {
+                        "my_key": "value",
+                        "market": "DE",
+                    },
+                    "metadata": {
+                        "source": "loop",
+                        "step": 0,
+                        "writes": None,
+                    },
+                    "next": ["prepare"],
+                    "tasks": [{"id": AnyStr(), "name": "prepare", "interrupts": ()}],
+                },
+            },
+            {
+                "type": "task",
+                "timestamp": AnyStr(),
+                "step": 1,
+                "payload": {
+                    "id": "ca572c3b-b805-5fc6-a19e-3d79f52dde70",
+                    "name": "prepare",
+                    "input": {"my_key": "value", "market": "DE"},
+                    "triggers": ["start:prepare"],
+                },
+            },
+            {
+                "type": "task_result",
+                "timestamp": AnyStr(),
+                "step": 1,
+                "payload": {
+                    "id": "ca572c3b-b805-5fc6-a19e-3d79f52dde70",
+                    "name": "prepare",
+                    "result": [("my_key", " prepared")],
+                    "error": None,
+                    "interrupts": [],
+                },
+            },
+            {
+                "type": "checkpoint",
+                "timestamp": AnyStr(),
+                "step": 1,
+                "payload": {
+                    "config": {
+                        "tags": [],
+                        "metadata": {"thread_id": "1"},
+                        "callbacks": None,
+                        "recursion_limit": 25,
+                        "configurable": {
+                            "thread_id": "1",
+                            "checkpoint_ns": "",
+                            "checkpoint_id": AnyStr(),
+                        },
+                    },
+                    "values": {
+                        "my_key": "value prepared",
+                        "market": "DE",
+                    },
+                    "metadata": {
+                        "source": "loop",
+                        "step": 1,
+                        "writes": {"prepare": {"my_key": " prepared"}},
+                    },
+                    "next": ["tool_two_slow"],
+                    "tasks": [
+                        {"id": AnyStr(), "name": "tool_two_slow", "interrupts": ()}
+                    ],
+                },
+            },
+        ]
         assert await tool_two.aget_state(thread1) == StateSnapshot(
             values={"my_key": "value prepared", "market": "DE"},
-            tasks=(
-                PregelTask(
-                    AnyStr(),
-                    "tool_two_slow",
-                    interrupts=(Interrupt("before"),),
-                ),
-            ),
+            tasks=(PregelTask(AnyStr(), "tool_two_slow"),),
             next=("tool_two_slow",),
             config=(await tool_two.checkpointer.aget_tuple(thread1)).config,
             created_at=(await tool_two.checkpointer.aget_tuple(thread1)).checkpoint[
@@ -5302,13 +5646,7 @@ async def test_branch_then() -> None:
         }
         assert await tool_two.aget_state(thread2) == StateSnapshot(
             values={"my_key": "value prepared", "market": "US"},
-            tasks=(
-                PregelTask(
-                    AnyStr(),
-                    "tool_two_fast",
-                    interrupts=(Interrupt("before"),),
-                ),
-            ),
+            tasks=(PregelTask(AnyStr(), "tool_two_fast"),),
             next=("tool_two_fast",),
             config=(await tool_two.checkpointer.aget_tuple(thread2)).config,
             created_at=(await tool_two.checkpointer.aget_tuple(thread2)).checkpoint[

libs/langgraph/tests/test_state.py

@@ -2,10 +2,11 @@ from typing import Annotated as Annotated2
 from typing import Any
 
 import pytest
+from langchain_core.runnables import RunnableConfig
 from pydantic.v1 import BaseModel
 from typing_extensions import Annotated, TypedDict
 
-from langgraph.graph.state import _warn_invalid_state_schema
+from langgraph.graph.state import StateGraph, _warn_invalid_state_schema
 
 
 class State(BaseModel):
@@ -46,3 +47,42 @@ def test_doesnt_warn_valid_schema(schema: Any):
     # Assert the function does not raise a warning
     with pytest.warns(None):
         _warn_invalid_state_schema(schema)
+
+
+def test_state_schema_with_type_hint():
+    class InputState(TypedDict):
+        question: str
+
+    class OutputState(TypedDict):
+        input_state: InputState
+
+    def complete_hint(state: InputState) -> OutputState:
+        return {"input_state": state}
+
+    def miss_first_hint(state, config: RunnableConfig) -> OutputState:
+        return {"input_state": state}
+
+    def only_return_hint(state, config) -> OutputState:
+        return {"input_state": state}
+
+    def miss_all_hint(state, config):
+        return {"input_state": state}
+
+    graph = StateGraph(input=InputState, output=OutputState)
+    actions = [complete_hint, miss_first_hint, only_return_hint, miss_all_hint]
+
+    for action in actions:
+        graph.add_node(action)
+
+    graph.set_entry_point(actions[0].__name__)
+    for i in range(len(actions) - 1):
+        graph.add_edge(actions[i].__name__, actions[i + 1].__name__)
+    graph.set_finish_point(actions[-1].__name__)
+
+    graph = graph.compile()
+
+    input_state = InputState(question="Hello World!")
+    output_state = OutputState(input_state=input_state)
+    for i, c in enumerate(graph.stream(input_state, stream_mode="updates")):
+        node_name = actions[i].__name__
+        assert c[node_name] == output_state

libs/langgraph/tests/test_store.py

@@ -0,0 +1,35 @@
+import asyncio
+from typing import Any, Optional
+
+from pytest_mock import MockerFixture
+
+from langgraph.store.base import BaseStore
+from langgraph.store.batch import AsyncBatchedStore
+
+
+async def test_async_batch_store(mocker: MockerFixture) -> None:
+    aget = mocker.stub()
+    alist = mocker.stub()
+
+    class MockStore(BaseStore):
+        async def aget(
+            self, pairs: list[tuple[str, str]]
+        ) -> dict[tuple[str, str], Optional[dict[str, Any]]]:
+            aget(pairs)
+            return {pair: 1 for pair in pairs}
+
+        async def alist(self, prefixes: list[str]) -> dict[str, dict[str, Any]]:
+            alist(prefixes)
+            return {prefix: {prefix: 1} for prefix in prefixes}
+
+    store = AsyncBatchedStore(MockStore())
+
+    # concurrent calls are batched
+    results = await asyncio.gather(
+        store.alist(["a", "b"]),
+        store.alist(["c", "d"]),
+    )
+    assert results == [{"a": {"a": 1}, "b": {"b": 1}}, {"c": {"c": 1}, "d": {"d": 1}}]
+    assert [c.args for c in alist.call_args_list] == [
+        (["a", "b", "c", "d"],),
+    ]

libs/sdk-py/langgraph_sdk/client.py

@@ -40,8 +40,14 @@ from langgraph_sdk.schema import (
 logger = logging.getLogger(__name__)
 
 
+RESERVED_HEADERS = ("x-api-key",)
+
+
 def get_client(
-    *, url: Optional[str] = None, api_key: Optional[str] = None
+    *,
+    url: Optional[str] = None,
+    api_key: Optional[str] = None,
+    headers: Optional[dict[str, str]] = None,
 ) -> LangGraphClient:
     """Get a LangGraphClient instance.
 
@@ -53,6 +59,7 @@ def get_client(
                 2. LANGGRAPH_API_KEY
                 3. LANGSMITH_API_KEY
                 4. LANGCHAIN_API_KEY
+        headers: Optional custom headers
     """
     transport: Optional[httpx.AsyncBaseTransport] = None
     if url is None:
@@ -65,17 +72,12 @@ def get_client(
             url = "http://localhost:8123"
     if transport is None:
         transport = httpx.AsyncHTTPTransport(retries=5)
-    headers = {
-        "User-Agent": f"langgraph-sdk-py/{langgraph_sdk.__version__}",
-    }
-    api_key = _get_api_key(api_key)
-    if api_key:
-        headers["x-api-key"] = api_key
+
     client = httpx.AsyncClient(
         base_url=url,
         transport=transport,
         timeout=httpx.Timeout(connect=5, read=60, write=60, pool=5),
-        headers=headers,
+        headers=_get_headers(api_key, headers),
     )
     return LangGraphClient(client)
 
@@ -1695,3 +1697,23 @@ def _get_api_key(api_key: Optional[str] = None) -> Optional[str]:
         if env := os.getenv(f"{prefix}_API_KEY"):
             return env.strip().strip('"').strip("'")
     return None  # type: ignore
+
+
+def _get_headers(
+    api_key: Optional[str], custom_headers: Optional[dict[str, str]]
+) -> dict[str, str]:
+    """Combine api_key and custom user-provided headers."""
+    custom_headers = custom_headers or {}
+    for header in RESERVED_HEADERS:
+        if header in custom_headers:
+            raise ValueError(f"Cannot set reserved header '{header}'")
+
+    headers = {
+        "User-Agent": f"langgraph-sdk-py/{langgraph_sdk.__version__}",
+        **custom_headers,
+    }
+    api_key = _get_api_key(api_key)
+    if api_key:
+        headers["x-api-key"] = api_key
+
+    return headers

libs/sdk-py/pyproject.toml

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